大模型应用开发实战（2）——手撕Transformer

‍♂️ 个人主页小李同学_LSH的主页✍ 作者简介LLM学习者 希望大家多多支持我们一起进步如果文章对你有帮助的话欢迎评论 点赞 收藏 加关注目录一、为什么会有 TransformerRNN 到底卡在哪了二、先看整体Transformer 到底长什么样Encoder-Decoder 整体结构三、真正的核心Self-Attention 是怎么工作的1先生成 Q、K、V2计算相关性分数3缩放 Softmax4加权求和得到新表示从自注意力到多头注意力一个最小 PyTorch 版本四、为什么还要 Multi-Head Attention一个简化版多头代码五、注意力后面为什么还要接 FFN一个简化版 FFN六、Add Norm为什么这俩小东西这么重要1残差连接2层归一化Encoder Layer 内部结构七、没有位置编码Transformer 连词序都分不清为什么用 sin / cos一个最小位置编码实现八、解码器和编码器有什么不一样编码器层解码器层1为什么要 Masked Self-Attention2什么是 Cross-Attention九、把所有模块拼起来一个 Transformer 层到底在干嘛如果说大模型是今天 AI 应用开发的“大脑”那 Transformer 就是这颗大脑的“基础电路板”。很多人学大模型时一上来就被各种名词砸晕Self-Attention、Multi-Head、Add Norm、FFN、Positional Encoding……看着都眼熟但一到真正解释“Transformer 到底是怎么工作的”就很容易卡住。这篇文章我不准备只停留在“概念扫盲”而是想按照“先看整体再拆模块最后拼回去”的思路真正把 Transformer 手撕一遍。你看完至少要搞清楚三件事Transformer 为什么能替代 RNN一个 Transformer 层里到底做了什么为什么它后来会演化成大语言模型的基座架构一、为什么会有 TransformerRNN 到底卡在哪了在 Transformer 之前处理序列最经典的路线是 RNN、LSTM 这类循环神经网络。它们的优点很直观天生适合按顺序处理文本一个词接一个词读下去理论上可以把前文信息不断传递下去。问题也同样明显它必须按时间步顺序计算第 t步要等第 t−1步先算完这直接限制了并行训练能力同时长序列里还容易出现长期依赖难学、训练效率低的问题。从这个瓶颈切入引出 Transformer “彻底抛弃循环结构、完全依赖注意力机制”的转变原始论文也是这样定义 Transformer 的。你可以把这件事理解成RNN像一个人一字一句顺着读Transformer像一个人直接把整句话铺开同时看所有词之间的关系。也正因为这样Transformer 一出现就天然更适合 GPU 并行和大规模训练。原始论文在两项机器翻译任务上都强调了这一点不仅效果更好而且训练更快、更可并行。二、先看整体Transformer 到底长什么样最初的 Transformer 不是为聊天机器人设计的而是为端到端机器翻译设计的所以它采用的是一个非常经典的Encoder-Decoder结构编码器负责“读懂输入”解码器负责“逐步生成输出”。Hello-Agents 对这部分的描述很清楚编码器为每个输入词元生成上下文化表示解码器则一边参考已经生成的前文一边“咨询”编码器输出再决定下一个词生成什么。Encoder-Decoder 整体结构最初的 Transformer 模型是为端到端任务机器翻译而设计的。它在宏观上遵循了一个经典的编码器-解码器 (Encoder-Decoder)架构。Transformer 整体架构图我们可以将这个结构理解为一个分工明确的团队编码器 (Encoder)任务是“理解”输入的整个句子。它会读取所有Token最终为每个词元生成一个富含上下文信息的向量表示。解码器 (Decoder)任务是“生成”目标句子。它会参考自己已经生成的前文并“咨询”编码器的理解结果来生成下一个词。模块作用一句话理解Embedding把 token 变成向量让词进入神经网络世界Positional Encoding注入位置信息告诉模型“谁在前谁在后”Self-Attention建模词与词关系每个词都能看全句FFN做逐位置特征变换对每个位置单独深加工Add Norm稳定训练残差连接 归一化Decoder Mask防止偷看未来保证生成时只能看前文三、真正的核心Self-Attention 是怎么工作的Transformer 最关键的创新就是注意力机制。用一个很形象的例子解释它在句子 “The agent learns because it is intelligent.” 里模型在处理 “it” 时会自动更关注前面的 “agent”因为这能帮助理解 “it” 指代谁。自注意力就是把这种“当前词该看谁”的过程数学化。1先生成 Q、K、V对于输入矩阵 X先通过三个可学习的线性变换得到这里QQuery表示“我正在问什么”KKey表示“我是什么标签别人可以来匹配我”VValue表示“我真正携带的信息”Q/K/V 的角色划分。2计算相关性分数接下来用当前词的 Q 去和所有词的 K 做点积得到注意力分数矩阵点积越大说明两个词越相关。3缩放 Softmax因为维度大时点积值会变大导致 Softmax 过于尖锐Transformer 论文加入了一个缩放项很多人第一次看到会想“都是从输入 XXX 线性变换出来的为什么不直接一个矩阵搞定非要分成三份”先说最直观的理解可以把它理解成QQuery我现在想找什么KKey我这里有什么特征适不适合被你找到VValue如果你真的关注我我把什么信息给你Q 用来提问K 用来匹配V 用来传递内容。这三件事本来就不是一回事。举个句子例子The animal didnt cross the street becauseitwas tired.模型在处理it的时候想做的是它要问我这个代词更可能指谁句子里别的词要提供一种可被匹配的标签真正被选中之后还要提供它自己的语义信息如果这三件事混成一个向量来做模型会很别扭。拆成三套投影以后模型就可以分别学什么样的特征适合拿来“发问”什么样的特征适合拿来“被匹配”什么样的特征适合拿来“输出内容”这就更灵活了。这就是最经典的Scaled Dot-Product Attention公式原始论文以这条公式作为注意力核心。4加权求和得到新表示Softmax 后得到的是每个词对其它词的注意力权重再用这些权重去加权求和 V就得到了融合全局信息后的新向量表示。import torch import torch.nn as nn import math # --- 占位符模块将在后续小节中实现 --- class PositionalEncoding(nn.Module): 位置编码模块 def forward(self, x): pass class MultiHeadAttention(nn.Module): 多头注意力机制模块 def forward(self, query, key, value, mask): pass class PositionWiseFeedForward(nn.Module): 位置前馈网络模块 def forward(self, x): pass # --- 编码器核心层 --- class EncoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super(EncoderLayer, self).__init__() self.self_attn MultiHeadAttention() # 待实现 self.feed_forward PositionWiseFeedForward() # 待实现 self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, mask): # 残差连接与层归一化 # 1. 多头自注意力 attn_output self.self_attn(x, x, x, mask) x self.norm1(x self.dropout(attn_output)) # 2. 前馈网络 ff_output self.feed_forward(x) x self.norm2(x self.dropout(ff_output)) return x # --- 解码器核心层 --- class DecoderLayer(nn.Module): def __init__(self, d_model, num_heads, d_ff, dropout): super(DecoderLayer, self).__init__() self.self_attn MultiHeadAttention() # 待实现 self.cross_attn MultiHeadAttention() # 待实现 self.feed_forward PositionWiseFeedForward() # 待实现 self.norm1 nn.LayerNorm(d_model) self.norm2 nn.LayerNorm(d_model) self.norm3 nn.LayerNorm(d_model) self.dropout nn.Dropout(dropout) def forward(self, x, encoder_output, src_mask, tgt_mask): # 1. 掩码多头自注意力 (对自己) attn_output self.self_attn(x, x, x, tgt_mask) x self.norm1(x self.dropout(attn_output)) # 2. 交叉注意力 (对编码器输出) cross_attn_output self.cross_attn(x, encoder_output, encoder_output, src_mask) x self.norm2(x self.dropout(cross_attn_output)) # 3. 前馈网络 ff_output self.feed_forward(x) x self.norm3(x self.dropout(ff_output)) return x从自注意力到多头注意力上图对应的就是 “生成 Q/K/V → 算相关性 → 缩放归一化 → 加权求和”的四步流程。一个最小 PyTorch 版本import torch import math def scaled_dot_product_attention(Q, K, V, maskNone): scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(Q.size(-1)) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn torch.softmax(scores, dim-1) output torch.matmul(attn, V) return output, attn这段代码本质上就是上面公式的直接实现。四、为什么还要 Multi-Head Attention如果只有一个注意力头模型每次只能学一种“关注模式”。比如处理 “it” 时它可能学会关注主语但语言里不只有指代关系还有时态关系、搭配关系、从句关系。对多头注意力的解释很到位把一次注意力拆成多次让多个“专家”从不同子空间看问题再把结果合起来。这意味着模型可以同时关注多个不同维度的关系。阶段张量形状输入 (X)((B, L, 512))线性映射后的 (Q,K,V)((B, L, 512))拆成 8 个头((B, 8, L, 64))每个头做注意力((B, 8, L, 64))拼接回来((B, L, 512))一个简化版多头代码import torch.nn as nn class MultiHeadAttention(nn.Module): def __init__(self, d_model512, num_heads8): super().__init__() assert d_model % num_heads 0 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model) def split_heads(self, x): B, L, D x.size() return x.view(B, L, self.num_heads, self.d_k).transpose(1, 2) def combine_heads(self, x): B, H, L, DK x.size() return x.transpose(1, 2).contiguous().view(B, L, self.d_model) def forward(self, Q, K, V, maskNone): Q self.split_heads(self.W_q(Q)) K self.split_heads(self.W_k(K)) V self.split_heads(self.W_v(V)) attn_output, _ scaled_dot_product_attention(Q, K, V, mask) output self.W_o(self.combine_heads(attn_output)) return output五、注意力后面为什么还要接 FFN很多人第一次看 Transformer 都会疑惑前面不是已经全局建模了吗后面再接一个前馈网络干嘛答案是注意力层负责聚合信息FFN 负责加工信息。如果说注意力是在整个序列范围内“动态聚合相关信息”那 FFN 就是在每个位置上进一步抽取高阶特征。更重要的是它是position-wise的也就是对每个位置独立处理但共享同一组参数。然后再映射回去。这种“先扩张、再压缩”的设计能让模型学到更丰富的非线性特征。一个简化版 FFNclass PositionWiseFeedForward(nn.Module): def __init__(self, d_model512, d_ff2048, dropout0.1): super().__init__() self.linear1 nn.Linear(d_model, d_ff) self.linear2 nn.Linear(d_ff, d_model) self.dropout nn.Dropout(dropout) self.relu nn.ReLU() def forward(self, x): x self.linear1(x) x self.relu(x) x self.dropout(x) x self.linear2(x) return x六、Add Norm为什么这俩小东西这么重要Transformer 每个子层外面都会套一个Add Norm也就是残差连接Add层归一化LayerNormHello-Agents 把这部分单独拎出来讲是对的因为很多人只关注 Attention却忽略了真正让深层 Transformer 能稳定训练的是这些“看起来不起眼”的工程结构。1残差连接它的作用是让梯度在反向传播时能绕过复杂子层直接往前传缓解深层网络训练中的梯度消失问题。2层归一化LayerNorm 会在特征维度上做归一化让每一层输入分布更稳定训练更容易收敛。概括成“让每层输入分布保持稳定加速收敛并提高训练稳定性”。Encoder Layer 内部结构先自注意力再残差归一化再 FFN再残差归一化。七、没有位置编码Transformer 连词序都分不清注意力机制有一个天然缺陷它只关心“词和词之间的关系强不强”但它本身并不知道词序。对于纯 Attention 来说“agent learns”和“learns agent”可能被看成同一组词的排列而这显然不对。位置编码部分自注意力本身不包含位置信息所以必须额外给每个 token 向量加一个表示位置的向量。原始 Transformer 论文采用的是正弦—余弦位置编码为什么用 sin / cos你不用一上来就深究它的严格数学性质先记住两个直觉就够了不同位置会得到不同编码不同维度用不同频率的波形模型更容易从中学习到相对位移关系。一个最小位置编码实现import torch import torch.nn as nn import math class PositionalEncoding(nn.Module): def __init__(self, d_model, dropout0.1, max_len5000): super().__init__() self.dropout nn.Dropout(pdropout) pe torch.zeros(max_len, d_model) position torch.arange(max_len).unsqueeze(1) div_term torch.exp( torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model) ) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) self.register_buffer(pe, pe.unsqueeze(0)) def forward(self, x): x x self.pe[:, :x.size(1)] return self.dropout(x)八、解码器和编码器有什么不一样编码器层和解码器层最大的区别不在 FFN也不在 Add Norm而在注意力结构。编码器层只有两部分主干Multi-Head Self-AttentionFFN因为它的任务只是“理解输入”。解码器层有三部分主干Masked Multi-Head Self-AttentionCross-AttentionFFN这里多出来两个关键点。1为什么要 Masked Self-Attention因为解码器在生成第 t个词时不能偷看第 t1个词。在 Softmax 之前把未来位置的分数打成一个很大的负数经过 Softmax 后这些位置概率就变成 0从数学上阻止模型看未来。2什么是 Cross-Attention解码器除了看自己已经生成的前文还要看编码器输出的“源句子理解结果”。这一步就叫交叉注意力Query 来自当前解码器状态Key / Value 来自编码器输出于是解码器一边看自己写到哪了一边看输入句子表达了什么。九、把所有模块拼起来一个 Transformer 层到底在干嘛如果你现在还觉得乱可以把一个 Encoder Layer 理解成下面这个执行链输入 token 向量加位置编码做自注意力让每个词看见全句做残差连接和归一化做 FFN进一步提特征再做一次残差连接和归一化输出给下一层整个 Encoder 堆 NNN 层后输出一组带全局上下文的表示。解码器则在“前文 编码器输出”的双重条件下一步步生成目标序列。