从零实现Transformer：用PyTorch代码逆向理解注意力机制

很多开发者都有过这样的困惑：明明看了无数篇讲解Transformer的文章，公式推导也能看懂，但真要自己动手实现时却无从下手。这种"一看就懂，一写就懵"的困境，根源在于我们习惯了被动接受知识，而缺乏从代码层面逆向拆解原理的实践。本文将带你用PyTorch从零实现一个简化版Transformer，在编写forward函数、调试张量形状的过程中，让Q/K/V矩阵、多头注意力、残差连接等抽象概念变得触手可及。

1. 环境准备与基础架构

1.1 初始化项目环境

在开始编码前，我们需要配置好开发环境。建议使用Python 3.8+和PyTorch 1.12+版本：

conda create -n transformer python=3.8
conda activate transformer
pip install torch torchvision torchaudio
pip install numpy matplotlib ipykernel

创建一个基础的Transformer类骨架：

import torch
import torch.nn as nn

class Transformer(nn.Module):
    def __init__(self, d_model=512, n_head=8, num_layers=6):
        super().__init__()
        self.d_model = d_model
        self.n_head = n_head
        self.num_layers = num_layers
        
    def forward(self, x):
        return x

1.2 理解输入输出维度

Transformer处理的是三维张量：[batch_size, seq_len, d_model]。举个例子，当处理一批32个句子，每个句子50个词，词向量维度512时，输入形状就是[32, 50, 512]。

# 示例输入
batch_size = 32
seq_len = 50
d_model = 512
x = torch.randn(batch_size, seq_len, d_model)
print(x.shape)  # torch.Size([32, 50, 512])

2. 实现核心注意力机制

2.1 缩放点积注意力

这是Transformer最核心的运算，公式为： $$ \text{Attention}(Q,K,V) = \text{softmax}(\frac{QK^T}{\sqrt{d_k}})V $$

PyTorch实现：

class ScaledDotProductAttention(nn.Module):
    def __init__(self, dropout=0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, q, k, v, mask=None):
        d_k = q.size(-1)  # 获取特征维度
        scores = torch.matmul(q, k.transpose(-2, -1)) / (d_k ** 0.5)
        
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
            
        attn = torch.softmax(scores, dim=-1)
        attn = self.dropout(attn)
        output = torch.matmul(attn, v)
        return output, attn

调试技巧：在开发过程中，可以使用以下代码检查注意力矩阵：

q = torch.randn(2, 5, 64)  # [batch, seq_len, d_model]
k = torch.randn(2, 5, 64)
v = torch.randn(2, 5, 64)
attn = ScaledDotProductAttention()
output, attn_weights = attn(q, k, v)
print(f"注意力权重形状: {attn_weights.shape}")  # 应为[2,5,5]

2.2 多头注意力机制

多头注意力将输入拆分为多个子空间，让模型在不同表示子空间中学习信息：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model=512, n_head=8, dropout=0.1):
        super().__init__()
        assert d_model % n_head == 0
        self.d_k = d_model // n_head
        self.n_head = n_head
        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.attention = ScaledDotProductAttention(dropout)
        self.fc = nn.Linear(d_model, d_model)
        
    def forward(self, q, k, v, mask=None):
        batch_size = q.size(0)
        
        # 线性变换
        q = self.w_q(q).view(batch_size, -1, self.n_head, self.d_k)
        k = self.w_k(k).view(batch_size, -1, self.n_head, self.d_k)
        v = self.w_v(v).view(batch_size, -1, self.n_head, self.d_k)
        
        # 转置以计算注意力
        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
        
        # 计算注意力
        scores, attn = self.attention(q, k, v, mask)
        
        # 拼接多头结果
        concat = scores.transpose(1, 2).contiguous() \
                     .view(batch_size, -1, self.n_head * self.d_k)
        
        output = self.fc(concat)
        return output, attn

常见问题排查：

• 如果遇到维度不匹配错误，检查view和transpose操作的顺序
• 确保最终输出的维度与输入一致（[batch, seq, d_model]）
• 使用torch.einsum可以更清晰地表达矩阵运算

3. 构建Transformer完整层

3.1 位置编码实现

由于Transformer没有递归结构，需要显式添加位置信息：

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, 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)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        return x + self.pe[:, :x.size(1)]

3.2 前馈网络

Transformer中的前馈网络是两个线性变换加ReLU激活：

class FeedForward(nn.Module):
    def __init__(self, d_model, d_ff=2048, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ff, d_model)
        
    def forward(self, x):
        x = self.dropout(torch.relu(self.linear1(x)))
        x = self.linear2(x)
        return x

3.3 编码器层整合

将多头注意力、前馈网络和残差连接组合起来：

class EncoderLayer(nn.Module):
    def __init__(self, d_model, n_head, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, n_head, dropout)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        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, mask=None):
        # 自注意力子层
        residual = x
        x, attn = self.self_attn(x, x, x, mask)
        x = self.norm1(residual + self.dropout1(x))
        
        # 前馈子层
        residual = x
        x = self.ffn(x)
        x = self.norm2(residual + self.dropout2(x))
        return x, attn

4. 训练与调试技巧

4.1 初始化策略

Transformer对参数初始化敏感，推荐使用Xavier初始化：

def initialize_weights(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_uniform_(m.weight)
        if m.bias is not None:
            nn.init.constant_(m.bias, 0)
    elif isinstance(m, nn.LayerNorm):
        nn.init.constant_(m.weight, 1)
        nn.init.constant_(m.bias, 0)

model.apply(initialize_weights)

4.2 学习率调度

使用带warmup的学习率调度器：

class WarmupScheduler:
    def __init__(self, optimizer, d_model, warmup_steps=4000):
        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

4.3 梯度裁剪

防止梯度爆炸：

torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

5. 可视化与理解

5.1 注意力权重可视化

理解模型关注的重点：

import matplotlib.pyplot as plt

def plot_attention(attention, sentence):
    fig = plt.figure(figsize=(10,10))
    plt.imshow(attention, cmap='hot', interpolation='nearest')
    plt.xticks(range(len(sentence)), sentence, rotation=90)
    plt.yticks(range(len(sentence)), sentence)
    plt.colorbar()
    plt.show()

5.2 张量形状检查

在forward方法中添加调试语句：

print(f"输入形状: {x.shape}")
x = self.self_attn(x, x, x)
print(f"注意力后形状: {x.shape}")
x = self.ffn(x)
print(f"前馈后形状: {x.shape}")

6. 性能优化技巧

6.1 内存优化

使用梯度检查点减少内存占用：

from torch.utils.checkpoint import checkpoint

def custom_forward(x):
    # 定义前向计算
    return x

x = checkpoint(custom_forward, x)

6.2 混合精度训练

加速训练过程：

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast():
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

7. 扩展功能实现

7.1 掩码处理

实现序列padding和未来位置掩码：

def create_padding_mask(seq, pad_idx):
    return (seq != pad_idx).unsqueeze(1).unsqueeze(2)  # [batch, 1, 1, seq_len]

def create_lookahead_mask(size):
    mask = torch.triu(torch.ones(size, size), diagonal=1)
    return mask == 1  # [seq_len, seq_len]

7.2 位置感知前馈网络

增强位置信息处理能力：

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super().__init__()
        self.conv1 = nn.Conv1d(d_model, d_ff, kernel_size=1)
        self.conv2 = nn.Conv1d(d_ff, d_model, kernel_size=1)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x):
        x = x.transpose(1, 2)  # [batch, d_model, seq_len]
        x = self.dropout(torch.relu(self.conv1(x)))
        x = self.conv2(x).transpose(1, 2)
        return x

在实现过程中，最让我印象深刻的是调试多头注意力机制时的维度转换问题。最初没有正确理解transpose和view操作的配合，导致注意力权重计算出现偏差。通过逐步打印每个步骤的张量形状，最终理清了维度变换的逻辑链条。这种从代码实践中获得的理解，远比单纯阅读论文公式要深刻得多。