疯狂的SRE

深色模式

Transformer 架构与注意力机制 ​

摘要：本文面向需要理解"模型为什么这样算、成本为什么这样涨"的工程师。我们从《Attention Is All You Need》（Vaswani et al., NeurIPS 2017）出发，拆解缩放点积注意力、多头注意力、位置编码与编码器-解码器堆叠，并用一张架构图串起训练与推理的关键路径。所有公式与原始超参数均来自论文，便于对照；现代变体（RoPE、Pre-Norm、Decoder-only）单独标注。适用版本：原理通用。

核心概念：为什么是注意力而不是 RNN ​

在 Transformer 之前，序列建模靠 RNN/LSTM：必须按时间步逐个处理，token_t 的依赖要穿过 t-1 步的隐藏状态。这带来两个问题（论文原文）：

  1. 无法并行：训练时一个样本内的位置不能并行计算，浪费 GPU。
  2. 长程依赖衰减：梯度要穿过很多步才能连接远距离 token。

Transformer 彻底抛弃循环与卷积，只靠注意力让任意两个位置一步连通。代价是每层注意力对序列长度 n 是 O(n²·d) 的计算与内存——这是后来 FlashAttention、长上下文优化的全部出发点。

架构与原理 ​

1. 缩放点积注意力（Scaled Dot-Product Attention） ​

这是整个架构的基石。给定 Query(Q)、Key(K)、Value(V) 矩阵：

text
Attention(Q, K, V) = softmax( Q·Kᵀ / √d_k ) · V
  • 每个 token 的 Q 与所有 token 的 K 求点积，衡量"相关性"。
  • 除以 √d_k（key 维度）是为了防止维度增大时点积数值过大、把 softmax 推入梯度消失区——这是论文里一个具体且必要的工程细节。
  • softmax 把分数归一化为权重（每行和为 1）。
  • 用权重对 V 做加权求和，得到融合了全局信息的表示。

2. 多头注意力（Multi-Head Attention） ​

单一注意力只能学一种"相关性"。论文用 h 个头，在各自降维子空间并行注意力，再拼接、投影：

text
MultiHead(Q,K,V) = Concat(head₁,…,head_h)·W_O
head_i = Attention(Q·W_Qⁱ, K·W_Kⁱ, V·W_Vⁱ)

不同头会学到不同关系（如语法依存、指代、语义相似）。原始论文超参数：base 模型 d_model=512、h=8，故每头 d_k=d_v=64。

3. 位置编码（Positional Encoding） ​

注意力本身无视顺序（集合对称）。论文用固定的正弦/余弦位置编码加到嵌入上：

text
PE(pos, 2i)   = sin( pos / 10000^(2i/d_model) )
PE(pos, 2i+1) = cos( pos / 10000^(2i/d_model) )

注意：现代 LLM（如 LLaMA、Qwen、Mistral）已普遍改用 RoPE（旋转位置编码，Su et al. 2021, arXiv:2104.09864） 等相对位置方案，并多用 Pre-Norm（残差在子层之前归一化）。这些是重要演进，但"注意力 + FFN + 残差 + LayerNorm"的核心骨架不变。[版本相关]

4. 编码器-解码器堆叠 ​

论文结构（base）：6 层编码器 + 6 层解码器，每层含：

  • 编码器：多头自注意力 → 残差+LayerNorm → 逐位置前馈网络(FFN) → 残差+LayerNorm。
  • 解码器：带掩码的自注意力（防止看到未来 token，因果性）→ 交叉注意力（ attend 编码器输出）→ FFN。每个子层外包残差连接与 LayerNorm。

逐位置 FFN 是同参数、逐 token 独立的两层 MLP：FFN(x)=max(0, xW₁+b₁)W₂+b₂，提供注意力本身缺少的非线性。

生产实践：论文结果与今天的映射 ​

原始论文在 WMT 2014 英德翻译达 28.4 BLEU、英法 41.0 BLEU，在 8×P100 GPU 上训练（最佳结果约 3.5 天）。[版本相关：不同报告对耗时有 12 小时–3.5 天之说，取决于模型规模与配置，未实测]

现代大语言模型如何映射：

论文设计现代 LLM 做法说明
编码器-解码器多为 Decoder-only（GPT 系）自回归生成只需要掩码自注意力
正弦绝对位置RoPE / ALiBi更好的长度外推
Post-NormPre-Norm深层训练更稳
6+6 层数十~上百层规模放大

对工程师的意义

理解注意力是 O(n²) 的，就能理解为什么"长上下文"贵、为什么需要 FlashAttention、为什么 KV cache 随序列线性增长——这些都在 上下文窗口与长文本处理 展开。

操作步骤：用 PyTorch 手算一次注意力 ​

下面是一段教学用的缩放点积注意力实现（无 GPU 优化），帮助确认公式：

python
import torch
import torch.nn.functional as F

def scaled_dot_product_attention(q, k, v, mask=None):
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / d_k ** 0.5   # Q·Kᵀ / √d_k
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float("-inf"))  # 因果/填充掩码
    attn = F.softmax(scores, dim=-1)                 # 行归一化
    return attn @ v                                  # 加权求和

# 示例：3 个 token，每个 d_k=64
q = k = v = torch.randn(1, 3, 64)
out = scaled_dot_product_attention(q, k, v)
print(out.shape)   # torch.Size([1, 3, 64])

这只是原理演示

生产推理不要手写注意力。要用 FlashAttention-2/3、PagedAttention（vLLM）、或框架内置内核，否则内存与速度都无法接受。eager 实现会把 N×N 分数矩阵整个写进 HBM，128K 序列时直接 OOM。

验证 ​

bash
# 用 Hugging Face 加载一个小型模型，确认架构字段与论文一致
python - <<'PY'
from transformers import AutoConfig
cfg = AutoConfig.from_pretrained("bert-base-uncased")  # BERT 保留了 Encoder 结构
print("hidden_size:", cfg.hidden_size)      # 通常 768
print("num_heads:", cfg.num_attention_heads)
print("num_layers:", cfg.num_hidden_layers)
PY

回滚与清理 ​

  • 模型/库版本升级可能改变注意力实现（如 FlashAttention 2→3）。升级后在回归集上验证输出一致性（logits 或生成结果），再全量上线。
  • 自托管镜像用 digest 锁定，避免 transformers/vllm 版本漂移导致数值差异。

故障排查 ​

现象可能原因排查
OOM（长序列）eager 注意力 N×N 矩阵换 FlashAttention / 限制序列长度
输出无序/无视顺序位置编码缺失或错配检查 RoPE 配置与 base 值
训练不稳定残差/LayerNorm 顺序（Pre/Post）核对模型配置
跨长距丢失信息原始绝对位置编码外推差改用 RoPE + YaRN（见上下文窗口篇）

安全与合规 ​

注意力层面的攻击面

  • 提示注入：注意力会"平等"地关注系统指令与外部内容，无法天然区分可信度。把不可信文本放进上下文即可能被其主导生成。
  • 上下文投毒：长上下文中塞入恶意指令片段（尤其在"中间"区域），利用 Lost-in-the-Middle 绕过开头/结尾的对齐约束。
  • 缓解：系统指令与用户内容物理隔离、对检索内容做来源标注、关键决策不单凭长上下文。

成本与性能 ​

  • 注意力每层 FLOPs ≈ 2·n²·d_model（不含 FFN）。序列翻倍，注意力计算翻 4 倍，KV cache 翻倍。
  • 推理瓶颈常是显存带宽而非算力（decode 阶段逐 token 生成，每 token 要读一遍全部权重）。量化权重直接降带宽需求（见 模型量化基础）。
  • FlashAttention 不改数学结果，只把 softmax 分块在 SRAM 算，消除 O(n²) 的 HBM 流量，是长上下文推理的必选项。

参考资料 ​

基于 VitePress 构建 · 托管于阿里云 ECS

Copyright © 2026 CrazySRE