Transformer 深度解析

2017 年，Google 发表论文《Attention Is All You Need》，提出了 Transformer 架构。谁也没想到，这篇论文会彻底改变 AI 领域。

今天的 GPT、Claude、Gemini、Llama……几乎所有主流大语言模型，本质上都是 Transformer 的变体。

理解了 Transformer，就理解了现代 LLM 的 90%。

这篇文章我们会深入 Transformer 的每一个组件：自注意力、多头注意力、位置编码、前馈网络、层归一化……不仅讲"是什么"，更讲"为什么"。

这是技术深度最高的模块之一。我们会用代码演示核心计算，确保你不仅看懂公式，还能亲手实现。

Transformer 前传：RNN 的局限

在 Transformer 出现之前，序列任务（如翻译、文本生成）主要用 RNN（循环神经网络）和它的变体 LSTM、GRU。

RNN 是怎么工作的

RNN 的核心思想是"逐个处理"：输入序列一个词一个词地进入网络，每一步的输出包含之前所有词的信息。

RNN 看起来合理，但它有三个致命问题：

问题一：梯度消失，难以记忆长距离依赖

RNN 是"链式"结构，每一步的梯度要反向传播回第一步。多次相乘后，梯度会指数级衰减，变成几乎 0。

比如句子"我在 2010 年去了巴黎，……，那是我最喜欢的城市"——"城市"要指代"巴黎"，但中间隔了太多词，RNN 很难学到这种长距离依赖。

LSTM 和 GRU 缓解了这个问题，但没有彻底解决。

问题二：无法并行计算，训练慢

RNN 必须等第 t-1 步算完，才能算第 t 步。这意味着：

• 即使有 100 个 GPU，也只能一个词一个词地算
• 序列越长，训练时间越长
• 很难扩展到超大规模数据

问题三：位置靠前的信息容易被"覆盖"

RNN 的隐藏状态是一步步更新的，后面的信息会不断覆盖前面的。句子开头的重要信息，到最后可能已经"稀释"得差不多了。

Transformer 的出现，一次性解决了这三个问题。

自注意力机制（Self-Attention）

自注意力是 Transformer 的核心。它的思想很简单：每个词都要和句子里的所有词"交流"一下，看看谁对自己重要，然后按重要程度加权求和。

先看整体架构图：

Q、K、V 的直觉理解

自注意力用三个向量来描述每个词：

• Q（Query，查询）——"我在找谁？"——这个词想关注谁
• K（Key，键）——"我是谁？"——这个词的身份标识
• V（Value，值）——"我有什么？"——这个词的实际内容

计算过程是：

1. 每个词用 Q 去和所有词的 K 做"匹配"，得到注意力分数
2. 用 Softmax 把分数归一化，总和为 1
3. 用归一化后的分数加权求和所有词的 V

看图更清楚：

注意力分数计算公式

完整的公式是：

Attention(Q, K, V) = softmax( Q·Kᵀ / √dₖ ) · V

让我们用纯 Python 一步一步实现它：

运行这段代码，你会看到完整的计算过程：Q/K/V 生成、注意力分数计算、缩放、Softmax、最后加权输出。

为什么要除以 √dₖ？

这是一个重要的细节。假设 Q 和 K 的每个元素都是均值 0、方差 1 的随机变量，那么它们的点积的方差是 dₖ。

如果 dₖ 很大（比如 512），点积的值会很大，导致 Softmax 进入"饱和区"——梯度几乎为 0，训练不动。

除以 √dₖ 后，方差变回 1，Softmax 的梯度就健康了。

缩放点积是 Transformer 训练稳定的关键技巧之一。

多头注意力（Multi-Head Attention）

单头自注意力虽然强大，但有个问题：每个词只能"关注"一次。如果我们想让词同时关注多个不同方面呢？

比如在句子"动物园里，老虎追着兔子跑"中：

• "老虎"可能需要一个头关注"动物园"（地点）
• 另一个头关注"兔子"（动作对象）
• 第三个头关注"追"（动作）

多头注意力就是把 Q/K/V 拆分成多个"头"，每个头学不同的注意力模式，最后拼起来。

多头注意力的计算步骤

1. 把 Q/K/V 按最后一维拆分成 h 份（h 是头数）
2. 每个头独立做自注意力
3. 把所有头的输出拼起来
4. 过一个线性投影层得到最终输出

公式是：

MultiHead(Q, K, V) = Concat(head₁, head₂, ..., headₕ) · Wₒ
  其中 headᵢ = Attention(Q·Wqᵢ, K·Wkᵢ, V·Wvᵢ)

多头注意力的强大之处在于：不同的头会自动学到不同类型的注意力模式。

例如在翻译任务中：

• 有的头关注语法关系（主谓宾）
• 有的头关注指代关系（代词指的是谁）
• 有的头关注局部相邻词

这些模式不是人为设计的，是模型在训练中自动学到的。

位置编码（Positional Encoding）

自注意力有个问题：它是"词袋"模型——不关心词的顺序。"我爱你"和"你爱我"在自注意力看来可能是一样的。

但顺序对语言至关重要。我们需要给模型"注入"位置信息。

正弦位置编码

原始 Transformer 用的是正弦位置编码：

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

其中：

• pos 是位置（0, 1, 2, ...）
• i 是维度索引
• d_model 是模型维度

这个设计很巧妙：

• 每个位置有唯一的编码
• 相对位置可以用线性组合表示（sin(a+b) = sin a cos b + cos a sin b）
• 可以外推到比训练时更长的序列

RoPE：旋转位置编码

正弦位置编码虽然经典，但后来有了更好的方案。现在的主流大模型（如 LLaMA）用的是 RoPE（Rotary Positional Embedding）。

RoPE 的思想是：把位置信息"旋转"进 Query 和 Key 里。

具体做法是对 Q 和 K 的每一对维度做旋转：

[ q₁, q₂, q₃, q₄, ... ]
  ↓ 旋转
[ q₁·cos mθ - q₂·sin mθ,
  q₁·sin mθ + q₂·cos mθ,
  q₃·cos mθ' - q₄·sin mθ',
  q₃·sin mθ' + q₄·cos mθ',
  ... ]

RoPE 的优势：

• 自然包含相对位置信息
• 外推性好（长度超过训练时也能用）
• 效果通常比正弦位置编码好

ALiBi：按距离衰减注意力

另一个简单有效的方案是 ALiBi（Attention with Linear Biases）。

它不加位置编码，而是直接给注意力分数加一个偏置：离得越远，分数越低。

ALiBi 的优势是简单，而且对长度外推特别友好。

前馈网络（FFN）层

每个 Transformer 层里，注意力之后还有一个前馈网络（Feed-Forward Network）。

FFN 的结构

原始 Transformer 的 FFN 是：

FFN(x) = max(0, x·W₁ + b₁) · W₂ + b₂

也就是：线性层 → ReLU 激活 → 线性层。

通常中间维度比输入维度大（比如 d_model=512，中间维度=2048）。

ReLU vs SwiGLU

原始用的是 ReLU，但现在的主流是 SwiGLU：

SwiGLU(x) = (x·W₁ ⊗ sigmoid(x·V₁)) · W₂

其中 ⊗ 是逐元素相乘。

SwiGLU 在很多任务上表现比 ReLU 好，LLaMA、PaLM 等都用它。

为什么 FFN 重要？

• 注意力是"聚合信息"，FFN 是"加工信息"
• FFN 有大量参数，是模型"记忆"知识的主要地方
• 没有 FFN，Transformer 的表达能力会大大降低

层归一化（Layer Norm）

Transformer 里还有一个重要组件：层归一化（Layer Normalization）。

层归一化做什么

层归一化把每个样本的特征归一化到均值 0、方差 1：

LN(x) = γ ⊗ ((x - μ) / √σ²) + β

其中：

• μ 是 x 的均值，σ² 是 x 的方差
• γ（scale）和 β（shift）是可学习参数
• ⊗ 是逐元素相乘

Pre-Norm vs Post-Norm

层归一化有两种放法：

• Post-Norm：子层 → 残差 → 归一化（原始 Transformer）
• Pre-Norm：归一化 → 子层 → 残差（现在主流）

Pre-Norm 训练更稳定，现在的大模型几乎都用它。

RMS Norm

RMS Norm 是层归一化的简化版，只算均方根，不减均值：

RMS(x) = sqrt(mean(x²))
RMSNorm(x) = γ ⊗ (x / RMS(x))

RMS Norm 计算更快，效果差不多，LLaMA 等模型用它。

Encoder-Decoder vs Decoder-Only 架构

原始 Transformer 是 Encoder-Decoder 架构，但现在的大语言模型大多是 Decoder-Only。

Encoder-Decoder：用于翻译等任务

Encoder-Decoder 架构：

• Encoder：读入输入序列（如源语言句子），生成上下文表示
• Decoder：基于 Encoder 的输出生成目标序列（如翻译结果）

Decoder 里有两种注意力：

1. 掩码自注意力：只能看之前的词（不能看未来）
2. 交叉注意力：去关注 Encoder 的输出

适合的任务：机器翻译、语音识别、摘要生成（需要把输入"重写成"输出）。

Decoder-Only：GPT 系列，用于生成

Decoder-Only 架构只有 Decoder，没有 Encoder。它的特点：

• 只有掩码自注意力（因果注意力）
• 自回归生成：一个词一个词地输出
• 适合预训练 + 微调的范式

现在的 GPT、LLaMA、Claude、Gemini……几乎都是 Decoder-Only。

Encoder-Only：BERT 系列，用于理解

Encoder-Only 架构只有 Encoder。它的特点：

• 双向注意力：每个词能看到所有词
• 适合理解任务：分类、命名实体识别、提取式问答
• 用"掩码语言模型"（MLM）预训练

为什么 Decoder-Only 成为主流？因为它简单，而且在大规模预训练下，生成能力极强。"一刀切"反而比"专门设计"效果更好。

Flash Attention：让注意力计算又快又省内存

标准自注意力的时间和内存复杂度都是 O(n²)，序列长了之后（比如 n=8192），既慢又显存不够。

Flash Attention 是一个关键优化，它在 2022 年的出现，让长序列训练变得可行。

标准注意力的内存瓶颈

标准自注意力需要保存中间结果（注意力权重矩阵），用来做反向传播。对于 n=8192，这是 8192×8192 ≈ 6700 万个元素，光这个就占几十 GB 显存。

Flash Attention 的优化原理

Flash Attention 的核心思想：

1. 分块计算：把 Q/K/V 分成小块，一块一块算
2. 在线 Softmax：不保存完整的注意力矩阵，用数学技巧在反向时重新计算
3. 利用 SRAM：把分块放进 GPU 的高速缓存（SRAM），减少 HBM 访问

结果是：

• 内存复杂度从 O(n²) 降到 O(n)
• 速度快 2-4 倍
• 精度完全一样（数学上等价）

现在的大模型训练几乎都用 Flash Attention 或它的变种（Flash Attention v2、Memory-Efficient Attention 等）。

KV Cache：推理加速原理

训练优化完了，推理（生成）也需要优化。KV Cache 是推理加速的关键。

问题：自回归生成的重复计算

自回归生成时，每生成一个新词，都要重新算一遍之前所有词的 K 和 V——这很浪费。

比如生成第 100 个词时，前 99 个词的 K/V 之前已经算过了，不需要重算。

KV Cache 怎么做

KV Cache 的思想很简单：

• 缓存之前所有位置的 K 和 V
• 生成新词时，只算新词的 K/V，拼到缓存里
• 用完整的 K/V 做注意力

这样每步的计算量从 O(n²) 降到 O(n)，生成速度大幅提升。

KV Cache 有个小代价：需要额外显存存 K 和 V。对于 7B 模型、4k 上下文，KV Cache 大约占 1-2GB 显存——这是值得的。

还有更高级的优化：

• Multi-Query Attention：多个头共享 K/V
• Grouped Query Attention：分组共享 K/V

这些能进一步减少显存占用，提升推理速度。

小结：Transformer 完整流程

让我们把所有组件串起来，总结一下 Transformer 的完整流程。

一个 Decoder-Only Transformer 层

x → LayerNorm → Masked Multi-Head Attention → 残差连接 →
    → LayerNorm → FFN → 残差连接 → 输出

完整的 Decoder-Only 模型

1. 输入 token IDs → 词嵌入
2. 加位置编码
3. 通过 N 个 Transformer 层
4. 最终 LayerNorm
5. 线性层映射到 vocab 大小
6. Softmax 得到下一个词的概率

关键超参数

代码实现指引

如果你想从零实现一个简化的 Transformer，步骤是：

1. 实现层归一化
2. 实现自注意力（带掩码）
3. 实现多头注意力
4. 实现 FFN
5. 实现 Transformer 层
6. 实现位置编码
7. 组装完整模型
8. 实现训练循环

实际中，我们会用 PyTorch 等框架，而不是纯 Python。这里的纯 Python 实现是为了演示原理，让你理解每一步在做什么。

理解 Transformer 的最好方式是亲手实现一次。哪怕是简化版，也会让你对很多细节有全新认识。