Transformer 从零推导：注意力到底在算什么

前情说明：技术深度 D1。这是《大模型原理面试冲刺》系列第 1 篇。这个系列是我自己面试前的复盘笔记，目标是把每个概念用人话讲一遍，再补公式和代码，最后落到"面试官会怎么问"。如果你也在准备大模型岗，应该能省你一些时间。系列会覆盖 Transformer / GPT 演进 / 对齐 / Long Context / LLM Eval 五块地基——推理 Infra 部分我在这个博客其他文章里已经写得很深了，本系列不重复。

为什么第一篇必须写 Transformer

面试大模型岗，不管你面的是训练、推理还是应用，Transformer 是唯一保证会被问到的东西。我见过的问法从"手写一下 attention"到"为什么 QKV 要分开"，再到"MHA 的每一个矩阵形状你能不看代码写出来吗"——都是同一个东西的不同深度。

写完这篇你应该能做到：

白板上手写 scaled dot-product attention 的公式
讲清楚 Q/K/V 为什么要分开、为什么要除 √d_k
说明白 MHA 和单头 attention 的区别在哪、代价是什么
解释 causal mask、padding mask 分别在解决什么问题

不追求"把 Transformer 讲全"——这篇只讲 attention 这一块，positional encoding、FFN、LayerNorm 单独会有专门篇幅（或放在第 2 篇 GPT 演进里）。

一、动机：RNN 到底不够用在哪

在 attention 之前，处理序列的主力是 RNN/LSTM。它有两个致命问题：

问题 1：串行计算，GPU 用不起来

RNN 的 hidden state 是 h_t = f(h_{t-1}, x_t)，第 t 步必须等第 t-1 步算完才能开始。哪怕你有 8 张 A100，一个 batch 内的时间维也是串行的。这直接把训练效率钉死。

问题 2：长距离依赖会衰减

信息要从第 1 个 token 传到第 1000 个 token，必须经过 999 次矩阵乘 + 激活函数。梯度回传时同样要经过 999 次乘法，数值上要么爆炸要么消失。LSTM 的门控机制缓解了一些，但没根治。

Attention 的核心 insight 简单粗暴：任何两个位置之间都直接连一条边，信息从 1 到 1000 一步到位。代价是计算量从 O(n) 涨到 O(n²)——但矩阵乘刚好是 GPU 最擅长的事。

二、Attention 的直觉：一次带权检索

先忘掉公式。Attention 在做的事，本质是一次带权重的软性字典查询。

想象你在图书馆里查资料：

你手上有个问题（query）
书架上每本书有个索引卡（key）和内容（value）
你拿问题去和每张索引卡比对相似度，得到一个分布
按这个分布，把所有书的内容加权取回来

在 Transformer 里：

Q（Query）：当前 token 想问的问题——"我需要什么信息来更好地表示自己？"
K（Key）：每个 token 的索引——"我这里有什么可以被查的？"
V（Value）：每个 token 的内容——"如果你想要我，我实际给你的是什么"

Q、K、V 是同一个输入 x 通过三个不同的线性层投影出来的。为什么要分开？因为"问什么"、"我能匹配什么"、"我实际给什么"是三件不同的事，让模型自己学出三个不同的映射比强行共用一个要灵活得多。这也是为什么共享 QK 的变种（比如某些 low-rank 优化）通常掉点。

三、Scaled Dot-Product Attention：公式来了

数学上就一个式子：

Attention(Q, K, V) = softmax(QK^T / √d_k) · V

拆开看每一步在做什么：

Step 1: QK^T——算相似度

Q 的形状是 [n, d_k]（n 个 token，每个 d_k 维）。K 的形状一样。QK^T 得到 [n, n] 的矩阵——第 i 行第 j 列就是 token i 的 query 和 token j 的 key 的内积，数值越大表示越相关。

Step 2: / √d_k——归一化，防止 softmax 饱和

这一步是最容易被面试官深挖的。为什么除的是 √d_k 而不是 d_k 或者别的？

数学上，如果 Q 和 K 的每个分量都是均值 0、方差 1 的独立随机变量，那么它们的内积 q · k = Σ q_i · k_i 的方差是 d_k（方差可加）。方差随维度线性增长，意味着 d_k 越大，内积的数值绝对值倾向于越大。

softmax 对输入尺度非常敏感——如果输入里有一个数比其他数大很多（比如 20 vs 5），softmax 输出会变得极度 one-hot 化，反向传播时梯度几乎为 0（饱和区）。

除以 √d_k 后，内积的方差回到 1，softmax 保持在"软"的状态，梯度能正常流动。这一步的动机纯粹是训练稳定性，和表达能力无关。

Step 3: softmax(...)——变成概率分布

沿着最后一维做 softmax，得到的 [n, n] 矩阵每一行加起来是 1。这就是每个 query 对所有 key 的注意力权重。

Step 4: · V——加权求和

用注意力权重去加权 V。V 的形状是 [n, d_v]（通常 d_v = d_k），输出还是 [n, d_v]。

一句话总结：每个位置的输出，是所有位置的 V 按"当前位置的 Q 和各位置的 K 有多像"来加权平均。

四、代码：30 行讲清楚

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    Q, K, V: [batch, n_heads, seq_len, d_k]
    mask:   [batch, 1, seq_len, seq_len] or None
    """
    d_k = Q.size(-1)

    # [batch, n_heads, seq_len, seq_len]
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)

    # 有 mask 就把要屏蔽的位置设成 -inf，softmax 后变 0
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))

    attn = F.softmax(scores, dim=-1)         # 权重矩阵
    output = torch.matmul(attn, V)           # 加权求和
    return output, attn

面试如果要你手写，写到这一步就够拿到"能上手"的分。再往下就是 MHA 的包装。

五、Multi-Head Attention：一次不够，那就多来几次

单头 attention 有个直觉上的问题：每个位置对其他位置只有一种"关注模式"。但语言里，一个 token 可能同时需要关注句法关系、指代关系、语义相似性——这些是不同类型的信号，压在一个 attention 里彼此干扰。

MHA 的解法很朴素：把 Q/K/V 切成 h 份，每份独立算 attention，最后拼起来。

假设 hidden dim 是 d_model = 512，head 数 h = 8，那么每个 head 的 d_k = 64：

1. 输入 x 经过三个大线性层，得到 Q, K, V，形状都是 [n, 512]
2. reshape 成 [n, 8, 64]，转置成 [8, n, 64] —— 8 个头并行
3. 每个头独立算 attention，得到 8 个 [n, 64] 的输出
4. 拼回 [n, 512]，再过一个输出线性层 W_O

关键点：MHA 的总计算量和单头 d_model = 512 是一样的——只是把计算切成 8 份并行。参数量也差不多（多了个 W_O）。所以 MHA 是"免费的午餐"，几乎所有实现都用 MHA 不用单头。

代码骨架：

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.d_k = d_model // n_heads
        self.n_heads = n_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        B, N, _ = x.shape
        # 投影 + 切多头: [B, N, d_model] -> [B, n_heads, N, d_k]
        Q = self.W_q(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        K = self.W_k(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)
        V = self.W_v(x).view(B, N, self.n_heads, self.d_k).transpose(1, 2)

        out, _ = scaled_dot_product_attention(Q, K, V, mask)

        # 拼回来: [B, n_heads, N, d_k] -> [B, N, d_model]
        out = out.transpose(1, 2).contiguous().view(B, N, -1)
        return self.W_o(out)

六、两种 Mask：causal 和 padding

面试时"attention mask"是个高频词，但很多人分不清是哪种 mask。

Causal Mask（因果掩码）

只有解码器（decoder，GPT 类模型）才需要。作用：让第 i 个 token 只能看到 1..i，看不到 i+1..n。

为什么？因为 GPT 是自回归的——训练时用 teacher forcing，一次性把整个序列喂进去，但每个位置的预测必须只依赖它之前的 token，否则就是"抄答案"。

实现是个下三角矩阵：

[[1, 0, 0, 0],
 [1, 1, 0, 0],
 [1, 1, 1, 0],
 [1, 1, 1, 1]]

上三角部分（未来位置）在 softmax 前被设成 -inf，softmax 后变 0。

Padding Mask（填充掩码）

batch 里不同样本长度不同，短的要 pad 到统一长度。pad 出来的 token 是无意义的，不应该被真实 token attend 到。

padding mask 就是把 pad 位置的 key 屏蔽掉——所有 query 看 pad 的 key 时都是 -inf。

注意：padding mask 是屏蔽 key（列），causal mask 是屏蔽整个未来子矩阵。两者经常同时存在，实现时把两个 mask 逐元素相与。

七、复杂度：那个众所周知的 O(n²)

Attention 的时间和空间复杂度都是 O(n² · d)。这个 n² 就是 Transformer 长上下文的核心痛点：

序列 1k：attention 矩阵 1M 元素
序列 32k：attention 矩阵 1G 元素（单头，单精度就是 4GB）
序列 128k：16G 元素——单头就爆显存

长上下文的所有优化——FlashAttention、Ring Attention、Linear Attention 变种——本质都是在跟这个 O(n²) 掰手腕。这也是为什么我这个博客里 FlashAttention 单独写了一篇《FlashAttention 三代演进》——kernel 侧怎么把 O(n²) 的内存墙拆掉。本系列第 7 篇会从模型侧讲 Long Context 外推的另一条路。

八、总结：一张图

Attention 你需要在脑子里刻死的一张图：

输入 x  [n, d_model]
   ↓
  W_q      W_k      W_v         ← 三个线性投影
   ↓        ↓        ↓
   Q [n,d]  K [n,d]  V [n,d]
   ↓        ↓
   Q · K^T / √d_k    [n, n]     ← 相似度矩阵
        ↓
     softmax          [n, n]    ← 归一化成概率
        ↓
        · V           [n, d]    ← 加权求和
        ↓
       W_o                       ← 输出投影
        ↓
   输出 [n, d_model]

MHA 就是把中间的 d 切成 h × d_k 并行做 h 遍。

面试锚点：大概率会被问到的问题

以下问题按频次从高到低排，每题给一个"及格线答案"，你自己再往深处扩展。

Q1: 为什么 attention 要除以 √d_k？

A: 因为 Q 和 K 的点积方差随维度 d_k 线性增长，数值绝对值变大会让 softmax 饱和（输出几乎 one-hot），梯度趋近 0，训练不稳定。除以 √d_k 把方差归一化回 1，让 softmax 保持在"软"的状态。这是纯粹的训练稳定性优化，不影响表达能力。

Q2: Q/K/V 为什么要用三个不同的投影矩阵？能不能共用？

A: 三个投影对应三件不同的事——"我要问什么"、"我能被匹配什么"、"我要给出什么"。共用会限制表达能力。实际中有些优化（比如 MQA、GQA）会共享 K 和 V 减少 KV cache，但 Q 一般不动。共享 QK 的方案通常会掉点。

Q3: MHA 和单头 d_model 大小的 attention，计算量一样吗？

A: 一样。MHA 是把 d_model 切成 h 份并行，总的 FLOPs 和参数量基本一致（只多一个 W_O）。好处是能让模型同时关注多种不同的模式。

Q4: Causal mask 和 padding mask 有什么区别？

A: causal mask 保证自回归性质（第 i 位只能看 1..i），只在 decoder 里用，是下三角矩阵。padding mask 是屏蔽 batch 内被 pad 出来的无效位置，encoder/decoder 都要用。实现时两者可以叠加。

Q5: Attention 的复杂度是多少？为什么是这个复杂度？

A: O(n² · d)。因为要算 Q · K^T，这是 [n, d] × [d, n] 的矩阵乘，结果是 [n, n]，所以 n² · d 次乘加。空间上 attention 矩阵本身就是 O(n²)，这是长上下文的核心瓶颈。

Q6（进阶）: FlashAttention 优化了什么？

A: FlashAttention 没有降低理论复杂度，它优化的是显存访问。标准 attention 会把 [n, n] 的中间矩阵完整写回 HBM（高带宽显存），再读回来做 softmax，再写回，IO 开销很大。FlashAttention 用 tiling + online softmax 的技巧，让整个 attention 在 SRAM 里流式完成，不落 HBM。所以它是 IO-aware 的实现，不是新算法。这是 attention 优化的现代默认方案，几乎所有推理框架都在用。想深挖看我另一篇《FlashAttention 三代演进》。

Q7（进阶，和推理面试强相关）: 推理阶段的 attention 和训练阶段有什么区别？

A: 训练是并行的——一次性输入整段 sequence，attention 矩阵一次算完。推理是自回归的——decode 阶段每次只生成一个 token，新 token 的 Q 只需要和历史所有 token 的 K/V 做 attention。如果每步都重算 K/V 会浪费，所以有 KV Cache：把历史 K/V 存起来，decode 时只算新 token 的 Q，和 cache 里的 K/V 做 attention。这也是为什么 KV Cache 是推理显存的头号大户——推理侧的展开我写在了《现代 LLM 推理引擎剖析》和《vLLM 三年演进史》里。

一句话记住 Transformer 的 attention

Attention 是一次可微的、带权重的软性字典查询——每个 token 的输出是所有 token 的加权平均，权重来自 query 和 key 的相似度。

MHA 就是并行多次这个过程，让模型能同时看多种关系。剩下的 FFN、LayerNorm、residual 都是常规组件——真正让 Transformer 成为 Transformer 的就是这个 O(n²) 的全连接式注意力。

系列下一篇：《从 GPT-1 到 GPT-4：架构演进与 Scaling Law》——为什么 Decoder-only 赢了，Chinchilla 说了什么，以及"更大真的更好吗"。

Transformer 从零推导：注意力到底在算什么

Transformer 从零推导：注意力到底在算什么

为什么第一篇必须写 Transformer

一、动机：RNN 到底不够用在哪

二、Attention 的直觉：一次带权检索

三、Scaled Dot-Product Attention：公式来了

四、代码：30 行讲清楚

五、Multi-Head Attention：一次不够，那就多来几次

六、两种 Mask：causal 和 padding

Causal Mask（因果掩码）

Padding Mask（填充掩码）

七、复杂度：那个众所周知的 O(n²)

八、总结：一张图

面试锚点：大概率会被问到的问题

一句话记住 Transformer 的 attention

评论