426 lines
15 KiB
Markdown
426 lines
15 KiB
Markdown
# Transformer 架构详解
|
||
|
||
> 2017 年，"Attention Is All You Need" 论文用自注意力取代了 RNN/LSTM，开启了现代大语言模型时代。
|
||
|
||
---
|
||
|
||
## 一、Transformer 之前：RNN 的困境
|
||
|
||
```
|
||
RNN 的处理方式（串行）：
|
||
Token₁ → Token₂ → Token₃ → ... → Tokenₙ
|
||
5ms 5ms 5ms 5ms = 5×n ms（无法并行）
|
||
|
||
Transformer 的处理方式（并行）：
|
||
Token₁ ↘
|
||
Token₂ → 所有 Token 同时互相"看到" → 并行计算
|
||
Token₃ ↗ ≈ 5ms（与 n 无关）
|
||
```
|
||
|
||
| 问题 | RNN/LSTM | Transformer |
|
||
|---|---|---|
|
||
| **并行性** | 串行，无法并行训练 | 完全并行，GPU 友好 |
|
||
| **长距离依赖** | 梯度消失/爆炸，远距离信息丢失 | 每个 Token 直接关注任意位置的 Token |
|
||
| **训练速度** | 慢（序列越长越慢） | 快得多 |
|
||
| **可解释性** | 难以解释隐藏状态 | 可直接可视化注意力权重 |
|
||
|
||
---
|
||
|
||
## 二、整体架构：编码器-解码器
|
||
|
||
原始 Transformer 用于机器翻译，包含编码器和解码器两部分。
|
||
|
||
```
|
||
┌─────────────────────────────────┐
|
||
│ 编码器（Encoder） │
|
||
│ ┌───────────────────────────┐ │
|
||
"我 爱 北京" → │ │ Encoder Block × 6 │ │
|
||
│ │ Self-Attn + FFN │ │
|
||
│ └───────────────────────────┘ │
|
||
└──────────────┬──────────────────┘
|
||
│ 编码后的表示（K, V）
|
||
▼
|
||
┌─────────────────────────────────┐
|
||
│ 解码器（Decoder） │
|
||
│ ┌───────────────────────────┐ │
|
||
│ │ Decoder Block × 6 │ │
|
||
│ │ ① Masked Self-Attn │ │ → "I love Beijing"
|
||
│ │ ② Cross-Attn (用 K,V) │ │
|
||
│ │ ③ FFN │ │
|
||
│ └───────────────────────────┘ │
|
||
└─────────────────────────────────┘
|
||
```
|
||
|
||
### 一个 Block 内部（原始设计）
|
||
|
||
```
|
||
输入 x
|
||
│
|
||
├──→ Multi-Head Self-Attention ──→ ⊕ ──→ LayerNorm ──→
|
||
│ ↑ ↑
|
||
│ 残差连接 x 输出
|
||
│
|
||
└──→ Feed-Forward Network ──→ ⊕ ──→ LayerNorm ──→ 输出
|
||
↑ ↑
|
||
残差连接 传给下一层
|
||
```
|
||
|
||
### 两种主流归一化位置
|
||
|
||
| 方式 | 归一化位置 | 代表模型 | 特点 |
|
||
|---|---|---|---|
|
||
| **Post-LN** | 残差之后 | 原始 Transformer | 训练不稳定，需要 warmup |
|
||
| **Pre-LN** | 子层之前 | GPT 系列、LLaMA | 训练更稳定，无需大量 warmup |
|
||
|
||
```
|
||
Post-LN：x + Attn(LayerNorm(x)) ← 原始论文
|
||
Pre-LN： x + Attn(LayerNorm(x)) 错 — 实际上是：
|
||
Pre-LN：LayerNorm(x + Attn(x)) 错 — Pre-LN 是：
|
||
|
||
Pre-LN 实际：x + Attn(LayerNorm(x)) — 等等，让我用图表示清楚：
|
||
|
||
Post-LN： x → SubLayer → Norm → + → 输出
|
||
Pre-LN： x → Norm → SubLayer → + → 输出
|
||
```
|
||
|
||
> Pre-LN 先将输入归一化再送入子层，然后做残差。GPT/LLaMA 全系使用 Pre-LN（或 RMSNorm + Pre-LN 变体）。
|
||
|
||
---
|
||
|
||
## 三、自注意力机制：完整推演
|
||
|
||
### 3.1 从"查字典"理解 Q/K/V
|
||
|
||
| 投影 | 类比 | 数学 | 含义 |
|
||
|---|---|---|---|
|
||
| **Q（Query）** | 你的搜索关键词 | `X · W_Q` | "我在找什么信息？" |
|
||
| **K（Key）** | 所有条目的索引 | `X · W_K` | "每个词身上有什么标签？" |
|
||
| **V（Value）** | 条目的实际内容 | `X · W_V` | "每个词的实际信息是什么？" |
|
||
|
||
### 3.2 逐步计算示例
|
||
|
||
假设输入句子 "我 爱 NLP"，模型维度 d=4，头维度 d_k=2：
|
||
|
||
**Step 1 — 输入嵌入**
|
||
```
|
||
我: [0.1, 0.3, 0.5, 0.2]
|
||
爱: [0.4, 0.1, 0.3, 0.6]
|
||
NLP: [0.2, 0.5, 0.1, 0.3]
|
||
```
|
||
|
||
**Step 2 — 计算 K 和 Q（只演示一个头的投影）**
|
||
```
|
||
K = X · W_K （3×4）×（4×2） = 3×2 矩阵
|
||
K_我 = [0.8, 0.2]
|
||
K_爱 = [0.3, 0.9]
|
||
K_NLP = [0.5, 0.4]
|
||
|
||
Q_爱 = [0.6, 0.3] （"爱"这个位置的 Query）
|
||
```
|
||
|
||
**Step 3 — 计算注意力分数（Q_爱 × K^T）**
|
||
```
|
||
Q_爱 · K_我 = 0.6×0.8 + 0.3×0.2 = 0.48 + 0.06 = 0.54
|
||
Q_爱 · K_爱 = 0.6×0.3 + 0.3×0.9 = 0.18 + 0.27 = 0.45
|
||
Q_爱 · K_NLP = 0.6×0.5 + 0.3×0.4 = 0.30 + 0.12 = 0.42
|
||
|
||
scores = [0.54, 0.45, 0.42]
|
||
```
|
||
|
||
**Step 4 — 缩放 + Softmax**
|
||
```
|
||
scaled = scores / √2 = [0.38, 0.32, 0.30]
|
||
weights = softmax(scaled) = [0.35, 0.33, 0.32]
|
||
```
|
||
|
||
**Step 5 — 加权求和 V**
|
||
```
|
||
V_我 = [0.7, 0.1]
|
||
V_爱 = [0.2, 0.8]
|
||
V_NLP = [0.4, 0.5]
|
||
|
||
output_爱 = 0.35×[0.7,0.1] + 0.33×[0.2,0.8] + 0.32×[0.4,0.5]
|
||
= [0.245, 0.035] + [0.066, 0.264] + [0.128, 0.160]
|
||
= [0.439, 0.459]
|
||
```
|
||
|
||
> 位置"爱"的输出融合了"我"(35%)、"NLP"(32%)和自身(33%)的信息——这正是"理解上下文"的本质。
|
||
|
||
### 3.3 矩阵形式（实际实现）
|
||
|
||
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V$$
|
||
|
||
- **QK^T**：N×N 的注意力矩阵，每对 Token 一个分数
|
||
- **÷ √d_k**：缩放因子。d_k 越大点积值越大 → softmax 趋向 one-hot → 梯度消失。除以 √d_k 保持方差 ≈ 1
|
||
- **Softmax**：沿最后维（Key 维度）归一化，每行和为 1
|
||
- **× V**：用注意力权重对 Value 加权求和
|
||
|
||
---
|
||
|
||
## 四、多头注意力：并行多个"视角"
|
||
|
||
### 4.1 为什么需要多头？
|
||
|
||
```
|
||
单头可能只学到一种注意力模式：
|
||
"我" → 关注了 "爱"
|
||
"爱" → 关注了 "北京"
|
||
所有头都只关注相邻词 ← 浪费了注意力的表达能力
|
||
|
||
多头允许并行学习多种模式：
|
||
头 1：语法结构（主语→谓语）
|
||
头 2：长距离共指（"它"→ 前文的"北京"）
|
||
头 3：语义关联（"医生"→ "医院"→ "病人"）
|
||
头 4：位置模式（局部上下文窗口）
|
||
```
|
||
|
||
### 4.2 数学公式
|
||
|
||
$$\text{head}_i = \text{Attention}(Q W_i^Q, K W_i^K, V W_i^V)$$
|
||
|
||
$$\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h) W^O$$
|
||
|
||
每个头有自己独立的投影矩阵 W_Q、W_K、W_V，维度为 d_model × d_k（其中 d_k = d_model / h）。
|
||
|
||
### 4.3 计算量并不是 h 倍！
|
||
|
||
| | 单头 (d_model=512) | 8头 (每头 d_k=64) |
|
||
|---|---|---|
|
||
| Q投影参数 | 512×512 = 262K | 8×512×64 = 262K |
|
||
| K投影参数 | 同上 | 同上 |
|
||
| V投影参数 | 同上 | 同上 |
|
||
| 总参数 | ≈ 786K | ≈ 786K + W_O(512×512) ≈ 1M |
|
||
|
||
> 计算量大致相同，只是**把工作分给多个头并行做**，每个头在更小的空间（d_k）中关注不同的模式。
|
||
|
||
### 4.4 生产实践中的头数配置
|
||
|
||
| 模型 | d_model | 头数 | d_k |
|
||
|---|---|---|---|
|
||
| 原始 Transformer | 512 | 8 | 64 |
|
||
| GPT-3 | 12288 | 96 | 128 |
|
||
| LLaMA 70B | 8192 | 64 | 128 |
|
||
| DeepSeek-V3 | — | — | 使用 MLA 替代标准多头 |
|
||
|
||
---
|
||
|
||
## 五、位置编码：告诉模型"顺序"
|
||
|
||
注意力机制本身是**置换不变**的——打乱输入顺序，注意力分数矩阵只是行列交换，输出值不变。所以必须显式注入位置信息。
|
||
|
||
### 方案对比
|
||
|
||
| 方案 | 原理 | 优点 | 缺点 | 代表模型 |
|
||
|---|---|---|---|---|
|
||
| **Sinusoidal（正弦）** | 不同频率的正弦/余弦波叠加 | 可外推、无参 | 固定模式 | 原始 Transformer |
|
||
| **Learned** | 可学习的 pos_embed 表 | 灵活 | 不能超出训练长度 | BERT、ViT |
|
||
| **RoPE（旋转位置编码）** | 用旋转矩阵编码相对位置 | 天然编码相对位置、可外推 | 长序列扩展需技巧 | LLaMA、Qwen、DeepSeek |
|
||
| **ALiBi** | 给注意力分数加线性偏置 | 简单、外推最强 | 表达能力有限 | BLOOM |
|
||
|
||
### RoPE 核心思想
|
||
|
||
> **不是往 embedding 里加位置向量，而是把 Q 和 K "旋转"一个角度，角度大小取决于位置。**
|
||
|
||
```
|
||
位置 m 的 Q：旋转 m×θ
|
||
位置 n 的 K：旋转 n×θ
|
||
|
||
Q·K = (rotated_Q_m)·(rotated_K_n) = 关于 (m-n) 的函数
|
||
↑
|
||
只依赖相对位置！
|
||
```
|
||
|
||
RoPE 扩展：
|
||
- **NTK-aware scaling**：调高基频（如 10000→500000），让旋转更慢，支持更长上下文
|
||
- **YaRN**：NTK + 额外微调，将 Llama 从 4K 扩展到 128K
|
||
|
||
---
|
||
|
||
## 六、三种架构变体：选哪种？
|
||
|
||
### 6.1 编码器-only（如 BERT）
|
||
|
||
```
|
||
输入："我 [MASK] 北京" → 编码器 → 每个位置的上下文表示 → 预测 [MASK] = "爱"
|
||
|
||
特点：
|
||
- 双向注意力（能看到整个句子）
|
||
- 适合：分类、实体识别、语义搜索
|
||
- 不能做生成（没有解码器）
|
||
```
|
||
|
||
### 6.2 解码器-only（如 GPT、LLaMA、DeepSeek V4）
|
||
|
||
```
|
||
输入："今天天气" → 因果注意力（只看左边）→ 预测下一个 Token "真"
|
||
"今天天气真" → ... → 预测 "好"
|
||
|
||
特点：
|
||
- 因果掩码（Causal Mask）：每个 Token 只能看自己和左边
|
||
- 适合：对话、代码生成、创意写作
|
||
- 自回归生成，输入=输出（训练时位移一位即可）
|
||
- 现代 LLM 几乎全部使用此方案
|
||
```
|
||
|
||
**因果掩码示意**：
|
||
|
||
```
|
||
注意力分数矩阵（× = 被 mask 为 -∞）：
|
||
|
||
我 爱 北京 天安门
|
||
我 ✓ × × × ← "我"只能看到自己
|
||
爱 ✓ ✓ × × ← "爱"能看到"我"和自己
|
||
北京 ✓ ✓ ✓ × ← ...
|
||
天安门✓ ✓ ✓ ✓
|
||
```
|
||
|
||
### 6.3 编码器-解码器（如 T5）
|
||
|
||
```
|
||
编码器（双向）：处理源语言 → 生成 K、V
|
||
解码器（因果 + 交叉注意力）：用 K、V 辅助生成目标语言
|
||
|
||
交叉注意力（Cross-Attention）：
|
||
解码器的 Q ← 来自解码器自身（"这句话接下来该说什么？"）
|
||
解码器的 K,V ← 来自编码器（"原文说了什么？"）
|
||
```
|
||
|
||
| 架构 | 自注意力方式 | 能否生成 | 代表 | 适合任务 |
|
||
|---|---|---|---|---|
|
||
| 编码器-only | 双向 | ❌ | BERT | 理解/分类 |
|
||
| 解码器-only | 因果（单向） | ✅ | GPT, LLaMA | 生成/对话 |
|
||
| 编码器-解码器 | 双向 + 交叉 | ✅ | T5, BART | 翻译/摘要 |
|
||
|
||
---
|
||
|
||
## 七、训练与推理
|
||
|
||
### 7.1 训练：Teacher Forcing
|
||
|
||
```
|
||
输入： "今天 天气 真" 目标： "天气 真 好"
|
||
↑
|
||
模型被训练来预测每个位置的下一个 Token
|
||
|
||
实际一次性计算：
|
||
输入 [今天, 天气, 真]
|
||
→ 因果掩码确保"今天"看不到后面
|
||
→ 输出 [P(天气), P(真), P(好)]
|
||
→ Loss = CrossEntropy(预测, 实际目标)
|
||
```
|
||
|
||
> 训练时**一次前向**算出所有位置的预测，因为目标已知、且有因果掩码保护。
|
||
|
||
### 7.2 推理：自回归生成
|
||
|
||
```
|
||
训练时可以并行（目标已知），推理必须串行（目标未知）：
|
||
|
||
Step 1: "今天" → 预测 "天气"
|
||
Step 2: "今天天气" → 预测 "真"
|
||
Step 3: "今天天气真" → 预测 "好"
|
||
Step 4: "今天天气真好" → 预测 <eos>
|
||
|
||
每次都要重新计算全序列？→ 不，用 KV Cache！
|
||
```
|
||
|
||
### 7.3 KV Cache：推理加速的关键
|
||
|
||
```
|
||
未使用 KV Cache：
|
||
Step 1：计算 "今天" 的 Q、K、V ──────────── 1个Token
|
||
Step 2：计算 "今天天气" 的 Q、K、V ──────── 2个Token ← 重新算了"今天"！
|
||
Step 3：计算 "今天天气真" 的 Q、K、V ────── 3个Token ← 又重算了！
|
||
...
|
||
|
||
使用 KV Cache：
|
||
Step 1：计算 "今天" → 缓存 K₁、V₁
|
||
Step 2：只算 "天气" 的 Q、K、V + 用缓存的 K₁、V₁ → 缓存 K₁K₂、V₁V₂
|
||
Step 3：只算 "真" 的 Q、K、V + 用缓存的 K₁K₂、V₁V₂ → ...
|
||
...
|
||
每个 Step 只算 1 个新 Token！开销从 O(n²) 降到 O(n)（单步）。
|
||
```
|
||
|
||
> KV Cache 的代价：显存占用随序列长度线性增长。MLA（多头潜在注意力）通过低秩压缩解决此问题——详见 [[MoE 混合专家架构#MLA（Multi-head Latent Attention，多头潜在注意力）|MLA 详解]]。
|
||
|
||
---
|
||
|
||
## 八、高效注意力变体
|
||
|
||
标准自注意力是 **O(n²)** 复杂度（n = 序列长度）。对长序列需优化：
|
||
|
||
### 8.1 FlashAttention（事实标准）
|
||
|
||
> **IO 感知的精确注意力算法**——数学结果完全等价于标准注意力，但通过算子融合大幅减少显存读写。
|
||
|
||
```
|
||
标准计算： FlashAttention：
|
||
QK^T → 写回 HBM 将 Q、K、V 分块（tile）
|
||
Softmax → 写回 HBM 每个块在 SRAM 中完成全部计算
|
||
乘 V → 写回 HBM 只有最终结果写回 HBM
|
||
↑ 3次 HBM 读写 ↑ 约 1/n 的 HBM 读写
|
||
```
|
||
|
||
实测效果：相比朴素实现，吞吐量提升 **~32%**（A100），现为 PyTorch 2.0+ 默认后端。
|
||
|
||
### 8.2 稀疏/线性注意力（长序列场景）
|
||
|
||
| 方案 | 策略 | 复杂度 | 代表 |
|
||
|---|---|---|---|
|
||
| **Longformer** | 局部滑动窗口 + 全局 Token | O(n) | 文档级 NLP |
|
||
| **BigBird** | 局部 + 全局 + 随机 | O(n) | — |
|
||
| **Performer** | 随机特征近似 softmax | O(n) | — |
|
||
| **Mamba/SSM** | 状态空间模型替代注意力 | O(n) | 超长序列 |
|
||
|
||
> 经验法则：序列 < 4K Token → 标准注意力 + FlashAttention 即可。长文档/代码库场景才需稀疏注意力或 SSM。
|
||
|
||
---
|
||
|
||
## 九、现代生产优化总览
|
||
|
||
| 优化手段 | 效果 | 代价 |
|
||
|---|---|---|
|
||
| **混合精度（bf16/fp16）** | 训练/推理速度翻倍 | 质量损失可忽略 |
|
||
| **梯度检查点** | 显存降低 50%+ | 训练慢 20% |
|
||
| **FlashAttention** | 推理加速 ~30%，显存降低 | 无（数学等价） |
|
||
| **算子融合（RMSNorm, xFormers）** | 减少 kernel launch 开销 | 无 |
|
||
| **序列分桶** | 减少 padding 浪费 | 实现复杂度增加 |
|
||
| **量化（INT8/INT4）** | 显存降低 4-8× | 轻微质量下降 |
|
||
|
||
---
|
||
|
||
## 十、总结
|
||
|
||
```
|
||
Transformer 的核心思想就三条：
|
||
|
||
1. 自注意力（Self-Attention）
|
||
Q·K 算关联 → Softmax 归一化 → ×V 合成上下文
|
||
每个 Token 直接"看到"所有 Token
|
||
|
||
2. 多头（Multi-Head）
|
||
并行跑多个注意力，每个头学不同的关系模式
|
||
总计算量不变，表达能力翻倍
|
||
|
||
3. 位置编码（Positional Encoding）
|
||
注意力不感知顺序 → 注入位置信息
|
||
RoPE 是当前主流：Q/K 旋转角度 = 位置函数
|
||
|
||
Transformer 不只是一篇论文——
|
||
它是整个现代大语言模型时代的起点。
|
||
```
|
||
|
||
---
|
||
|
||
## 相关链接
|
||
|
||
- [[大模型架构基础]] — 大模型的各组件总览
|
||
- [[FFN 前馈网络详解]] — Transformer 中参数量最大的组件
|
||
- [[MoE 混合专家架构]] — 如何用稀疏专家替代 FFN
|
||
- [[多模态大模型]] — Transformer 如何扩展到图像/音频
|
||
- [[DeepSeek V4 系列]] — Transformer + MoE + MLA 的实际应用
|
||
|
||
---
|
||
|
||
*最后更新：2026-08-05*
|