写这一篇笔记之前我已经把 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_k64为例点积的标准差是 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_model512时配 8 个头是常见的d_model768时配 12 个头也常见。核心约束是d_model必须能被 h 整除且d_k不要太小不然每个头能表达的信息太有限。我试过把d_model256强行拆 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 的长度是 1K 的长度是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-4loss 曲线可以完全是两个世界。原始论文用的是 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 纹丝不动甚至开始上涨。这种问题我几乎每周都会遇到排查的顺序我整理成了下面这张表现象可能原因排查动作loss 从第一步就是不降学习率过大或过小、数据有大量噪声先用原论文的 Noam 调度跑 200 步看曲线loss 震荡剧烈或出现 NaN学习率太大、梯度裁剪失效、FP16 数值下溢检查 loss scaler 状态调低学习率loss 降到一定程度后停滞模型容量不够、数据重复、mask 配错检查训练/验证集重合度打印 attention 矩阵验证 loss 升、训练 loss 降过拟合加 dropout、加数据增强、减小模型出现过拟合但模型很小数据太简单、任务本身简单先确认是否为实现 bug再谈优化最常用也最好用的调试手段是“过拟合一个 batch”拿一个 batch 的数据不做 dropout、不设 label smoothing、把正则化全关掉看模型能不能把这一批数据背下来。如果连一个 batch 都学不进去那问题一定在模型实现或数据预处理环节如果轻松过拟合再逐步打开正则化组件很快就能定位到哪个环节出了问题。6.2 推理输出重复、退化的处理生成任务里最扎心的现象是模型重复输出同一句话或同一个词语气僵化像是卡住了。这个问题的根源常常不是推理代码而是训练阶段的目标函数在鼓励“高概率的乏味输出”。处理手段优先级从低到高可以这样排先把温度从默认的 1.0 调低一点比如 0.8-0.9但太低会让输出更呆板再用 top-k 采样只从概率最高的前 k 个 token 里采样这也是一种去尾再用 top-pnucleus sampling动态选择累积概率达到 p 的 token 集合最后可以试试重复惩罚repetition penalty对已经出现过的 token在采样时给它们的分数乘一个小于 1 的惩罚系数。实际操作里我会把 top-k、top-p、温度三个参数组合起来调而不是单一调某个。比如temperature0.8, top_k40, top_p0.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 时我强烈建议不要一上来就上大规模数据和超大模型。先用一个几十万条的小语料、一两层的小模型把训练流程完整跑通再逐步放大。这一步看似浪费时间实际上能帮你把注意力全部集中到模型本身而不是被数据清洗、分布式训练这些外围问题淹没。等你把这一整套流程跑顺了再回去看论文里的各种改进会发现自己终于能看懂它们到底在改什么了。