1. 从零开始为什么我们需要亲手实现注意力机制如果你和我一样最初接触Transformer时是被那篇著名的《Attention is All You Need》论文吸引的。但说实话光看论文里的公式比如那个缩放点积注意力公式你可能觉得懂了但关上论文让你自己用代码把它写出来大脑是不是瞬间一片空白这就是理论和实践之间的鸿沟。我当年也卡在这里很久直到我发现了《The Annotated Transformer》这个项目。它就像一位耐心的导师把论文里每一个抽象的数学符号都对应成了PyTorch里的一行行具体代码。所以这一讲我们不空谈理论我们就跟着这个项目的代码一行一行地把Transformer最核心的“注意力机制”给搭建出来。相信我当你亲手敲完这些代码并看到它跑起来时那种理解是看十篇论文都比不上的。这个过程的目标非常明确我们不是要调用现成的nn.Transformer模块而是要自己从最基础的张量操作开始构建出ScaledDotProductAttention和MultiHeadAttention。为什么非要这么“自讨苦吃”呢因为只有深入到代码层面你才能真正理解“查询Query”、“键Key”、“值Value”这三个张量到底是如何互动注意力权重是如何计算和应用的以及“多头”到底是怎么并行工作的。这些细节是你在日后调参、魔改模型或者解决诡异bug时最宝贵的直觉来源。好了废话不多说我们打开编辑器新建一个Python文件开始我们的构建之旅吧。2. 搭建舞台理解注意力机制的输入与核心公式在动手写代码之前我们得先搞清楚我们要实现的东西到底是什么它的输入输出又是什么。你可以把注意力机制想象成一个信息检索系统。假设你有一本百科全书Value这本书没有目录但有一个索引表Key。现在你心中有一个问题Query。注意力机制的工作就是用你的问题Query去索引表Key里查找最相关的条目然后根据查找结果的“相关度分数”去百科全书Value里提取相应的内容最后把这些内容加权组合起来形成你的答案。在Transformer的语境下Query、Key、Value都不是一个单词而是一系列向量。比如你输入一个句子“I love AI”经过嵌入层后你会得到三个形状相同的张量Q, K, V。它们的形状通常是[batch_size, sequence_length, d_model]其中d_model是模型的特征维度比如512。注意力机制的核心就是计算Query和Key的相似度然后对Value进行加权求和。论文里给出的缩放点积注意力公式是这个样子的[ \text{Attention}(Q, K, V) \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V ]这个公式看着简单但里面有几个关键点。第一为什么是点积点积可以衡量两个向量的相似度值越大越相似。第二为什么要除以(\sqrt{d_k})这是一个非常巧妙的“缩放”操作。因为当d_kKey的维度很大时点积的结果可能会变得非常大这会导致softmax函数的梯度变得极其微小也就是梯度消失问题除以这个缩放因子可以让数值分布更稳定有利于训练。第三softmax的作用是把相似度分数转化为一个概率分布所有权重和为1这样我们就能知道对于当前Query应该“关注”每个Key对应Value的多少比例。理解了这些我们再看代码就会清晰很多。在《The Annotated Transformer》中这个核心函数被优雅地实现为一个独立的函数。接下来我们就进入实战环节看看如何用PyTorch的张量操作把这个公式“翻译”出来。2.1 逐行解剖缩放点积注意力函数我们直接来看代码我会在每一行后面加上详细的“人话”解释。import torch import torch.nn.functional as F import math def attention(query, key, value, maskNone, dropoutNone): 计算缩放点积注意力。 参数: query: 形状为 [batch_size, h, seq_len_q, d_k] key: 形状为 [batch_size, h, seq_len_k, d_k] value: 形状为 [batch_size, h, seq_len_v, d_v] (通常 seq_len_k seq_len_v) mask: 可选的掩码张量用于在特定位置如填充位置或未来位置屏蔽注意力 dropout: 可选的Dropout层 返回: 加权后的value以及注意力权重 # 第一步获取key的最后一个维度也就是d_k用于缩放 d_k query.size(-1) # 第二步计算Query和Key的点积分数 # query.shape: [..., seq_len_q, d_k] # key.transpose(-2, -1).shape: [..., d_k, seq_len_k] # 矩阵相乘后scores.shape: [..., seq_len_q, seq_len_k] # 这个分数矩阵的每一个元素[i, j]就代表了第i个Query对第j个Key的相似度。 scores torch.matmul(query, key.transpose(-2, -1)) # 第三步缩放这是稳定训练的关键。 # 将分数除以 sqrt(d_k)防止点积值过大导致softmax梯度太小。 scores scores / math.sqrt(d_k) # 第四步应用掩码如果提供了的话 # 掩码在解码器中非常关键用于防止当前位置“看到”未来的信息。 # 掩码通常是一个布尔张量需要被填充的位置是1True我们给这些位置赋一个极小的负数如-1e9。 # 这样在后续softmax时这些位置的指数会趋近于0权重也就为0了。 if mask is not None: scores scores.masked_fill(mask 0, -1e9) # 第五步对最后一个维度seq_len_k做softmax将分数转化为概率分布注意力权重 # p_attn.shape: [..., seq_len_q, seq_len_k] p_attn F.softmax(scores, dim-1) # 第六步可选地应用Dropout。这是一个防止过拟合的小技巧在训练时随机“丢弃”一部分注意力权重。 if dropout is not None: p_attn dropout(p_attn) # 第七步用注意力权重对Value进行加权求和得到最终的输出。 # p_attn.shape: [..., seq_len_q, seq_len_k] # value.shape: [..., seq_len_v, d_v] (seq_len_k seq_len_v) # 矩阵相乘后返回的张量形状为 [..., seq_len_q, d_v] return torch.matmul(p_attn, value), p_attn我来分享一个我初次实现时踩过的坑矩阵乘法的维度对齐。在scores torch.matmul(query, key.transpose(-2, -1))这一行key.transpose(-2, -1)是把key的最后两个维度调换顺序。为什么是-2和-1这是PyTorch中非常方便的相对索引-1永远代表最后一个维度。假设key的形状是[batch, heads, seq_len, d_k]那么转置后形状就是[batch, heads, d_k, seq_len]。这样query的[..., seq_len_q, d_k]才能和转置后的key[..., d_k, seq_len_k]进行矩阵乘法得到[..., seq_len_q, seq_len_k]的形状。如果你不小心转置错了维度程序会直接报错这是一个很好的检查点。另一个需要注意的点是掩码的应用时机。一定要在softmax之前应用掩码。因为我们的目的是让某些位置的注意力权重变为零而softmax之前这些位置的分数是一个很大的负数比如-1e9经过e^{-1e9}计算后结果无限接近于0完美实现了“屏蔽”效果。如果你在softmax之后才应用掩码虽然也能把权重设为零但会破坏概率分布总和为1的性质可能会带来意想不到的问题。3. 从单头到多头并行化的艺术实现了基础的注意力函数后我们来到了Transformer的另一个精髓设计多头注意力Multi-Head Attention。论文里说单次的注意力就像一次检索可能只关注到一种模式的关系比如语法结构。而多头注意力允许模型同时进行多次检索每次关注不同的模式比如语义、指代、语法等最后把结果合并起来这样模型就能捕捉到更丰富、更细微的上下文信息。具体是怎么做的呢想象一下你有一个维度为d_model比如512的输入。在单头注意力里我们直接用这个512维的向量去做Q、K、V。但在多头注意力里我们把这个512维的向量“切”成h个头比如8个头每个头的维度就是d_k d_model / h 512/864。然后这8个头并行地执行我们上面写好的attention函数。最后我们把8个头的输出再“拼接”起来通过一个线性变换映射回d_model维度。这个过程听起来有点绕但用代码来表示就非常直观了。它本质上是一个“分而治之再合并”的策略。下面我们就来看看《The Annotated Transformer》中是如何构建这个MultiHeadedAttention类的。3.1 构建多头注意力模块这个类会稍微复杂一点因为它包含了几个线性变换层和“分头”、“合并”的逻辑。我们一步步来拆解。import torch.nn as nn from torch.nn import Dropout class MultiHeadedAttention(nn.Module): def __init__(self, h, d_model, dropout0.1): 初始化多头注意力层。 参数: h: 注意力头的数量 d_model: 模型的维度也是输入输出的维度 dropout: Dropout概率 super(MultiHeadedAttention, self).__init__() # 确保d_model可以被h整除这样每个头的维度d_k才是整数 assert d_model % h 0 # 我们假设每个头的维度d_v等于d_k self.d_k d_model // h self.h h # 创建四个线性层。 # 前三个linears[0], linears[1], linears[2]分别用于将输入投影到Q, K, V空间。 # 注意这里并不是直接分成h个头而是先做一个大的线性变换然后在计算注意力时再分头。 # 最后一个线性层linears[3]用于将多个头的输出拼接后投影回d_model维度。 self.linears clones(nn.Linear(d_model, d_model), 4) self.attn None # 用于存储注意力权重方便后续可视化 self.dropout nn.Dropout(pdropout) def forward(self, query, key, value, maskNone): 前向传播。 参数形状: query, key, value - [batch_size, seq_len, d_model] 返回形状: - [batch_size, seq_len, d_model] if mask is not None: # 同样的掩码需要应用到所有头上所以增加一个维度head维度 mask mask.unsqueeze(1) batch_size query.size(0) # 1. 线性投影并分头 (batch_size, seq_len, d_model) - (batch_size, seq_len, h, d_k) # 然后转置为 (batch_size, h, seq_len, d_k)方便并行计算 query, key, value [ lin(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2) for lin, x in zip(self.linears, (query, key, value)) ] # 注意这里我们只用了前三个线性层分别处理Q, K, V。 # linears[0](query), linears[1](key), linears[2](value) # 2. 在分好头的张量上应用注意力函数 # x的形状: (batch_size, h, seq_len_q, d_k) # self.attn的形状: (batch_size, h, seq_len_q, seq_len_k) 存储了注意力权重 x, self.attn attention(query, key, value, maskmask, dropoutself.dropout) # 3. 将多个头的输出“拼接”起来 # 首先把头的维度转回最后 (batch_size, h, seq_len, d_k) - (batch_size, seq_len, h, d_k) x x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k) # contiguous()是为了确保张量在内存中是连续存储的某些view操作需要这个前提。 # 4. 通过最后一个线性层投影回原始维度并返回 return self.linears[-1](x)这里有几个实现细节值得深究。首先是分头的操作。代码中并没有真的准备8个不同的线性层而是先用一个大的线性层nn.Linear(d_model, d_model)把输入投影一下。然后通过.view(batch_size, -1, self.h, self.d_k)操作把d_model维度的输出“重塑”成[h, d_k]的形状。.transpose(1, 2)是为了把“头”的维度h提到序列长度维度前面变成[batch, h, seq_len, d_k]这样在调用attention函数时batch和h维度就被自动视为“批处理”维度实现了真正的并行计算。这个设计非常高效。其次是**clones函数**。这是一个工具函数用于创建N个相同的模块。在Transformer中编码器、解码器都由多个相同的层堆叠而成这个函数让代码非常简洁。def clones(module, N): 生成N个相同的层。 return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])最后self.attn None和self.attn attn这一步非常巧妙。它把每次前向传播计算出的注意力权重保存了下来。这有什么用呢可视化在模型训练或推理之后我们可以直接取出model.encoder.layers[0].self_attn.attn这个张量把它画成热力图。你就能亲眼看到当模型处理“The animal didnt cross the street because it was too tired”这句话时对于“it”这个词模型到底更关注“animal”还是“street”。这种直观的反馈对于理解模型行为和调试至关重要也是《The Annotated Transformer》项目的一大亮点。4. 让网络更稳定残差连接与层归一化如果你以为实现了多头注意力就大功告成了那还差得远。Transformer之所以能堆叠得很深比如论文里的6层编码器6层解码器离不开两个“稳定器”残差连接Residual Connection和层归一化Layer Normalization。它们在《The Annotated Transformer》中被一起封装成了一个叫SublayerConnection的模块。我先打个比方。训练一个很深的神经网络就像用积木搭一个很高的塔越往上越容易晃。残差连接就像在每层积木之间加了一根坚固的柱子直接连通底部和顶部。这样梯度可以理解为调整指令就能沿着这根柱子直接传到底层避免了在中间层层传递时可能发生的“指令消失或爆炸”问题。层归一化则像一个水平仪确保每一层积木的输出都保持在一个相对稳定、分布良好的范围内不会因为某一层的输出过大或过小而影响下一层。在代码中这两个技术被紧密地结合在一起。标准的做法是“Pre-Norm”先归一化再进入子层还是“Post-Norm”先经过子层再归一化在学术界有过讨论原始Transformer论文和《The Annotated Transformer》实现的是Post-Norm也就是LayerNorm(x Sublayer(x))。不过现在很多改进模型如BERT、GPT更喜欢用Pre-Norm因为它通常训练更稳定。我们这里先忠实于原论文的实现。4.1 实现子层连接模块让我们看看这个关键的模块是如何用几行代码实现的。class SublayerConnection(nn.Module): 一个残差连接后面跟着一个层归一化。 注意为了简化代码归一化操作在前但这是Post-Norm的实现。 具体顺序在forward函数中体现 def __init__(self, size, dropout): super(SublayerConnection, self).__init__() self.norm nn.LayerNorm(size) self.dropout nn.Dropout(dropout) def forward(self, x, sublayer): 将残差连接应用于任何与其输入输出形状相同的子层。 参数: x: 输入张量 sublayer: 一个函数例如一个注意力层或前馈网络接受一个输入并返回相同形状的输出。 # 注意这里的顺序先对x做层归一化然后传入子层如注意力层再经过dropout最后加上原始的输入x。 # 这就是 Post-Norm: x dropout(sublayer(norm(x))) return x self.dropout(sublayer(self.norm(x)))这个forward函数的设计非常Pythonic。它接受一个函数sublayer作为参数。这意味着无论是自注意力层还是前馈神经网络层都可以被这个相同的连接模块包裹。在编码器层中你会看到这样的用法# 假设self.self_attn是一个MultiHeadedAttention实例self.feed_forward是一个前馈网络 # self.sublayer是一个由两个SublayerConnection实例组成的ModuleList # 第一个子层连接处理自注意力 x self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # 第二个子层连接处理前馈网络 x self.sublayer[1](x, self.feed_forward)这里用lambda函数是因为注意力层需要传入Q, K, V三个参数而sublayer函数只期望一个输入。lambda x: self.self_attn(x, x, x, mask)创建了一个匿名函数它接受一个输入x然后将其同时作为Q, K, V传给注意力层。这种设计保持了接口的统一和简洁。通过这种方式Transformer的每一层都具备了抵抗梯度消失和稳定训练的能力为构建深层模型打下了坚实基础。5. 组装核心部件构建编码器层现在我们已经有了最核心的积木多头注意力MultiHeadedAttention和子层连接SublayerConnection。接下来我们要用它们来搭建Transformer的编码器层EncoderLayer。一个编码器层主要包含两个子层一个多头自注意力层和一个简单的前馈神经网络。每个子层外面都包裹着刚才提到的残差连接和层归一化。自注意力Self-Attention是Transformer理解上下文的关键。在编码器中“自”意味着它的Query, Key, Value都来自同一个来源即上一层的输出。通过自注意力序列中的每个位置都可以关注到序列中所有其他位置的信息从而建立起丰富的上下文表征。前馈神经网络则是一个简单的两层全连接层通常中间有一个ReLU激活函数它的作用是对每个位置的特征进行独立的、非线性的变换。让我们看看在代码中这两个部分是如何被组装成一个完整的编码器层的。理解了这一层解码器层也就触类旁通了。5.1 编码器层的实现class EncoderLayer(nn.Module): 编码器由自注意力和前馈网络两个子层构成 def __init__(self, size, self_attn, feed_forward, dropout): super(EncoderLayer, self).__init__() self.self_attn self_attn # 传入一个MultiHeadedAttention实例 self.feed_forward feed_forward # 传入一个PositionwiseFeedForward实例 # 创建两个SublayerConnection分别用于连接两个子层 self.sublayer clones(SublayerConnection(size, dropout), 2) self.size size # 通常是d_model用于层归一化 def forward(self, x, mask): 前向传播。 参数: x: 输入张量形状 [batch_size, seq_len, d_model] mask: 用于自注意力的掩码如填充掩码 # 第一个子层多头自注意力 # self.sublayer[0] 是一个SublayerConnection # lambda函数将输入x同时作为Q, K, V传给自注意力层 x self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # 第二个子层前馈网络 # 直接将前馈网络函数传入即可 x self.sublayer[1](x, self.feed_forward) return x代码非常清晰几乎就是论文中公式的直译。这里有个细节值得注意掩码mask的传递。在编码器中这个掩码通常是“填充掩码”Padding Mask。因为我们的输入序列长度可能不一样需要填充到相同长度比如用0填充。在计算注意力时我们需要屏蔽这些填充位置防止模型关注无意义的填充符。这个掩码会在注意力函数的scores scores.masked_fill(mask 0, -1e9)这一步起作用。那么前馈网络PositionwiseFeedForward又是什么呢它其实很简单就是一个两层线性变换加一个ReLU激活。class PositionwiseFeedForward(nn.Module): 实现FFN(x) max(0, xW1 b1)W2 b2 def __init__(self, d_model, d_ff, dropout0.1): super(PositionwiseFeedForward, self).__init__() self.w_1 nn.Linear(d_model, d_ff) # 第一个线性层扩大维度 self.w_2 nn.Linear(d_ff, d_model) # 第二个线性层投影回原维度 self.dropout nn.Dropout(dropout) def forward(self, x): # 对每个位置独立地进行相同的变换所以叫“Position-wise” return self.w_2(self.dropout(F.relu(self.w_1(x))))这里d_ff通常比d_model大得多论文中是2048而d_model是512。这种“先扩维再压缩”的设计为模型提供了强大的非线性拟合能力。每个位置的变换都是独立的没有序列间的信息交互这也是为什么它计算效率很高的原因。至此一个完整的编码器层就搭建好了。多个这样的层堆叠起来就构成了Transformer的编码器。通过这种逐层抽象模型能够从输入序列中提取出越来越复杂的特征表示。