0.前言
基于上一章的学习我们深刻了解了Transformer架构,并也进一步了解到其真实的训练过程。接着进一步我们继续深究Transformer的源码上进行深层次的学习和理解,明白其中的核心内容首先我们需要关注几篇文章和源码。
https://nlp.seas.harvard.edu/annotated-transformer/https://nlp.seas.harvard.edu/annotated-transformer/以及作者的源码路径
harvardnlp/annotated-transformer at debc9fd747bb2123160a98046ad1c2d4da44a567https://github.com/harvardnlp/annotated-transformer/tree/debc9fd747bb2123160a98046ad1c2d4da44a567#encoder-and-decoderstacks还有Transformers的架构图片(我们并不陌生)
1.背景
简单概述背景:
许多工具都以卷积神经网络作为基础构建单元并行计算所有输入和输出位置的隐藏表示。两个任意输入或输出位置信号所需的作次数随着位置间距离增加而增加,ConvS2S是线性增长,而ByteNet则是对数增长。-导致->学习远距离位置之间的依赖关系变得更加困难。在 Transformer 模型中,这一步的运算量被缩减至一个恒定数值;不过,由于对注意力加权的位置进行了均值化处理,模型的有效分辨率会有所下降,而我们则通过多头注意力机制来抵消这一影响。
2.核心源码+思路剖析
模型架构:CITE 编解码架构
CITE 架构遵循 “编码器 - 交互层 - 解码器” 的经典范式,上下文感知(Context-aware)、交互增强(Interaction-enhanced)、时序建模(Temporal modeling)、高效推理(Efficient inference)
1.位置编码(Positional Encoding)
核心思路:Transformer 没有 RNN 的时序依赖,无法捕捉序列的位置信息,因此需要手动注入位置编码,让模型知道每个 token 在序列中的位置。
核心代码用途:为输入嵌入张量添加位置信息,弥补 Transformer 无法捕捉序列顺序的缺陷,是 Transformer 能处理序列数据的基础。
import torch import torch.nn as nn import math class PositionalEncoding(nn.Module): def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000): super().__init__() self.dropout = nn.Dropout(p=dropout) # 生成位置编码矩阵:(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 = torch.zeros(max_len, d_model) pe[:, 0::2] = torch.sin(position * div_term) # 偶数维度用sin pe[:, 1::2] = torch.cos(position * div_term) # 奇数维度用cos pe = pe.unsqueeze(0) # 增加batch维度:(1, max_len, d_model) self.register_buffer('pe', pe) # 注册为缓冲区(不参与训练) def forward(self, x: torch.Tensor) -> torch.Tensor: """ x: 输入嵌入张量,形状 (batch_size, seq_len, d_model) """ # 仅取和输入序列长度匹配的位置编码 x = x + self.pe[:, :x.size(1), :] return self.dropout(x)关键要点:
register_buffer:位置编码是固定的(不训练),用该方法将 pe 存入模型,避免被优化器更新;- 奇偶维度分别用 sin/cos:保证不同位置的编码有唯一表征,且能通过线性变换捕捉相对位置;
- Dropout:防止过拟合,对注入位置编码后的张量做随机失活。
2.核心注意力模块(Scaled Dot-Product Attention)
核心思路
这是 Transformer 的核心单元,实现 “注意力权重计算”:
- 计算 Query(Q)和 Key(K)的相似度(点积);
- 缩放(除以dk):避免点积值过大导致 Softmax 梯度消失;
- 掩码(Mask):遮挡 padding 或未来位置的 token;
- Softmax:将相似度转为注意力权重;
- 加权求和 Value(V):得到最终的注意力输出。
核心代码用途:
实现 “注意力机制” 的核心逻辑:让模型在处理每个 token 时,自适应地关注输入序列中相关的 token,是 Transformer 能捕捉长距离依赖的关键。
def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor] = None, dropout: Optional[nn.Dropout] = None) -> Tuple[torch.Tensor, torch.Tensor]: """ 输入形状:Q/K/V: (batch_size, n_heads, seq_len, d_k) mask: (batch_size, 1, seq_len, seq_len)(广播兼容) """ d_k = q.size(-1) # 1. 计算Q·K^T / √d_k:(batch_size, n_heads, seq_len_q, seq_len_k) attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k) # 2. 应用掩码:掩码位置设为-∞,Softmax后权重为0 if mask is not None: attn_scores = attn_scores.masked_fill(mask == 0, -1e9) # 3. Softmax计算注意力权重 attn_weights = F.softmax(attn_scores, dim=-1) # 4. 可选Dropout if dropout is not None: attn_weights = dropout(attn_weights) # 5. 加权求和V:(batch_size, n_heads, seq_len_q, d_k) output = torch.matmul(attn_weights, v) return output, attn_weights关键要点:
- 维度变换:
k.transpose(-2, -1)将 K 的最后两个维度交换(seq_len_k ↔ d_k),保证点积维度匹配;- 缩放因子dk:d_k 是 Q/K 的维度,若 d_k 过大,点积结果会很大,Softmax 后梯度趋近于 0;
- 掩码类型:
- Padding Mask:遮挡输入中的 padding token(避免模型关注无效 token);
- Look-Ahead Mask(解码器专用):遮挡未来位置的 token(保证生成时只依赖过去 / 当前 token);
- 掩码实现:
masked_fill(mask == 0, -1e9)把掩码位置设为极小值,Softmax 后权重几乎为 0。
3.核心注意力模块(Multi-Head Attention)
核心思路
将 Q/K/V 拆分为多个 “头”,每个头独立计算 Scaled Dot-Product Attention,最后拼接结果。目的是:
- 让模型同时关注不同位置、不同维度的信息(比如一个头关注语法,一个头关注语义);
- 提升注意力机制的表达能力。
核心代码用途:
将单头注意力扩展为多头,提升模型对不同类型信息的捕捉能力,是 Transformer 注意力机制的核心落地实现。
class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, n_heads: int, dropout: float = 0.1): super().__init__() assert d_model % n_heads == 0, "d_model必须能被n_heads整除" self.d_model = d_model self.n_heads = n_heads self.d_k = d_model // n_heads # 每个头的维度 # 定义4个线性层:Q/K/V的投影 + 最终拼接后的投影 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) self.dropout = nn.Dropout(dropout) self.attn = None # 保存注意力权重(用于可视化) def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: batch_size = q.size(0) # 1. 线性投影 + 拆分多头:(batch_size, seq_len, d_model) → (batch_size, n_heads, seq_len, d_k) q = self.w_q(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) k = self.w_k(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) v = self.w_v(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2) # 2. 计算缩放点积注意力 attn_output, self.attn = scaled_dot_product_attention(q, k, v, mask, self.dropout) # 3. 拼接多头:(batch_size, n_heads, seq_len, d_k) → (batch_size, seq_len, d_model) attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.n_heads * self.d_k) # 4. 最终线性投影 output = self.w_o(attn_output) return output关键要点:
- 维度拆分:
d_model = n_heads * d_k,必须保证整除,否则无法均匀拆分(论文是8);- 维度变换:
transpose(1, 2)将 “头维度” 从第 3 维移到第 2 维,使形状变为(batch_size, n_heads, seq_len, d_k),方便每个头独立计算;- 拼接逻辑:
contiguous()保证内存连续,避免 view 报错,再通过view合并多头维度;- 三种多头注意力场景:
- 编码器自注意力:Q=K=V = 编码器输入,用 Padding Mask;
- 解码器自注意力:Q=K=V = 解码器输入,用 Padding + Look-Ahead Mask;
- 编码器 - 解码器注意力:Q = 解码器输入,K=V = 编码器输出,用 Padding Mask。
4.前馈网络模块(Position-wise Feed-Forward Networks)
核心思路
对每个位置的 token 做相同的两层线性变换(加 ReLU 激活),实现特征的非线性变换。“Position-wise” 表示每个位置独立计算,不跨位置交互。
核心代码用途
对注意力输出的特征做非线性变换,补充注意力机制的表达能力,是 Transformer 中 “特征提取” 的重要环节。
class PositionWiseFeedForward(nn.Module): def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1): super().__init__() self.linear1 = nn.Linear(d_model, d_ff) # 升维:d_model → d_ff self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(d_ff, d_model) # 降维:d_ff → d_model def forward(self, x: torch.Tensor) -> torch.Tensor: # 流程:Linear → ReLU → Dropout → Linear return self.linear2(self.dropout(F.relu(self.linear1(x))))关键要点:
- 维度设计:论文中
d_ff=2048,d_model=512,先升维再降维,增加模型的表达能力;- 位置无关:每个 token 的变换完全独立,输入形状
(batch_size, seq_len, d_model),输出形状相同;- 激活函数:用 ReLU 而非 Sigmoid/Tanh,避免梯度消失问题。
5.编码器模块(Encoder & EncoderLayer)
核心思路:
编码器由 N 个相同的 EncoderLayer 堆叠而成(论文中 N=6),每个 EncoderLayer 包含:
- 多头自注意力(带残差连接 + 层归一化);
- 前馈网络(带残差连接 + 层归一化)。核心范式:
LayerNorm(x + SubLayer(x))(残差 + 层归一化)。核心代码用途:
处理输入序列,生成包含全局上下文信息的序列表征,为解码器提供 “源语言信息”。
1. 编码器层(EncoderLayer)
class EncoderLayer(nn.Module): def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1): super().__init__() self.self_attn = MultiHeadAttention(d_model, n_heads, dropout) # 多头自注意力 self.ff = PositionWiseFeedForward(d_model, d_ff, dropout) # 前馈网络 # 层归一化 + 残差连接的封装(原代码用SublayerConnection) self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: # 1. 自注意力 + 残差 + 层归一化 attn_output = self.self_attn(x, x, x, mask) # Q=K=V=x(自注意力) x = self.norm1(x + self.dropout1(attn_output)) # 2. 前馈网络 + 残差 + 层归一化 ff_output = self.ff(x) x = self.norm2(x + self.dropout2(ff_output)) return x2. 完整编码器(Encoder)
class Encoder(nn.Module): def __init__(self, vocab_size: int, d_model: int, n_layers: int, n_heads: int, d_ff: int, dropout: float = 0.1): super().__init__() self.d_model = d_model self.embedding = nn.Embedding(vocab_size, d_model) # 输入嵌入层 self.pos_encoding = PositionalEncoding(d_model, dropout) # 位置编码 self.layers = nn.ModuleList([EncoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)]) # 堆叠N层 self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> torch.Tensor: seq_len = x.size(1) # 1. 嵌入层 + 位置编码(嵌入层输出需乘以√d_model,论文中的设计) x = self.embedding(x) * math.sqrt(self.d_model) x = self.pos_encoding(x) # 2. 逐层传递 for layer in self.layers: x = layer(x, mask) return x关键要点:
- 残差连接:
x + SubLayer(x)保证梯度能直接回传,避免深层网络的梯度消失;- 层归一化:
LayerNorm作用于最后一维(d_model),稳定训练;- 嵌入层缩放:
embedding(x) * √d_model平衡嵌入层和位置编码的幅值;- ModuleList:堆叠多个 EncoderLayer,保证每层参数独立且可训练。
6.解码器模块(Decoder & DecoderLayer)
核心思路
解码器由 N 个相同的 DecoderLayer 堆叠而成(论文中 N=6),每个 DecoderLayer 包含:
- 解码器自注意力(带 Mask,避免关注未来 token);
- 编码器 - 解码器注意力(关注编码器输出);
- 前馈网络(同编码器);每层都带残差连接 + 层归一化。
核心代码用途
结合编码器的源序列表征和自身的目标序列信息,生成逐步的目标序列表征,为最终预测做准备
1. 解码器层(DecoderLayer)
class DecoderLayer(nn.Module): def __init__(self, d_model: int, n_heads: int, d_ff: int, dropout: float = 0.1): super().__init__() self.self_attn = MultiHeadAttention(d_model, n_heads, dropout) # 解码器自注意力 self.cross_attn = MultiHeadAttention(d_model, n_heads, dropout) # 编码器-解码器注意力 self.ff = PositionWiseFeedForward(d_model, d_ff, dropout) # 前馈网络 self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.norm3 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout) self.dropout3 = nn.Dropout(dropout) def forward(self, x: torch.Tensor, enc_output: torch.Tensor, src_mask: Optional[torch.Tensor] = None, tgt_mask: Optional[torch.Tensor] = None) -> torch.Tensor: # 1. 解码器自注意力(Q=K=V=x,带Look-Ahead+Padding Mask) attn1 = self.self_attn(x, x, x, tgt_mask) x = self.norm1(x + self.dropout1(attn1)) # 2. 编码器-解码器注意力(Q=x,K=V=enc_output,带源序列Padding Mask) attn2 = self.cross_attn(x, enc_output, enc_output, src_mask) x = self.norm2(x + self.dropout2(attn2)) # 3. 前馈网络 ff_output = self.ff(x) x = self.norm3(x + self.dropout3(ff_output)) return x2. 完整解码器(Decoder)
class Decoder(nn.Module): def __init__(self, vocab_size: int, d_model: int, n_layers: int, n_heads: int, d_ff: int, dropout: float = 0.1): super().__init__() self.d_model = d_model self.embedding = nn.Embedding(vocab_size, d_model) # 目标序列嵌入层 self.pos_encoding = PositionalEncoding(d_model, dropout) # 位置编码 self.layers = nn.ModuleList([DecoderLayer(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)]) self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor, enc_output: torch.Tensor, src_mask: Optional[torch.Tensor] = None, tgt_mask: Optional[torch.Tensor] = None) -> torch.Tensor: seq_len = x.size(1) # 1. 嵌入层 + 位置编码(同编码器) x = self.embedding(x) * math.sqrt(self.d_model) x = self.pos_encoding(x) # 2. 逐层传递 for layer in self.layers: x = layer(x, enc_output, src_mask, tgt_mask) return x关键要点:
- 两种掩码:
tgt_mask:解码器自注意力的掩码(Look-Ahead + Padding),保证生成第 i 个 token 时只看前 i-1 个;src_mask:编码器 - 解码器注意力的掩码(源序列 Padding),避免关注源序列的 padding token;- 交叉注意力:Q 来自解码器,K/V 来自编码器,实现 “目标序列关注源序列的相关位置”;
- 三层归一化:对应三个子模块,保证每一步的梯度稳定。
7.完整 Transformer 模型 & 预测模块
核心思路
将编码器、解码器拼接,最后加一层线性层 + Softmax,将解码器输出映射到目标词汇表的概率分布。
核心代码用途
组装编码器和解码器,完成从 “源序列→目标序列概率分布” 的端到端映射,是 Transformer 的最终落地模型。
class Transformer(nn.Module): def __init__(self, src_vocab_size: int, tgt_vocab_size: int, d_model: int = 512, n_layers: int = 6, n_heads: int = 8, d_ff: int = 2048, dropout: float = 0.1): super().__init__() self.encoder = Encoder(src_vocab_size, d_model, n_layers, n_heads, d_ff, dropout) self.decoder = Decoder(tgt_vocab_size, d_model, n_layers, n_heads, d_ff, dropout) self.fc_out = nn.Linear(d_model, tgt_vocab_size) # 最终投影层 def forward(self, src: torch.Tensor, tgt: torch.Tensor, src_mask: Optional[torch.Tensor] = None, tgt_mask: Optional[torch.Tensor] = None) -> torch.Tensor: # 1. 编码器前向 enc_output = self.encoder(src, src_mask) # 2. 解码器前向 dec_output = self.decoder(tgt, enc_output, src_mask, tgt_mask) # 3. 投影到词汇表 output = self.fc_out(dec_output) # 注:训练时不做Softmax(CrossEntropyLoss内置Softmax),预测时加 return output关键要点:
- 最终线性层:将
d_model维度的表征映射到目标词汇表大小(如英语词汇表大小);- Softmax 时机:训练时直接输出 logits(CrossEntropyLoss 更高效),预测时对输出做 Softmax 得到概率;
- 输入输出形状:
- 输入:src (batch_size, src_seq_len),tgt (batch_size, tgt_seq_len);
- 输出:(batch_size, tgt_seq_len, tgt_vocab_size)。
8.掩码生成 / 数据处理
核心思路
为模型提供必要的工具函数,比如生成掩码、数据加载 / 预处理等,保证模型能正常训练 / 推理。
核心代码用途
为注意力机制提供必要的掩码,保证模型训练 / 推理时的逻辑正确性(不关注无效 token、不提前看未来 token)。
def generate_square_subsequent_mask(sz: int) -> torch.Tensor: """生成Look-Ahead Mask(上三角矩阵,对角线以下为1,以上为0)""" mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1) mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0)) return mask def create_mask(src: torch.Tensor, tgt: torch.Tensor, pad_idx: int) -> Tuple[torch.Tensor, torch.Tensor]: """生成源序列和目标序列的掩码""" src_seq_len = src.size(1) tgt_seq_len = tgt.size(1) # 源序列Padding Mask:(batch_size, 1, 1, src_seq_len) src_mask = (src != pad_idx).unsqueeze(1).unsqueeze(2) # 目标序列Look-Ahead + Padding Mask:(batch_size, 1, tgt_seq_len, tgt_seq_len) tgt_mask = (tgt != pad_idx).unsqueeze(1).unsqueeze(2) tgt_mask = tgt_mask & generate_square_subsequent_mask(tgt_seq_len).to(device) return src_mask, tgt_mask关键要点:
generate_square_subsequent_mask:生成上三角掩码,遮挡未来位置;create_mask:结合 Padding Mask 和 Look-Ahead Mask,适配解码器的需求。
9.总结:
总结
Transformer 源码的核心要点可归纳为 3 点:
- 核心组件:位置编码注入序列顺序,Scaled Dot-Product Attention 实现注意力计算,Multi-Head Attention 提升表达能力,残差 + 层归一化保证深层训练稳定;
- 编解码结构:编码器堆叠实现源序列上下文提取,解码器堆叠结合自注意力 + 交叉注意力实现目标序列生成;
- 关键设计:掩码机制保证生成逻辑的合理性,Position-wise 前馈网络补充非线性表达,全注意力架构摆脱 RNN/CNN 的时序限制。
3.现代transformer的优化点
| 模块名称 | 经典方案 | 优化方向 | 代表技术 / 模型 | 核心优势 |
|---|---|---|---|---|
| 位置编码 | 正弦 / 余弦固定绝对位置编码 | 1. 相对位置编码2. 旋转位置编码(RoPE)3. 多段旋转编码(mrope)4. 无参数位置编码 | T5 相对偏置、LLaMA/Qwen(RoPE)、MinerU2.5(mrope)、ALiBi | 1. 支持超长序列外推(突破 2048 限制,达 16K+)2. 强化长距离相对位置感知3. 降低内存占用 |
| 注意力机制 | 全量 Scaled Dot-Product Attention(O (n²) 复杂度) | 1. 稀疏注意力(局部窗口 / 全局 + 局部混合 / 膨胀)2. 线性注意力(核函数分解 / 低秩近似)3. 硬件感知优化(显存 IO 优化) | Longformer、Performer、Linformer、Flash Attention、vLLM(Paged Attention) | 1. 复杂度降至 O (n) 或 O (kn),支持 100K + 长序列2. 推理速度提升 3-10 倍3. 解决 KV 缓存内存碎片问题 |
| 前馈网络(FFN) | 两层线性变换(d_model→d_ff→d_model)+ ReLU | 1. 低秩分解精简参数2. 激活函数升级(GELU/SwiGLU)3. 动态隐藏层维度4. 硬件并行优化 | GPT-3(GELU)、PaLM(SwiGLU)、低秩 FFN 变体 | 1. 参数量降低 50% 以上,保持表达能力2. 非线性表达提升 20%3. 适配 GPU 并行计算,提升吞吐量 |
| 整体架构 | 编码器 - 解码器对称架构 + Post-Norm | 1. 架构简化(Decoder-only)2. 归一化策略革新(Pre-Norm)3. 混合架构(Transformer+SSM) | GPT/LLaMA 系列、GPT-4(Pre-Norm)、Mamba/Nemotron-H | 1. 减少 50% 参数量,适配生成任务2. 支持 100 + 层深层模型训练,梯度传播更稳定3. 1M tokens 长序列推理速度提升 30 倍 |
| 工程化优化 | 基础训练 / 推理流程 | 1. 训练加速(Teacher Forcing / 量化感知训练 / 分布式并行)2. 推理优化(KV 缓存 / 推测解码 / 编译器优化)3. 硬件适配(Tensor Core/TPU 专用优化) | Megatron-LM(分布式训练)、Medusa(推测解码)、TensorRT/TVM | 1. 训练速度提升 5-10 倍,支持万亿参数模型2. 推理延迟降低 50% 以上,吞吐率提升 10 倍3. 充分利用专用硬件算力 |
表格核心说明
- 优化逻辑:所有优化围绕降复杂度、扩序列长度、提效率、适配场景四大目标,且多为 “算法创新 + 工程落地” 的协同优化。
- 优先级:工程化优化(如 Flash Attention、KV 缓存)是当前工业落地的核心,算法层面优化(如 RoPE、Decoder-only)是模型性能提升的基础。
- 趋势:注意力机制向 “稀疏 / 线性” 演进,架构向 “Decoder-only + 混合模块” 演进,工程化向 “编译器 + 专用硬件” 深度耦合演进。
总结:
本章,笔者详细的总结了transformer的代码详解和其相应的现代的优化点以及优化趋势,希望给大家带来帮助。