Annotated Transformer 源码解析：从零理解 Transformer 架构核心

1次阅读

共计 2480 个字符，预计需要花费 7 分钟才能阅读完成。

Transformer 架构由 Vaswani 等人在 2017 年的论文《Attention Is All You Need》中提出，彻底改变了自然语言处理领域。其核心思想是使用自注意力机制完全替代传统的循环神经网络（RNN）和卷积神经网络（CNN），解决了长距离依赖问题和并行计算的瓶颈。

自注意力机制的优势在于：

能够直接建模序列中任意两个位置的关系
计算过程高度可并行化
避免了 RNN 的梯度消失 / 爆炸问题

原始论文中的实现较为抽象，而 Annotated Transformer 是一个教育性质的实现，主要区别在于：

代码组织方式
论文实现：模块耦合度高，不利于理解
Annotated：模块化设计，每个组件独立清晰
可读性
论文实现：优化较多，代码较晦涩
Annotated：添加大量注释，变量命名更直观
教学目的
论文实现：追求最高效率
Annotated：强调可理解性，适当牺牲效率

多头注意力的计算可以分为以下步骤：

线性投影：将输入分别投影到 Q、K、V 空间
分头处理：将投影后的张量分割成多个头
缩放点积注意力计算
合并多头输出
最终线性投影

关键实现代码片段：

def attention(query, key, value, mask=None, dropout=None):
    """缩放点积注意力实现"""
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

Transformer 使用正弦余弦函数生成位置编码：

$$PE_{(pos,2i)} = sin(pos/10000^{2i/d_{model}})$$
$$PE_{(pos,2i+1)} = cos(pos/10000^{2i/d_{model}})$$

代码实现特点：

预先计算所有位置编码并缓存
支持超过训练时最大长度的位置插值
与输入 embedding 直接相加

前馈网络采用两层线性变换加 ReLU 激活的结构：

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.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):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

以下是关键组装步骤：

定义基础模块（注意力、前馈网络等）
实现编码器层和解码器层
堆叠多层构建完整模型
添加 embedding 和位置编码
实现 mask 生成逻辑

完整代码示例见项目仓库（附链接），主要结构如下：

class Transformer(nn.Module):
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(Transformer, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator

    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)

    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)

注意力矩阵大小：序列长度的平方
长序列处理方案：
局部注意力窗口
内存高效的注意力实现
梯度检查点技术

层归一化位置：Pre-LN vs Post-LN
学习率预热策略
梯度裁剪阈值

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    output = model(input)
    loss = criterion(output, target)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()