ylhan/note_sync

Fork 0

Files

T

ylhan 8a4210f380 将 note 笔记库同步到 note_sync，AI_compiler 作为 submodule

2026-08-20 17:42:14 +08:00

15 KiB

Raw Blame History

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，多头潜在注意力）。

八、高效注意力变体

标准自注意力是 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 不只是一篇论文——
它是整个现代大语言模型时代的起点。

15 KiB Raw Blame History Unescape Escape