共计 3328 个字符，预计需要花费 9 分钟才能阅读完成。
背景痛点：为什么需要 Transformer？
传统 RNN（如 LSTM）处理长序列时存在两个致命缺陷：
- 顺序计算瓶颈 ：必须逐个处理时间步，无法并行化，训练速度慢
- 长期依赖丢失 ：随着序列增长，早期信息在反向传播时梯度逐渐消失（Vanishing Gradient）
2017 年《Attention Is All You Need》论文提出的 Transformer 架构，通过 Self-Attention 机制完美解决了这两个问题：
- 任意两个位置的距离都是 O(1)（常数级计算复杂度）
- 所有时间步可并行计算
- 通过注意力权重显式建模词与词的关系
词嵌入：从静态到动态
传统方法的局限
Word2Vec/GloVe 等静态词嵌入（Static Embedding）存在明显不足：
- 同一个词在不同语境中含义不同（如『苹果』在水果 vs. 公司场景）
- 无法捕获位置信息（『猫抓老鼠』和『老鼠抓猫』的语义完全不同）
Transformer 的解决方案
Transformer 采用动态词嵌入（Dynamic Embedding）方案：
-
Token Embedding：与传统方法类似，用可训练的矩阵将词 ID 映射为向量
# PyTorch 实现示例 embedding = nn.Embedding(vocab_size, d_model) token_embeddings = embedding(input_ids) # [batch, seq_len, d_model] -
位置编码（Positional Encoding）：通过正弦函数生成位置信息
数学公式：
$$
PE_{(pos,2i)} = \sin(pos/10000^{2i/d_{model}}) \
PE_{(pos,2i+1)} = \cos(pos/10000^{2i/d_{model}})
$$
class PositionalEncoding(nn.Module):
def __init__(self, d_model: int, max_len: int = 5000):
super().__init__()
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)
pe[:, 1::2] = torch.cos(position * div_term)
self.register_buffer('pe', pe) # [max_len, d_model]
def forward(self, x: Tensor) -> Tensor:
return x + self.pe[:x.size(1)] # [batch, seq_len, d_model]核心实现：多头注意力机制
自注意力（Self-Attention）三要素
- 查询 - 键 - 值（QKV）：
- Query：当前词想要查询的信息
- Key：其他词的身份标识
-
Value：实际传递的信息
-
计算流程 ：
def scaled_dot_product_attention(q: Tensor, # [batch, heads, seq_len, d_k] k: Tensor, v: Tensor, mask: Optional[Tensor] = None ) -> Tuple[Tensor, Tensor]: attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k) # [batch, heads, seq_len, seq_len] if mask is not None: attn_scores = attn_scores.masked_fill(mask == 0, -1e9) attn_weights = F.softmax(attn_scores, dim=-1) return torch.matmul(attn_weights, v), attn_weights -
多头拼接（Multi-Head）：
class MultiHeadAttention(nn.Module): def __init__(self, d_model: int, num_heads: int): assert d_model % num_heads == 0 self.d_k = d_model // num_heads self.proj_q = nn.Linear(d_model, d_model) self.proj_k = nn.Linear(d_model, d_model) self.proj_v = nn.Linear(d_model, d_model) self.out_proj = nn.Linear(d_model, d_model) def forward(self, x: Tensor, mask: Tensor = None) -> Tensor: batch_size = x.size(0) # 分头处理 [batch, seq_len, d_model] -> [batch, heads, seq_len, d_k] q = self.proj_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2) k = self.proj_k(x).view(...) # 同上 v = self.proj_v(x).view(...) # 同上 attn_output, _ = scaled_dot_product_attention(q, k, v, mask) attn_output = attn_output.transpose(1, 2).contiguous() # [batch, seq_len, heads, d_k] return self.out_proj(attn_output.view(batch_size, -1, self.d_model))
残差连接与层归一化
防止深层网络梯度消失的关键设计：
class Sublayer(nn.Module):
def __init__(self, d_model: int, dropout: float = 0.1):
self.norm = nn.LayerNorm(d_model)
self.dropout = nn.Dropout(dropout)
def forward(self, x: Tensor, sublayer: nn.Module) -> Tensor:
"""先 Norm 再子层计算"""
return x + self.dropout(sublayer(self.norm(x)))避坑指南
高频错误 TOP3
- 注意力分数未缩放
- 症状：训练初期 loss 突然变为 NaN
- 原因：点积结果方差随维度增大而增长，softmax 进入饱和区导致梯度消失
-
修复：必须除以√d_k
-
掩码应用时机错误
- 正确做法：在 softmax 前将 padding 位置设为负无穷（-1e9）
-
错误示范：在 softmax 后应用会破坏概率分布性质
-
共享 QKV 矩阵的陷阱
- 实验现象：模型性能显著下降
- 原理分析：Query 和 Key 需要不同的语义空间
性能优化技巧
Flash Attention 原理
- 核心思想 ：通过分块计算避免频繁访问显存
- 实现效果 ：
- 内存占用从 O(N²) 降到 O(N)
- 速度提升 2 - 4 倍
- 适用条件 ：
- 需要 CUDA 硬件支持
- 序列长度 >512 时效果显著
# 使用示例（需安装 flash-attn 包）from flash_attn import flash_attention
output = flash_attention(q, k, v)延伸思考
开放实验建议
- 位置编码对比实验 ：
- 固定式（正弦）vs 可学习式（Learnable）
-
在 IWSLT 德语 - 英语翻译任务上的 BLEU 差异
-
注意力头数调优 ：
- 绘制头数（4/8/16）与验证集 perplexity 的关系曲线
- 观察不同头学习的注意力模式差异
结语
实现完整 Transformer 需要组合上述所有组件（编码器 / 解码器各 6 层）。建议从简化版开始：
- 先实现单层 Encoder
- 在文本分类任务验证正确性
- 逐步扩展到完整结构
完整代码已开源在 GitHub（伪链接）：
https://github.com/xxx/mini-transformer正文完
