图解Transformer：从零开始理解自注意力机制（附PyTorch代码示例）

图解Transformer从零开始理解自注意力机制附PyTorch代码示例如果你接触过自然语言处理或者对深度学习模型有所了解那么“Transformer”这个名字对你来说一定不陌生。自从2017年那篇名为《Attention Is All You Need》的论文发表以来Transformer架构已经彻底改变了整个领域。它不仅是BERT、GPT系列等大语言模型的基石更是在计算机视觉、语音识别等多个方向展现出惊人的通用性。然而对于许多初学者而言理解Transformer尤其是其核心的自注意力机制常常被复杂的矩阵运算和抽象的“查询-键-值”概念所困扰。本文的目标就是为你拨开这层迷雾。我们将从一个全新的、视觉化的视角出发一步步拆解Transformer的自注意力计算过程。我不会仅仅复述论文中的公式而是通过直观的图解和可运行的PyTorch代码片段让你亲手“触摸”到数据在模型中的流动。无论你是已经熟悉卷积神经网络和循环神经网络但尚未深入Transformer的开发者还是希望巩固基础、理解其内在运作机制的研究者这篇文章都将为你提供一个坚实、直观的起点。我们将重点关注多头注意力、位置编码这些核心概念让你不仅知道它们是什么更能理解它们为什么有效以及如何用代码实现。1. 从“黑盒”到“白盒”Transformer的整体架构在深入细节之前让我们先对Transformer建立一个宏观的认识。你可以把它想象成一个高级的翻译机器输入一种语言的句子输出另一种语言的句子。这个“黑盒”内部主要由两个对称的部分构成编码器Encoder和解码器Decoder。它们各自像叠罗汉一样由多个完全相同的层堆叠而成原论文中使用了6层但这并非固定不变。编码器的任务是理解输入序列。它接收一串词向量经过层层处理输出一串富含上下文信息的“编码”。每一层编码器内部都包含两个核心子层多头自注意力层这是Transformer的灵魂。它允许序列中的每个词在编码自身时“观察”序列中所有其他词从而动态地捕捉词与词之间的依赖关系无论它们相距多远。前馈神经网络层这是一个简单的全连接网络独立地应用于每个位置。它的作用是进行非线性变换增强模型的表达能力。这两个子层都被一个“残差连接”包裹并紧接着一个“层归一化”操作。这就像为信息流动搭建了高速公路和交通规则能有效缓解深层网络中的梯度消失问题加速训练。解码器的任务是基于编码器的输出生成目标序列。它的结构比编码器稍复杂一些包含三个子层掩码多头自注意力层与编码器的自注意力类似但增加了一个“掩码”确保在预测当前位置时只能“看到”之前已生成的词而不能“偷看”未来的词这保证了生成过程的因果性。编码器-解码器注意力层这是连接编码器和解码器的桥梁。解码器利用这一层将自身的“查询”与编码器输出的“键”和“值”进行交互从而在生成每个目标词时都能聚焦于输入序列中最相关的部分。前馈神经网络层与编码器中的相同。解码器的输出经过一个线性层和一个Softmax层最终转化为目标词汇表上的概率分布从而确定生成的词。注意编码器和解码器中的自注意力层虽然同名但功能有微妙差异。编码器的自注意力是“全局视野”而解码器的第一个自注意力层是“单向视野”这是理解其工作原理的关键。为了让你对数据流有更具体的感受我们来看一个简化的维度示例。假设我们的词嵌入维度是d_model512一个包含3个词的序列经过嵌入层后会变成一个形状为[3, 512]的矩阵。这个矩阵将流经编码器的各个层。下面的表格对比了编码器和解码器在处理这个输入时的核心差异组件输入形状 (示例)核心操作输出形状 (示例)关键特性编码器 (单层)[3, 512]1. 多头自注意力2. 前馈网络[3, 512]全局上下文并行计算解码器 (单层)[2, 512](已生成部分)1. 掩码自注意力2. 编码器-解码器注意力3. 前馈网络[2, 512]因果掩码关注编码器输出最终输出层[1, 512](最后一个解码器输出)线性层 Softmax[1, vocab_size]产生下一个词的概率分布这个宏观框架是理解所有细节的基础。接下来我们将打开第一个也是最关键的“黑盒子”自注意力机制。2. 自注意力机制让模型学会“动态聚焦”自注意力是Transformer理解上下文的核心。它的目标很简单为序列中的每个位置生成一个新的表示这个表示是所有位置信息的加权和而权重则由位置之间的相关性动态决定。2.1 直觉从“it”指代什么说起理解自注意力最好的方式是通过一个例子。考虑这个句子“The animal didnt cross the street because it was too tired.” 对于人类来说很容易判断“it”指的是“animal”而不是“street”。但对于一个模型呢传统的RNN通过隐状态逐步传递信息距离较远的词间关系容易衰减。而自注意力机制让模型在处理“it”时可以直接去计算它与“animal”、“street”甚至“tired”的关联度从而做出准确判断。这个过程可以类比于信息检索系统。当你在搜索引擎中输入一个查询Query系统会将其与数据库中文档的关键词Key进行匹配然后返回最相关的文档内容Value。自注意力机制中序列中的每个词都同时扮演了这三种角色它既作为查询去检索其他词也作为键和值被其他词检索。2.2 分步计算从向量到矩阵现在让我们把直觉转化为数学和代码。假设我们有一个包含两个词的微型序列其嵌入向量分别为x1和x2维度为d_model4。第一步创建查询Q、键K、值V向量每个输入向量x会通过三个不同的可学习权重矩阵W_Q,W_K,W_V进行线性变换生成对应的q,k,v向量。通常这些向量的维度d_k,d_v会小于原始嵌入维度以减少计算量。import torch import torch.nn as nn import math # 假设嵌入维度为4我们使用更小的注意力维度 d_model 4 d_k d_v 3 # 注意力头维度 batch_size 1 seq_len 2 # 随机初始化输入嵌入 (batch_size, seq_len, d_model) x torch.randn(batch_size, seq_len, d_model) # 定义可学习的权重矩阵 W_Q nn.Linear(d_model, d_k, biasFalse) W_K nn.Linear(d_model, d_k, biasFalse) W_V nn.Linear(d_model, d_v, biasFalse) # 计算 Q, K, V Q W_Q(x) # 形状: (1, 2, 3) K W_K(x) # 形状: (1, 2, 3) V W_V(x) # 形状: (1, 2, 3)第二步计算注意力分数注意力分数衡量了查询向量q_i与所有键向量k_j的相关性。最常用的方法是计算点积。例如对于第一个词我们需要计算q1与k1、k2的点积。# 计算注意力分数 (简化未考虑缩放和掩码) # 使用矩阵乘法一次性计算所有分数 # (batch, seq_len, d_k) (batch, d_k, seq_len) - (batch, seq_len, seq_len) attention_scores torch.matmul(Q, K.transpose(-2, -1)) # 转置最后两个维度 print(原始注意力分数:\n, attention_scores)第三步与第四步缩放与Softmax点积的结果可能数值很大导致Softmax函数的梯度非常小。因此通常会将分数除以键向量维度的平方根sqrt(d_k)进行缩放然后应用Softmax函数将分数转化为概率分布所有权重为正且和为1。# 缩放 scaled_attention_scores attention_scores / math.sqrt(d_k) # 应用Softmax在最后一个维度seq_len上进行 attention_weights torch.nn.functional.softmax(scaled_attention_scores, dim-1) print(缩放后的注意力权重:\n, attention_weights)输出会是一个[1, 2, 2]的矩阵。第一行[a11, a12]表示第一个词分配给第一个词和第二个词的注意力权重第二行[a21, a22]同理。第五步与第六步加权求和将注意力权重作为系数对值向量V进行加权求和得到自注意力层的输出Z。# 加权求和 # (batch, seq_len, seq_len) (batch, seq_len, d_v) - (batch, seq_len, d_v) Z torch.matmul(attention_weights, V) print(自注意力输出 Z:\n, Z)Z的每一行都是原始序列所有值向量的加权组合它包含了整个序列的上下文信息。将以上步骤合并就得到了自注意力最经典的公式Attention(Q, K, V) softmax(QK^T / sqrt(d_k)) V这个公式简洁而强大它允许模型在常数级操作内捕获任意两个位置间的依赖关系彻底摆脱了RNN的顺序计算限制。3. 多头注意力并行化的“专家委员会”单一的注意力机制可能只关注到一种类型的依赖关系。比如在上面的“it”例子中一个注意力头可能专注于学习指代关系另一个头可能专注于学习时态一致性还有一个头可能关注短语结构。多头注意力Multi-Head Attention就是为了捕获这种不同子空间中的信息而设计的。3.1 多头机制的原理与其只使用一组(W_Q, W_K, W_V)矩阵我们使用h组例如8组不同的权重矩阵。每一组都独立地进行上一节描述的自注意力计算产生一个输出head_i。因为每组权重是随机初始化并独立学习的所以每个“头”会倾向于关注输入信息的不同方面。class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads): super(MultiHeadAttention, self).__init__() assert d_model % num_heads 0, d_model must be divisible by num_heads self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads # 每个头的维度 # 定义生成Q, K, V的线性层输出维度仍是d_model self.W_Q nn.Linear(d_model, d_model, biasFalse) self.W_K nn.Linear(d_model, d_model, biasFalse) self.W_V nn.Linear(d_model, d_model, biasFalse) # 最后的输出线性层 self.W_O nn.Linear(d_model, d_model, biasFalse) def forward(self, Q, K, V, maskNone): batch_size Q.size(0) # 1. 线性投影并分头 # (batch, seq_len, d_model) - (batch, seq_len, num_heads, d_k) # 然后转置为 (batch, num_heads, seq_len, d_k) 以便并行计算 Q self.W_Q(Q).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) K self.W_K(K).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) V self.W_V(V).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) # 2. 计算缩放点积注意力 (为每个头单独计算) # 我们使用一个简化的单头注意力函数这里展示核心矩阵运算 attention_scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: attention_scores attention_scores.masked_fill(mask 0, -1e9) attention_weights torch.nn.functional.softmax(attention_scores, dim-1) # 3. 应用注意力权重到V上 context torch.matmul(attention_weights, V) # (batch, num_heads, seq_len, d_k) # 4. 合并多头 # 转置回 (batch, seq_len, num_heads, d_k) 然后合并最后两个维度 context context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) # 5. 最终线性投影 output self.W_O(context) return output, attention_weights3.2 多头输出的合并与意义每个头产生一个[seq_len, d_k]的输出。我们有h个头所以会得到h个这样的矩阵。为了送入后续的前馈网络我们需要将它们合并。合并的方法是先将所有头的输出在特征维度上拼接起来得到一个[seq_len, h * d_k]的矩阵因为h * d_k d_model然后再通过一个可学习的线性层W_O进行变换最终输出维度仍为d_model。这种设计带来了两个核心优势增强模型容量不同的注意力头可以学习到不同类型的依赖模式模型表达能力更强。保持计算效率虽然头变多了但每个头的维度d_k变小了总的计算复杂度与单头全维度注意力相近却获得了更丰富的表征。在实际的Transformer模型中你常常可以通过可视化不同头的注意力权重发现一些有趣的现象有的头主要关注局部信息相邻词有的头关注句法结构如主谓一致有的头则关注长距离的语义关联。4. 位置编码为无序的注意力注入顺序信息自注意力机制有一个“先天缺陷”它对输入序列的顺序是不敏感的。无论词序如何打乱只要词的集合不变点积注意力计算出的结果就是一样的。这显然不符合语言规律“猫追老鼠”和“老鼠追猫”的意思截然不同。因此Transformer需要一种方法来显式地注入位置信息。4.1 正弦余弦位置编码原论文采用了一种非常巧妙且固定的方法——正弦余弦位置编码Sinusoidal Positional Encoding。它不是通过训练学习得到的而是通过一个确定的公式生成的。对于位置pos和维度i其编码值PE(pos, 2i)和PE(pos, 2i1)计算如下PE(pos, 2i) sin(pos / 10000^(2i/d_model))PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中pos是词在序列中的位置0, 1, 2...i是维度索引。这个公式设计得非常精妙它使得模型能够轻松地学习到相对位置关系因为对于一个固定的偏移量kPE(posk)可以表示为PE(pos)的线性函数。def get_positional_encoding(seq_len, d_model): 生成正弦余弦位置编码矩阵 pe torch.zeros(seq_len, d_model) position torch.arange(0, seq_len, dtypetorch.float).unsqueeze(1) # (seq_len, 1) div_term torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)) pe[:, 0::2] torch.sin(position * div_term) # 偶数维度用sin pe[:, 1::2] torch.cos(position * div_term) # 奇数维度用cos return pe # 形状: (seq_len, d_model) # 示例生成长度为10维度为512的位置编码 seq_len 10 d_model 512 pos_encoding get_positional_encoding(seq_len, d_model) print(f位置编码矩阵形状: {pos_encoding.shape})生成的位置编码矩阵会直接加到词嵌入向量上input_embedding positional_encoding。这样模型在后续的注意力计算中就能同时感知到词的语义信息和位置信息。4.2 可视化与理解为了直观感受位置编码我们可以将其可视化。下图展示了一个小型位置编码矩阵seq_len50, d_model128的热力图。你可以看到沿着位置维度纵轴模式是连续变化的沿着特征维度横轴频率从高到低变化。这种结构使得模型不仅能知道绝对位置还能通过点积操作推导出相对位置。注此处为文字描述实际应用中可通过matplotlib绘制热力图进行观察。你会发现图像呈现出明显的条纹状周期模式。除了正弦余弦编码也有可学习的位置嵌入Learned Positional Embedding等变体例如在BERT中就直接使用一个可训练的嵌入层来学习位置表示。两种方法各有优劣正弦编码可以泛化到比训练时更长的序列而可学习嵌入可能在训练数据长度内获得更优化的表示。5. 实战用PyTorch搭建一个简易Transformer编码器层理解了所有核心组件后是时候动手将它们组合起来了。我们将实现一个完整的Transformer编码器层它包含多头自注意力、前馈网络、残差连接和层归一化。import torch.nn as nn import torch.nn.functional as F class TransformerEncoderLayer(nn.Module): def __init__(self, d_model, num_heads, dim_feedforward2048, dropout0.1): super(TransformerEncoderLayer, self).__init__() # 多头自注意力子层 self.self_attn MultiHeadAttention(d_model, num_heads) # 前馈网络子层两个线性变换 一个激活函数 self.linear1 nn.Linear(d_model, dim_feedforward) self.linear2 nn.Linear(dim_feedforward, d_model) # 层归一化 self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) # Dropout用于防止过拟合 self.dropout nn.Dropout(dropout) def forward(self, src, src_maskNone): src: 输入序列形状 (batch_size, seq_len, d_model) src_mask: 源序列掩码形状 (batch_size, seq_len, seq_len) # 子层1: 多头自注意力 残差 层归一化 attn_output, _ self.self_attn(src, src, src, src_mask) src src self.dropout(attn_output) # 残差连接 src self.norm1(src) # 子层2: 前馈网络 残差 层归一化 ff_output self.linear2(self.dropout(F.relu(self.linear1(src)))) src src self.dropout(ff_output) # 残差连接 src self.norm2(src) return src # 让我们测试一下这个编码器层 d_model 512 num_heads 8 seq_len 10 batch_size 4 encoder_layer TransformerEncoderLayer(d_model, num_heads) x torch.randn(batch_size, seq_len, d_model) # 模拟输入 output encoder_layer(x) print(f输入形状: {x.shape}) print(f编码器层输出形状: {output.shape}) # 应该保持 (4, 10, 512)在这个实现中有几点值得强调残差连接src self.dropout(attn_output)。这允许梯度直接流过缓解了深层网络的优化难题。层归一化在残差相加之后进行有助于稳定训练过程。前馈网络这是一个位置级position-wise的全连接网络每个位置独立进行相同的变换。dim_feedforward通常设置为d_model的4倍。将多个这样的编码器层堆叠起来就构成了Transformer的编码器。解码器层的结构与此类似但增加了掩码自注意力和编码器-解码器注意力层。6. 超越理解训练技巧与实战洞察搭建出模型结构只是第一步让模型有效地学习才是关键。在这一部分我们探讨一些在训练Transformer时至关重要的技巧和实战中容易遇到的“坑”。6.1 学习率调度与优化器选择Transformer模型通常对优化设置非常敏感。原论文使用了Adam优化器并配合一个特殊的学习率调度策略——热身Warmup后平方根倒数衰减Inverse Square Root Decay。# 一个简化的Warmup调度器示例 class WarmupScheduler: def __init__(self, optimizer, d_model, warmup_steps): self.optimizer optimizer self.d_model d_model self.warmup_steps warmup_steps self.current_step 0 def step(self): self.current_step 1 lr (self.d_model ** -0.5) * min(self.current_step ** -0.5, self.current_step * (self.warmup_steps ** -1.5)) for param_group in self.optimizer.param_groups: param_group[lr] lr self.optimizer.step()这个策略在训练早期缓慢增加学习率热身阶段有助于模型稳定初始化之后随着步数增加逐步降低学习率。在实际项目中你可能需要根据任务和数据集调整warmup_steps。6.2 掩码Masking的艺术掩码在Transformer中扮演着多重角色是保证其正确工作的关键。填充掩码Padding Mask用于处理变长序列。在批处理时较短的序列会被填充pad到相同长度。在计算注意力时我们需要屏蔽这些填充位置防止它们影响有效词的注意力分布。序列掩码Sequence Mask / Look-ahead Mask用于解码器的自注意力层。确保在预测第t个词时只能看到第1到t-1个词而不能看到未来的词。def create_padding_mask(seq, pad_token_id0): 创建填充掩码pad_token_id的位置为0需要被mask mask (seq ! pad_token_id).unsqueeze(1).unsqueeze(2) # (batch, 1, 1, seq_len) return mask def create_look_ahead_mask(size): 创建前瞻掩码下三角矩阵 mask torch.triu(torch.ones(size, size), diagonal1).bool() # 需要mask的位置为True在softmax前加上一个很大的负数 return mask # 形状: (seq_len, seq_len) # 使用示例 seq torch.tensor([[1, 2, 3, 0, 0], [4, 5, 0, 0, 0]]) # 假设0是填充符 padding_mask create_padding_mask(seq) # 形状: (2, 1, 1, 5) look_ahead_mask create_look_ahead_mask(5) # 形状: (5, 5) # 在注意力计算中应用掩码 # attention_scores Q K.transpose(-2, -1) / sqrt(d_k) # attention_scores attention_scores.masked_fill(mask 0, -1e9) # 对于padding mask # attention_scores attention_scores.masked_fill(look_ahead_mask, -1e9) # 对于look-ahead mask6.3 梯度裁剪与标签平滑Transformer模型层数深、参数多训练初期容易产生梯度爆炸。梯度裁剪Gradient Clipping是一个简单有效的稳定训练的手段。# 在训练循环中backward之后step之前 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)此外在分类任务如机器翻译的词汇预测中使用标签平滑Label Smoothing可以缓解模型对“绝对正确”的过度自信起到一定的正则化效果往往能提升最终模型的泛化能力。criterion nn.CrossEntropyLoss(label_smoothing0.1) # PyTorch 1.10 支持在我自己的项目经验里Transformer的训练就像调试一台精密仪器。学习率策略不对模型可能根本不收敛掩码设置错了解码时就会产生混乱的结果而不做梯度裁剪在长序列任务上训练几个epoch后很可能就会看到Loss变成NaN。这些技巧虽然琐碎但往往是项目成功与否的关键。理解自注意力机制是掌握了Transformer的“道”而熟练运用这些训练技巧则是掌握了其“术”两者结合才能真正驾驭这个强大的模型。