Transformer从输入到推理：Attention、训练与KV Cache全链路解析

Transformer 这几年被聊得太多了，但大部分讲解要么停在“Self-Attention 就是让每个词看其他词”这种一句话层面，要么直接甩公式把人劝退。我自己带过几个刚入行的同学，发现一个很普遍的现象：单独问 Attention 怎么算，能说个大概；但一问“训练的时候到底在算什么、推理的时候又省掉了什么、为什么同一个模型训练和推理表现不一样”，就卡住了。问题不在于公式没背熟，而在于没有把输入、Attention、训练、推理这四件事串成一条线。

这篇东西就是想把这条线捋直。我会从一条输入句子怎么一步步变成输出讲起，把 Attention 的矩阵运算、训练时的并行与损失、推理时的自回归与 KV Cache 全部串起来，中间穿插参数量估算、显存占用、常见报错排查这些实操层面的东西。适合已经能跑通 Transformer 代码、但对其内部机制还停留在“大概知道”阶段的读者，也适合想从 RNN 思维切换到 Attention 思维的从业者。看完你应该能自己回答：为什么训练能并行而推理不能、为什么 Attention 是 O(n²)、KV Cache 到底缓存了什么。

1. 从一条输入到输出：Transformer 的整体数据流

1.1 输入是怎么变成向量的

先别急着看 Attention，先把输入这条线走通。假设我们输入一句话“我 爱 学习”，模型第一步不是直接处理这三个字，而是先做Tokenization，把文本切成 token，再通过词表映射成整数 ID。比如“我”对应 1024，“爱”对应 2048，“学习”对应 4096（具体数字取决于词表，这里只是示意）。这一步的产物是一个长度为 3 的整数序列。

接下来是Embedding 层，把每个整数 ID 映射成一个 d_model 维的向量。d_model 是 Transformer 的核心超参之一，常见取值 512、768、1024。假设 d_model=512，那么这一步之后我们得到一个 3×512 的矩阵。你可以把它理解成：每个 token 变成了一个 512 维空间里的点，语义相近的词在这个空间里距离更近。

但光有词向量还不够，因为 Attention 本身是位置无关的——它只看 token 之间的关系，不看顺序。如果打乱输入顺序，Attention 算出来的结果几乎一样（除了位置编码）。所以必须加上位置编码（Positional Encoding）。原始论文用的是正弦余弦函数，现在很多模型改用可学习的位置嵌入或者 RoPE。加完之后，输入矩阵还是 3×512，但每个向量里已经混入了“我在第几个位置”的信息。

注意：位置编码不是可选项。我见过有人做文本分类时图省事去掉位置编码，结果模型在需要顺序信息的任务上直接崩掉，准确率掉十几个点。哪怕你的任务看起来“词袋就够”，也建议保留。

到这里，输入部分就完成了：整数序列 → Embedding → 加位置编码 → 得到一个 n×d_model 的矩阵，n 是序列长度。这个矩阵就是后面所有层的输入。

1.2 Encoder 和 Decoder 的分工

Transformer 原始结构是 Encoder-Decoder，但现在的用法分化很严重。Encoder-only（如 BERT 类）擅长理解任务，Decoder-only（如 GPT 类）擅长生成任务，Encoder-Decoder（如 T5 类）擅长序列到序列任务。不管哪种，核心积木都是Multi-Head Attention + Feed Forward Network，区别在于 Attention 的掩码方式和连接方式。

Encoder 的每一层做两件事：先做 Self-Attention，让每个位置看到所有位置；再过 FFN，对每个位置独立做非线性变换。注意 Encoder 里没有掩码，所有位置互相可见。Decoder 的每一层做三件事：Masked Self-Attention（只能看当前位置及之前）、Cross-Attention（看 Encoder 的输出）、FFN。Decoder-only 模型则去掉了 Cross-Attention，只保留 Masked Self-Attention。

这里有个容易混淆的点：Decoder 的 Masked Self-Attention 和 Encoder 的 Self-Attention 计算方式完全一样，只是多了一个上三角掩码，把未来位置的注意力分数设成负无穷，softmax 之后这些位置权重为 0。这个掩码是训练能并行、推理必须串行的根源，后面会详细讲。

1.3 残差连接和 LayerNorm 的位置

每个子层（Attention、FFN）外面都包了一层残差连接 + LayerNorm。原始论文是 Post-LN（先残差再加 LN），后来大家发现 Pre-LN（先 LN 再进子层，最后残差相加）训练更稳定，现在主流模型基本都用 Pre-LN。这个改动看起来小，但对深层模型能不能训起来影响巨大。

残差连接的作用是给梯度一条高速公路，让深层网络不至于梯度消失。LayerNorm 则是对每个 token 的向量做归一化，稳定训练。我实测下来，Pre-LN 在 12 层以上的模型里，学习率可以设得更大，收敛也更稳；Post-LN 则需要 warmup 很久，否则容易发散。

2. Attention 到底在算什么：从公式到矩阵形状

2.1 Q、K、V 的来历

Attention 的核心是三个矩阵：Query、Key、Value。它们都是从同一个输入矩阵 X（n×d_model）通过三个不同的线性变换得到的：

• Q = X · W_Q，形状 n×d_k
• K = X · W_K，形状 n×d_k
• V = X · W_V，形状 n×d_v

W_Q、W_K、W_V 是可学习参数。d_k 和 d_v 通常等于 d_model / num_heads。比如 d_model=512、num_heads=8，那么每个头的 d_k=d_v=64。

为什么要有 Q、K、V 三个而不是一个？你可以这样理解：Q 是“我在找什么”，K 是“我有什么可以被找到”，V 是“我实际提供的内容”。用 Q 和 K 算相似度得到权重，再用权重对 V 加权求和。这个设计让模型可以灵活地决定“关注谁”和“取什么信息”分开。

2.2 缩放点积注意力的计算过程

公式是：

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

拆开看每一步的形状变化（以单头、n=3、d_k=64 为例）：

1. Q·Kᵀ：3×64 乘 64×3，得到 3×3 的分数矩阵。第 i 行第 j 列表示第 i 个 token 对第 j 个 token 的注意力分数。
2. 除以 √d_k：防止点积结果过大导致 softmax 梯度消失。d_k=64 时 √d_k=8，相当于把分数缩小 8 倍。
3. softmax：对每一行做归一化，使每行和为 1，得到注意力权重。
4. 乘 V：3×3 乘 3×64，得到 3×64 的输出。

为什么要除以 √d_k？因为当 d_k 较大时，点积的方差会随 d_k 线性增长，softmax 会变得非常尖锐，梯度趋近于 0。除以 √d_k 把方差拉回 1 附近，保证梯度健康。这个细节很多人背了公式但没想过为什么，面试被问到就露馅。

2.3 多头注意力：为什么要拆成多个头

Multi-Head Attention 就是把上面的过程并行做 h 次，每次用不同的 W_Q、W_K、W_V，最后把 h 个头的输出拼接起来，再过一次线性变换 W_O。

拆多头的好处是：不同的头可以关注不同的模式。比如有的头关注语法依赖，有的头关注语义相似，有的头关注位置邻近。如果只用单头，这些模式会挤在同一个注意力分布里，互相干扰。我做过一个对比实验，同样参数量下，8 头的模型比 1 头的在长句理解任务上高 3-5 个点。

计算上，多头并不是真的把矩阵拆开算，而是把 d_model 维切成 h 份，用一次大的矩阵乘法完成，再 reshape。这样 GPU 利用率更高。

2.4 复杂度为什么是 O(n²)

注意力分数矩阵是 n×n，所以计算量和显存都是 O(n²)。n=512 时是 26 万，n=4096 时是 1600 万，n=32768 时是 10 亿。这就是长文本处理的瓶颈所在。

实际显存占用估算：假设 batch=8、heads=12、n=4096、d_k=64，单层注意力分数矩阵占用 8×12×4096×4096×4 字节 ≈ 6.4 GB。这还只是一层，12 层就是 77 GB，直接爆显存。所以长上下文模型必须用 FlashAttention、稀疏注意力、滑动窗口等优化手段。

提示：如果你在训练时遇到 OOM，先算一下注意力矩阵的显存，往往它就是元凶。降低 batch size 或序列长度是最直接的办法，但会牺牲吞吐，需要权衡。

3. 训练阶段：并行、损失与反向传播

3.1 训练为什么能并行

这是 Transformer 相比 RNN 最大的优势。RNN 必须按时间步串行，第 t 步依赖第 t-1 步的隐藏状态，无法并行。Transformer 的 Self-Attention 在训练时，因为有了 Mask，每个位置的计算只依赖输入矩阵和掩码，所有位置可以同时算。

具体来说，训练时我们把完整的目标序列（右移一位）一次性喂给 Decoder，Mask 保证第 i 个位置只能看到前 i-1 个位置。这样一次前向就能算出所有位置的输出和损失，GPU 可以吃满。这也是为什么 Transformer 能在大规模数据上训练——并行度决定了吞吐。

3.2 损失函数与标签对齐

训练时用的是交叉熵损失，对每个位置的输出做 softmax 得到词表上的概率分布，然后和真实的下一个 token 算交叉熵。注意标签是输入序列右移一位：输入 [BOS, 我, 爱, 学习]，标签是 [我, 爱, 学习, EOS]。

这里有个细节：padding 位置要 mask 掉，不参与损失计算。否则模型会学着去预测 padding，浪费容量。实现上用一个 loss_mask，把 padding 位置的损失置零，最后求平均时除以有效 token 数而不是总长度。

3.3 反向传播与梯度问题

反向传播沿着计算图回传，更新所有参数。Transformer 训练中常见的梯度问题有两个：梯度消失和梯度爆炸。梯度消失靠残差连接和 Pre-LN 缓解；梯度爆炸靠梯度裁剪（gradient clipping），通常设 max_norm=1.0。

学习率调度也很关键。原始论文用 warmup + inverse sqrt decay：前几千步线性升温，之后按步数平方根倒数衰减。现在很多模型用 cosine decay 或带 warmup 的 cosine。我实测下来，warmup 步数设总步数的 1%-5% 比较稳，太少容易早期发散，太多浪费训练时间。

3.4 参数量与显存估算

以 d_model=512、layers=6、heads=8、词表 30000 为例，粗略估算：

训练显存除了参数，还要存梯度、优化器状态（Adam 是参数的 2 倍）、激活值。激活值往往是大头，和 batch、序列长度成正比。经验公式：训练显存 ≈ 参数量×4（fp32）+ 激活值。激活值可以用 gradient checkpointing 换时间省显存，代价是前向多算一次。

4. 推理阶段：自回归与 KV Cache

4.1 推理为什么必须串行

训练能并行，推理却不行，原因在于推理时没有真实标签，只能一个一个生成。生成第 t 个 token 时，输入是前 t-1 个 token，模型输出第 t 个 token 的概率分布，采样得到 token 后拼到输入里，再生成下一个。这个过程天然串行，无法像训练那样一次算完。

这就是所谓的自回归生成。每一步都要做一次完整前向，序列越长，步数越多，总延迟越高。这也是为什么推理优化（KV Cache、批处理、投机解码）这么重要。

4.2 KV Cache 缓存了什么

如果每生成一个 token 都重新算所有位置的 K、V，那计算量会爆炸。KV Cache 的思路是：已经算过的位置的 K、V 不会变，缓存起来，新 token 只需要算自己的 Q、K、V，然后用自己的 Q 去和缓存的 K 算注意力。

具体来说，生成第 t 个 token 时：

• 新 token 的 Q：1×d_k
• 缓存的 K：t×d_k，缓存的 V：t×d_v
• 注意力分数：1×t
• 输出：1×d_v

这样每步的计算量从 O(t²) 降到 O(t)，总生成 n 个 token 的复杂度从 O(n³) 降到 O(n²)。KV Cache 的显存占用是 2×layers×heads×n×d_k×batch×精度字节数。n=4096、layers=32、heads=32、d_k=128、batch=1、fp16 时，约 2×32×32×4096×128×2 ≈ 2.1 GB。长上下文时这个数字很可观。

注意：KV Cache 是推理显存的主要占用之一。如果并发请求多，batch 一大，显存很快见底。生产环境常用 PagedAttention 之类的技术来管理 KV Cache 的显存碎片。

4.3 采样策略：贪心、束搜索与温度

推理时怎么从概率分布里选 token，直接影响生成质量：

• 贪心：每步选概率最大的，简单快，但容易重复、缺乏多样性。
• 束搜索：保留 top-k 条候选路径，最后选整体概率最高的。质量通常更好，但计算量是 k 倍。
• 温度采样：把 logits 除以温度 T 再 softmax。T<1 更确定，T>1 更随机。T=0 等价于贪心。
• top-k / top-p：只从概率最高的 k 个或累计概率 p 的集合里采样，过滤掉低概率的垃圾 token。

我一般对话任务用 top-p=0.9 + 温度 0.7，代码生成用温度 0.2 甚至贪心。这些参数没有万能值，得按任务调。

4.4 训练与推理的差异排查

很多人遇到“训练时 loss 降得很好，推理结果却很差”的情况。常见原因：

1. 训练/推理不一致：比如训练用了 dropout，推理忘了 eval 模式，dropout 还在生效。
2. 位置编码外推问题：训练长度 512，推理长度 1024，位置编码没做过外推训练，性能骤降。
3. KV Cache 实现 bug：缓存了错误的 K、V，或者位置索引对不上。
4. 采样参数不当：温度太高导致胡言乱语，太低导致重复。

排查时先固定随机种子、用贪心解码，排除采样干扰；再对比训练和推理的输入预处理是否一致；最后检查位置编码和 Cache 逻辑。

5. 常见问题与排查速查

5.1 训练不收敛怎么办

先看 loss 曲线。如果 loss 一开始就 NaN，多半是学习率太大或数据里有异常值，先把学习率降 10 倍试试。如果 loss 震荡不降，检查 warmup 是否够、batch size 是否太小、LayerNorm 位置是否合适。如果 loss 降了又升，可能是过拟合，加 dropout 或 weight decay。

我踩过的一个坑：数据里混入了空序列，导致 attention 全被 mask，softmax 输出 NaN，整个 batch 的 loss 变 NaN。后来在数据管道里加了长度过滤才好。

5.2 显存不够怎么优化

按优先级排序：

1. 降低 batch size 或序列长度（最直接）
2. 用混合精度（fp16/bf16），显存减半
3. 开 gradient checkpointing，用时间换显存
4. 用梯度累积模拟大 batch
5. 换更省显存的注意力实现（如 FlashAttention）
6. 分布式训练（数据并行、张量并行）

5.3 推理速度慢怎么加速

• 上 KV Cache（必做）
• 批处理多个请求
• 用更快的注意力 kernel
• 量化（int8/int4），牺牲一点精度换速度
• 投机解码：用小模型草稿、大模型验证

5.4 常见报错速查表

6. 把整条链路串起来的几个关键认知

走到这里，输入、Attention、训练、推理四块应该能连成一条线了。我最后再补几个容易被忽略但很关键的认知点。

第一，Attention 的本质是加权平均，权重由 Q 和 K 的相似度决定。理解了这一点，就能理解为什么它位置无关、为什么需要位置编码、为什么复杂度是 O(n²)。

第二，训练和推理用的是同一套参数，但计算图完全不同。训练是 teacher forcing + 并行 + 全序列损失；推理是自回归 + 串行 + 单步输出。KV Cache 是连接两者的桥梁——它把训练时算过的 K、V 复用起来，避免推理时重复计算。

第三，所有优化都围绕两个瓶颈：O(n²) 的注意力和 O(n) 的串行步数。长上下文优化针对前者，推理加速针对后者。想清楚你的瓶颈在哪，再选优化手段，不要盲目堆技术。

第四，参数量不等于能力，数据质量和训练策略往往更关键。我见过参数量翻倍但效果没提升的模型，也见过小模型靠更好的数据和调度打赢大模型。调参之前先把数据管道和评估做扎实。

最后分享一个我自己的习惯：每次改模型结构或训练配置，先在一个小数据集上跑通全流程，确认 loss 能降、推理能出合理结果，再上大规模。这样能省下大量试错时间，也能快速定位是结构问题还是数据问题。