写这一篇笔记之前，我已经把 Transformer 的基础结构翻来覆去看过不下三遍：原始论文、各类图解、开源实现里的代码，每一步都能讲得头头是道。但真正动手去写一个可训练的小模型时才发现，结构图离“能跑起来”之间，隔着大量没人明说的细节。

上一篇笔记整理了整体架构和 attention 的基本定义，这篇换个角度：不再重复“什么是注意力”，而是把 QKV 的计算、缩放因子的由来、位置编码怎么选、归一化到底放哪、训练参数怎么配、loss 不降怎么查这些实操层面的坑一个个摊开。适合已经看过结构图、正准备动手复现小模型的读者；如果你连 scaled dot-product attention 都还没搞清，建议先回去把乘法过程看明白，再来看这篇。

1. 从“看懂结构图”到“写对注意力代码”

1.1 QKV 的计算到底在做什么

很多人记住了softmax(QK^T / sqrt(d_k)) V这个式子，写代码时却会犯一个很直觉的错误：把 Q、K、V 当成三份独立的输入去处理。实际上在标准自注意力里，Q、K、V 全部来自同一个输入序列 X，只是分别乘上三个不同的权重矩阵W_Q、W_K、W_V得到的。

这三次线性投影的本质，是对同一份信息做三次“视角切换”。Q 负责表达“我在找什么”，K 负责表达“我身上带有什么可被检索的属性”，V 负责表达“如果别人选中了我，我实际贡献的内容是什么”。点积就是在做一次检索匹配：拿 Q 去和所有 K 做相似度比较，相似度越高，注意力权重越大，最后把权重作为系数去加权求和 V。

一个容易忽略的细节是：W_Q、W_K、W_V这三个矩阵的形状。如果模型维度是d_model，输入序列长度是n，那么 X 的形状是[n, d_model]，三个权重矩阵都是[d_model, d_model]，输出 Q、K、V 也都是[n, d_model]。在多头上，这个投影会被拆成h份，每份维度d_k = d_model / h，于是形状就变成[n, h, d_k]，再转置成[h, n, d_k]去做批量点积。

我自己习惯在写代码前先把每个张量的形状写在注释里，不然 debug 的时候注意力全在 shape 报错上，根本看不清模型在学什么。这一步看似笨，但真的能帮我省掉大半低级错误。

1.2 缩放因子为什么是 sqrt(d_k)

这个问题值得单独拿出来说。不少初学者会觉得sqrt(d_k)只是一个让数值好算一点的“魔法常数”，但它的来源非常具体。

假设 Q 和 K 中的每个元素都是均值 0、方差 1 的独立随机变量，那么点积结果就是d_k个这样的乘积之和，方差会变成d_k。也就是说，维度越大，点积结果的方差越大，数值分布越分散。

如果不缩放，直接把这个大方差的结果丢进 softmax，会发生什么？以d_k=64为例，点积的标准差是 8，这意味着有些点积值可能落在 ±16 甚至更大的区间。softmax 对输入非常敏感，输入里只要有一个值稍微偏大，它的指数项就会碾压其他项，导致注意力权重退化成近似 one-hot。训练早期，模型参数还没有学到合理分布，这种“假自信”的注意力会带来非常陡峭的梯度，让训练很不稳定。

把点积除以sqrt(d_k)，相当于把方差重新拉回 1 的量级，softmax 的输入分布趋于平滑，梯度也就健康了。代码实现上建议直接用q @ k.transpose(-2, -1) * (d_k ** -0.5)，注意缩放因子乘在 softmax 之前、对 softmax 的输出没有影响，很多人容易顺手把它写到 softmax 后面，结果等于白写。

1.3 掩码的两种形态，别搞混了

注意力代码里最容易被忽略的就是 Mask。项目里训练文本模型时，至少存在两种形态完全不同的掩码，用途也不同：

• Padding Mask：由于 batch 内序列长度不一，短序列会用 pad token 补齐到 batch 最大长度。这些 pad 位置没有任何语义，注意力不应该分配到它们身上。实现时把它加进 softmax 前的分数里，用-inf（或一个绝对值很大的负数，比如-1e9）填充。
• Causal Mask：也叫自回归掩码，确保第 i 个位置只能看到 0 到 i 的位置，不能看到未来。形状是一个下三角全 true 的矩阵。

实际写代码时，这两者是同时使用的。做法是先把两种 mask 合并成一个combined_mask，比如causal_mask & padding_mask，再统一加到注意力分数上。很多新手只加了 causal mask，忘了 padding mask，导致 pad token 被“认真学习”了一遍，最终效果会明显变差。

另一个容易踩的点是：-inf参与 softmax 后，对应位置的概率是 0，这没问题；但如果在实现里用了一个较大的负有限数如-1e9，在混合精度训练下可能因为浮点范围问题变成 NaN，所以直接用float('-inf')最稳妥。

2. 多头注意力：拆开再拼回去的细节

2.1 头数怎么定，投影矩阵怎么用

多头注意力的核心逻辑是：不让模型只用一种“注意力关系”去看序列，而是让不同的头学会关注不同范围、不同类型的关系。比如一个头专注于局部相邻词的依赖，另一个头关注整个句子的长距离核心词。

工程实现上有两套做法。第一套是先做整体投影再拆分：输入先经过[d_model, d_model]的W_Q得到完整 Q，再用view拆成 h 个头，每个头维度d_k。第二套是直接一次性完成所有头的拆分，把W_Q设计成[d_model, h * d_k]，一次投影后 reshape。两种方式数学上等价，区别只在代码组织习惯。

头数怎么选？经验上d_model=512时配 8 个头是常见的，d_model=768时配 12 个头也常见。核心约束是d_model必须能被 h 整除，且d_k不要太小，不然每个头能表达的信息太有限。我试过把d_model=256强行拆 16 个头，每个头只有 16 维，结果模型训起来明显吃力，原因就是头的表达空间被压缩得太狠。

多头合并时，各个头的输出先按头拼接回d_model，再乘一个输出投影矩阵W_O。这个W_O常常被忽略，但它非常重要：它负责把多个头的信息做交叉融合，而不是简单地把各头结果堆在一起。

这里有个实际经验：如果做的是小规模文本生成，可以直接把所有头的注意力矩阵打印出来看看。你会发现一部分头确实在学语法结构，另一部分头在学词义搭配，这种可视化能帮你理解模型到底在做什么，比盯着 loss 曲线直观得多。

2.2 训练和推理时注意力的行为差异

训练时，Transformer 是并行计算的：整个序列一次性输入，所有位置的注意力一次性算完。因为在训练阶段要充分利用 batch 内的并行度，causal mask 只是把未来位置遮住，但计算本身并不真正“一个个地往后推”。

推理时就不一样了。生成第 t 个 token 时，理论上可以重新把前 t 个 token 全部送入模型重新计算注意力，但这样做会做大量重复计算：每生成一个 token 都要把前面所有 token 的 QKV 重新算一遍，复杂度是 O(n^2) 的二次增长。

这就是 KV Cache 存在的理由。推理时，前 t-1 个 token 的 K 和 V 已经算过了，把它们缓存下来，生成第 t 个 token 时只需要计算新 token 的 Q、K、V，再拿新 Q 去和缓存的全部 K 做匹配，新 V 再和权重加权求和，就能得到当前步的输出。Q 不需要缓存，因为每一步的 Q 都是新 token 自己的。

KV Cache 带来的显存开销也不小。缓存大小和层数、头数、序列长度成正比，长上下文场景下这个缓存很容易成为显存瓶颈。印象里做 2k 长度对话时，KV Cache 大概能占到总显存的 30%-40%，所以后面讲到的很多优化手段，本质上都是在和 KV Cache 抢显存。

2.3 KV Cache 的实现与踩坑

KV Cache 在开源推理框架里已经是标配，但自己手写过一个简化版之后，才真正理解它为什么要这样设计。

简单实现思路：初始化一个空缓存字典，每个层存两个张量k_cache和v_cache，形状是[batch, num_heads, past_len, d_k]。生成时每来一个新 token，把当前的 K、V 拼接到缓存尾部，注意力分数计算时 Q 的长度是 1，K 的长度是past_len + 1，这样就完成了跨时间的匹配。

踩坑的地方主要有三处：

• 缓存长度要一次性预留：如果每步cat一个单 token，会反复触发内存搬移，速度慢得离谱。正确做法是预分配一个max_seq_len长度的张量，用一个kv_len指针标记当前写到哪。
• int 形状和负值索引：拼接时如果是动态长度，务必保持缓存张量的 dtype 和 device 一致，否则在 GPU 上会频繁触发设备同步，性能瞬间打回原形。
• beam search 与 batch 的维度：beam search 会把 batch 维度翻倍，KV Cache 也需要相应复制或者按 beam 管理，很多框架里的 bug 都出在这套维度变换上。

印象最深的一次 debug 是：生成结果看起来“越来越慢”，单步延迟从 20ms 慢慢涨到 200ms。查了一圈，发现是缓存每次都在用torch.cat拼接，每步都在重新申请内存。改成预分配之后，单步延迟基本恒定，问题立刻消失。

3. 位置编码：模型凭什么知道顺序

3.1 绝对位置编码的局限

transformer 本身没有循环结构，输入顺序对它而言只是“一组集合的排列”。如果不加位置信息，把“我爱你”写成“你爱我”，得到的表示几乎一样，这对于语言任务来说是完全不可接受的。

原始论文用的是正弦余弦绝对位置编码，每个位置生成一个固定的向量，加到 token 的 embedding 上。这种编码有两个特点：一是不同频率的分量让模型有机会学到不同尺度的位置关系，二是理论上可以外推到训练时没见过的更长序列。

但绝对位置编码有一个明显的局限性：它把位置当作一个绝对坐标，模型需要自己去学“位置 3 和位置 7 的相对关系”。位置 100 和位置 101 的关系，和位置 3 和位置 4 的关系，编码形式上并不直接等价。对于长度外推和相对位置建模，绝对编码并不友好。

后来出现了可学习位置编码，把位置向量作为可训练参数直接学。这种方法在固定长度任务上表现很好，但一旦序列长度超过训练时的最大长度，就只能截断或者强行插值，外推能力几乎为零。做固定长度分类任务无所谓，做生成任务就会比较难受。

3.2 RoPE 和 ALiBi 的核心思路

近几年的实践里，旋转位置编码（RoPE）和注意力线性偏置（ALiBi）逐渐成为主流，它们解决的关键问题就是“相对位置关系”和“长度外推”。

RoPE 的思路很巧妙：它不把位置信息加到 embedding 上，而是通过旋转 Q 和 K 的向量来实现。具体地说，对位置 m 的 Q 向量，按一定频率做二维旋转变换；对位置 n 的 K 向量也做同样变换。这样 Q_m 和 K_n 的点积结果里会自然出现一个与 (m - n) 相关的项，也就是相对位置信息直接编码进了注意力分数里。

数学上 RoPE 的实现会让很多初学者看晕，但代码实现其实很固定：先预处理一组频率theta_i = 10000^(-2i/d_k)，然后对向量做“两两成组”的旋转变换。实际代码可以预先算好 cos/sin 表，推理时查表再乘上去。

ALiBi 则更直接：它完全不改 QKV 的计算，而是在 softmax 之前的分数上加一个和距离成正比的线性偏置。距离越远，偏置越负，注意力分数越低。实现极其简单，几乎不增加任何计算量，却天然带有相对位置信息，而且外推能力非常强。

我自己对比过这两种方案在小规模文本生成任务上的表现：RoPE 在“训练长度范围内的任务”上通常上限更高，模型学习长依赖关系的能力更强；ALiBi 则在“外推到更长序列”的时候明显占优，几乎不会因为序列变长而崩掉。

3.3 实际选型：我在实验里的取舍

选位置编码，本质是在“表达能力上限”和“外推鲁棒性”之间做权衡。

如果任务相对固定，序列长度比较稳定，比如做固定长度文本分类，那么可学习位置编码足够简单好用，训练也快。

如果任务是文本生成、对话、需要处理比训练时更长的输入，我建议优先选 RoPE。现在大量的开源模型和推理框架都对 RoPE 做了高度优化，算力开销很小，roformer 思路已经非常成熟。RoPE 还支持一种“长度插值”技巧：当序列长度超出训练范围时，把位置下标除以一个缩放系数再映射回去，可以在不平滑性能断崖的情况下把上下文窗口扩大几倍。这个技巧我在项目里用得非常频繁。

ALiBi 适合那种对外推要求极高、追求稳定性的场景。它几乎没有额外参数，实现也不会出错，适合快速实验和做 baseline 对比。

我个人的经验是：新项目默认选 RoPE，除非明确知道要做超长文本外推，才考虑 ALiBi。做对比实验时一定要固定其他所有超参，只切换位置编码，不然两个方案的效果差异会被学习率等变量掩盖。

4. 归一化与 FFN：容易被忽略的两个组件

4.1 Pre-LN 和 Post-LN，放的位置差很多

原始 Transformer 用的是 Post-LN：先做注意力（或 FFN），再加残差，最后做 LayerNorm。这个排列方式在原始论文里效果不错，但深层模型很容易训练不稳，原因在于残差路径上的梯度流要穿过多层的 LayerNorm 反向传播，范数容易退化或爆炸。

Pre-LN 把顺序反过来了：先做 LayerNorm，再做注意力/FFN，最后加残差。这样每条残差路径都直接跨层传递梯度，深层网络的训练稳定性好很多，所以现在绝大多数开源模型都是 Pre-LN 结构。

但 Pre-LN 不是没有代价。训练稳定不代表最终效果一定更好，很多实验表明 Pre-LN 在深层模型上容易“欠拟合”一部分能力，尤其是深层表征的学习效率略低。所以实际中有一种做法是在 Post-LN 结构里对每一层残差做缩放，或者对最后一层 LayerNorm 的输出做特殊初始化，来兼顾稳定和效果。

我自己重实现时通常直接选 Pre-LN，因为它“无脑稳定”，省去很多调参时间。如果你追求极致性能、资源充足，可以仔细对比两种排列在你数据上的差异，但别默认 Post-LN 是标准答案——现代实现里它反而不是默认选项。

4.2 RMSNorm 为什么越来越流行

LayerNorm 做的事情是：对每个 token 的特征向量，减去均值、除以标准差，再乘可学习的缩放系数加偏置。这套操作能稳定训练，但有一个问题：计算均值这一步要额外做一次全局归约，在长序列、大 batch 下有一定开销。

RMSNorm 做了一个简化：只除以 RMS（均方根），不减均值。它认为 LayerNorm 里的 centering 对训练稳定性贡献不大，核心价值在于 rescale。省掉均值计算后，速度有一定提升，而且在实践中和 LayerNorm 效果基本持平甚至略优。

现在开源模型用 RMSNorm 已经很普遍。实现时通常用x * torch.rsqrt(x.pow(2).mean(-1) + eps)，注意加了 eps 是为了防止除零，位置在根号里面。一个容易踩的坑是 eps 太小，在混合精度下会发生数值下溢，建议保持默认的 1e-6 或更大。

RMSNorm 会在模型每一层重复计算，所以它的一点性能优化都能被放大很多倍。如果你的实现是自己手写的，优先把 RMSNorm 的 kernel 优化好，收益比优化注意力更明显。

4.3 FFN 的结构变体与激活函数

注意力层之后通常接一个两层的 FFN：先把d_model升维到4 * d_model，激活，再降维回d_model。这个 FFN 的参数量在整个模型中占比非常大，通常在三分之二以上，所以 FFN 的设计直接影响模型容量和计算量。

原始结构用的是 ReLU 激活，后来很多实现换成 GELU，因为它在负数区域不是完全死掉，梯度更平滑。GELU 有一个常用的近似公式0.5x(1 + tanh(sqrt(2/pi)(x + 0.044715x^3)))，精度足够。

另一个重要的变体是门控结构，比如 SwiGLU。它的做法是把中间层输出拆成两份，一份经过激活函数，另一份作为门控信号直接相乘。这种结构在很多大模型里已经变成标配，效果通常优于普通 GELU-FFN，但代价是参数量会多一些，实际实现时往往会把升维倍数从 4 降到 8/3，以维持总参数量不变。

做项目时如果你不想引入额外复杂度，普通 GELU FFN 依然完全能用；如果想提高上限，可以试试 SwiGLU。但我强烈建议先把基础版本跑通、确认数据没问题，再往 FFN 上堆花样，否则一旦效果变差，你根本分不清是 FFN 的问题还是数据的问题。

5. 训练 Transformer 的实操参数

5.1 学习率与 warmup：不是拍脑袋定的

Transformer 对学习率极其敏感，这一点我体会非常深：同样是调一个中规模模型，学习率从 3e-4 改到 6e-4，loss 曲线可以完全是两个世界。

原始论文用的是 Noam 学习率调度：先线性 warmup，再按步数的平方根倒数衰减。公式是lr = d_model^(-0.5) * min(step^(-0.5), step * warmup_steps^(-1.5))。这个设计的动机是 warmup 阶段让模型先在一个较小的学习率下稳定更新，避免一开始就把随机初始化的参数带飞；warmup 之后再慢慢衰减，让模型在损失平面上精细收敛。

warmup 步数怎么定？经验上小模型用 1000-5000 步，大模型可能要到 1 万步以上。一个更直觉的判断方法是：如果训练刚起步时 loss 就出现 NaN 或剧烈震荡，说明学习率太大或 warmup 太短；如果整个训练过程 loss 下降得过于平缓，可能是 warmup 过长占用了太多有效训练步数。

现在很多人直接用余弦退火替代 Noam 调度，实践证明它表现更好。余弦退火通常在 warmup 结束后，按余弦曲线从峰值学习率衰减到接近 0。我自己更倾向余弦退火，因为它对训练后期的探索和收敛平衡得更好，而且实现起来只要几行代码。

5.2 初始化、梯度裁剪与混合精度

Transformer 的初始化一直是个敏感话题。常规做法是把线性层权重初始化为均值为 0、标准差为sqrt(1/d_model)的正态分布，考虑到残差路径上多层累积，一些人会给最后输出层的权重乘一个额外缩放因子，比如 0.02 甚至更小，来防止初始阶段残差信号过强。

梯度裁剪几乎是必选项。混合精度训练下，FP16 的梯度范围很窄，一个偶发的异常梯度很容易导致 loss 变成 NaN。把梯度的全局范数裁剪到 1.0 或 0.5 是常见做法，尤其对语言模型，我基本不会省略这一步。

混合精度是训练速度和显存的核心优化。简单理解：前向计算用 FP16，反向梯度用 FP16，参数和主梯度保存在 FP32，每次更新时把 FP16 的梯度累加到 FP32 上。这里有一个特别容易踩的坑：FP16 的表示范围有限，如果一个参数的值非常小，在 FP16 下会直接变成 0，这就是 loss 莫名其妙变成 NaN 或模型权重更新失效的常见原因。所以混合精度实现里一定要启用“动态 loss scaler”——初始用一个很大的缩放因子把 loss 放大，训练中梯度过小就增大缩放因子，出现 inf 或 NaN 就减小缩放因子。这些事情开源训练框架已经处理好了，但如果你手写训练循环，务必自己实现或对照框架的实现逻辑做一遍。

5.3 数据侧：batch、padding 与 label smoothing

训练语料准备阶段，最常见的坑是 padding 方式。序列按 batch 最大长度补齐，这个最大长度会直接影响计算量。如果 batch 里有一句话特别长，整个 batch 都要按这个长度计算，效率会骤降。通用的做法是“按长度分桶”：把相似长度的句子分到同一个 batch，分别 pad，这样每个 batch 的 pad 占比都能压得很低。实测下来，训练吞吐能提升 20%-50%，非常可观。

另一个细节是 label smoothing。语言模型在训练时如果用 hard label（正确词是 1，其他是 0），模型会越来越“自信”，在 open-vocab 文本生成中容易过头，导致输出过于确定、多样性变差。label smoothing 的做法是给错误类别分一小部分概率，比如 0.1，这样可以起到一定的正则化作用。在文本生成里，稍微一点 label smoothing 往往就能改善输出质量，但设得太大（比如 0.3）会让模型变得迟钝，反而学不进去。

数据清洗的优先级也值得强调：重复样本、语料 lang 混乱、未做词表过滤，这些数据问题对训练效果的影响，远大于 model 结构上的微调。我见过不少项目花大量时间调结构，结果后来发现是语料里混入了大量重复的噪声文本，导致 loss 一直下不去。建议每次开始训练前，先抽 100 条数据人工看一眼，比任何超参都重要。

6. 常见问题与排查实录

6.1 loss 不降或异常高的排查清单

从零开始训练 Transformer，最崩溃的时刻就是训练了一两千步，loss 纹丝不动，甚至开始上涨。这种问题我几乎每周都会遇到，排查的顺序我整理成了下面这张表：

最常用也最好用的调试手段是“过拟合一个 batch”：拿一个 batch 的数据，不做 dropout、不设 label smoothing、把正则化全关掉，看模型能不能把这一批数据背下来。如果连一个 batch 都学不进去，那问题一定在模型实现或数据预处理环节；如果轻松过拟合，再逐步打开正则化组件，很快就能定位到哪个环节出了问题。

6.2 推理输出重复、退化的处理

生成任务里最扎心的现象是：模型重复输出同一句话或同一个词，语气僵化，像是卡住了。这个问题的根源常常不是推理代码，而是训练阶段的目标函数在鼓励“高概率的乏味输出”。

处理手段优先级从低到高，可以这样排：先把温度从默认的 1.0 调低一点，比如 0.8-0.9，但太低会让输出更呆板；再用 top-k 采样，只从概率最高的前 k 个 token 里采样，这也是一种去尾；再用 top-p（nucleus sampling），动态选择累积概率达到 p 的 token 集合；最后可以试试重复惩罚（repetition penalty）：对已经出现过的 token，在采样时给它们的分数乘一个小于 1 的惩罚系数。

实际操作里，我会把 top-k、top-p、温度三个参数组合起来调，而不是单一调某个。比如temperature=0.8, top_k=40, top_p=0.9是一套比较稳妥的起点；如果还是容易重复，就把 repetition penalty 设为 1.1 左右，再往上就很容易伤害语义。

还有一点容易忽略：beam search 的 beam 数越大，越容易在长度较长时陷入重复。如果做对话生成，我通常建议用采样而不是 beam search，因为 beam search 会系统性地偏高概率路径，导致输出更趋同。

6.3 显存优化三件套

如果你在训练或推理时被显存卡住，先按顺序做这三件事：

第一，开启梯度检查点：在 forward 时只保存中间结果中的一部分，反向时再重新计算，用时间换显存。对 Transformer 来说，这个开关能省下大概 30%-50% 的显存，代价是训练速度慢 20%-40%，但胜在能跑起来。

第二，检查 KV Cache 的处理。推理时 KV Cache 是显存的大户。能共享前缀就共享，能用 int8 量化的就量化，能分批处理的就分批。我遇到过一个真实场景：单条长序列推理时，KV Cache 占了显存的一半以上，导致 batch 只能开到 1，吞吐低得令人发指。后来对 K、V 做了分组量化，batch 直接翻了四倍。

第三，使用 Flash Attention。注意力计算在算子层面做了大量重组，把内存访问次数显著降低。它约束了 mask 的实现方式（有些自定义 mask 不支持），但性能和显存收益太明显。只要能用，我几乎不会考虑手写 attention，除非要做特殊 mask 或者严格对齐某篇论文的实现细节。

这三个手段做完，如果还是不够，再考虑模型并行或流水线并行。但 90% 的情况下，上面三件套就能把你从显存焦虑里救出来。

最后分享一个我在实际操作中养成的习惯：每个实验跑完，除了记录 loss 曲线和最终指标，我会把当时调过的所有超参数、数据情况、踩过的坑写进一个固定的笔记模板里。这个习惯帮我避免了很多“重蹈覆辙”——尤其是我踩过的 mask 合并、FP16 下溢、KV Cache 拼接这些坑，如果不记录下来，过两周再看之前的代码，完全想不起来当时为什么那样写。

另外，新手复现 Transformer 时，我强烈建议不要一上来就上大规模数据和超大模型。先用一个几十万条的小语料、一两层的小模型把训练流程完整跑通，再逐步放大。这一步看似浪费时间，实际上能帮你把注意力全部集中到模型本身，而不是被数据清洗、分布式训练这些外围问题淹没。等你把这一整套流程跑顺了，再回去看论文里的各种改进，会发现自己终于能看懂它们到底在改什么了。