从0训练LLM · 现代大模型体系 阅读 + 实验 ≈ 40 分钟 页内实验 ×5
C8 稳定性:放大后为什么会炸 第 4 课

怎么按住 logit?

按住 logit:z-loss、QK-Norm、soft-cap、QK-Clip
÷√d 的推导假设 q、k 独立,训练会让它们对齐、变长,点积按 d 而不是 √d 增长;分数 8/7/9 只放大 10 倍,softmax 就成了 99.995% 的独占,梯度几乎为零。现代做法按对象分别按住:输出端用 z-loss 把 log Z 拉回 0;注意力端用 QK-Norm 钉住 q、k 的长度(Qwen3.8 每头上界 16·g²),或用 soft-cap 压、用 QK-Clip 在超标后缩权重。
❓

本课一问:上一课找到的病根是「logit 被一步步养大」。能不能不降学习率,直接把 logit 按住?按哪里、怎么按?

这一课先不碰:logit 大了以后在 BF16 里的舍入(C9 细算,本课只点到);FFN 激活的截断(SiTU-GLU、swiglu_limit,C15)。

从已有出发

从已有出发 · 累计数据流 点已学环节看形状与参数
解码层

模型里有两处 softmax,也就有两处 logit 会失控:注意力里的 q·k/√d(C3),和输出层的 logits(每个候选词一个分数)。C3 第 3 课用「小明 / 苹果 / 小王」这一行分数 8 / 7 / 9 算过 softmax,也在 Dig 里留了一句:除以 √d 不是一劳永逸。C5 第 7 课留下的数字是:Kimi K2 用原版 Muon,中等规模 run 的最大注意力 logit 很快超过 1,000。

先自己设计:只靠 ÷√d 和梯度裁剪

朴素的想法:注意力已经除以 √d 了,梯度也裁剪了,还要什么?先看 logit 被放大之后 softmax 会怎样——把 C3 那一行分数乘 k 倍:

  1. ① 朴素方案
  2. ② 它在哪失败
  3. ③ 改进方案
朴素方案只靠 ÷√d + 梯度裁剪
把一行分数放大 k 倍:softmax 会变成什么样

为什么 ÷√d 按不住

÷√d 的推导(C3 第 3 课)假设 q、k 的各维独立、方差为 1——那是初始化那一刻。训练中两件事会同时发生:WQ、WK 变大,q、k 变长;该互相关注的 q、k 越来越对齐。一旦对齐,d 项同号相加,点积按 d 增长,而不是 √d——μP 的论文专门指出过这一点 。

为什么 ÷√d 按不住:q、k 一旦对齐、变长

🎯

试三件事:① ρ 拖到 1、a 保持 1:每头 256 维时 logit 约为 16(理论值 a²·ρ·√d = 16,实测随这一对随机向量略有出入),不是 1;② 再把 a 拖到 4:红线按 a² 涨到两百多,绿线一动不动;③ 把 d 切到 32(我们项目的每头维数):同样完全对齐时理论值只有 5.66,QK-Norm 后恰好是 5.66——维数越大,对齐之后的 logit 越大,QK-Norm 的上界也越高。

改进一:输出端的 z-loss——把 log Z 拉回 0

输出层的 logits 有一个交叉熵完全看不见的方向:所有 logit 一起加一个常数。softmax 不变,loss 不变,梯度在这个方向上恰好为 0(交叉熵对 V 个 logit 的梯度之和是 Σ(pi − yi) = 1 − 1 = 0),于是 log Z = log Σ ezi 可以一直漂,没有任何力量把它拉回来。

输出层的 log Z:交叉熵不管它,z-loss 管

z-loss · Mesh-TensorFlow 代码库 → PaLM §5(系数 10⁻⁴);OLMo 2 沿用
= + ·
👆 公式里带下划线的每一项都可以点——点哪项，讲哪项。
归属:ST-MoE(2022)称它改编自 Mesh-TensorFlow 代码库里「用于最终 softmax logit 的 z-loss」;PaLM 是公开写明系数的大模型;MoE 路由器上的 router z-loss 是 ST-MoE 首提的另一种用法。

z-loss 管的是 Wortsman 说的第二类不稳定——输出 logit 与 log 概率发散 。上面组件最后一列提醒了另一件事:logit 漂到几十、上百后,BF16 相邻两个可表示数的间隔变成 0.25、0.5 乃至更大——两个本该分得清的 logit 被存成同一个数,C9 精度会回到这里。

改进二:注意力端——四种按法

QK-Norm:从源头钉住 q、k 的长度

把 q、k 在打分之前各自做一次 RMSNorm(带可学习的逐维增益),再点积:

QK-Norm 之后的注意力 logit(Henry et al. 2020 → ViT-22B 2023 → Qwen3 / Qwen3.8)
= ·
👆 公式里带下划线的每一项都可以点——点哪项，讲哪项。
QK-Norm 不改变「看谁」的排序信息里的方向部分(余弦相似度),它拿掉的是 q、k 的长度这个自由度——正是训练中会失控的那个。
锚点 Qwen3.8 的一个全注意力头:q、k 先各自 RMSNorm,再打分(RoPE 省略)
S = (RMSNorm(q)·gq)(RMSNorm(k)·gk)T / √dh
输出中间结果 渐变 = 分数(连续值) 留白 = 结构上恒为 0
边长按对数比例缩放,块边上的数字是真实值;同一维度符号的边在全图等长。
维度含义 真实规模 小规模
T 序列长度(例句 5 个 token) 5 5
dh 每头维数(head_dim) 256 4
维数取自 anchor-model.json;q_norm、k_norm 的形状 [256] 取自 R2 对官方 safetensors 的分项。Qwen3.8 只有 16 个全注意力层有这一步(另外 48 层是线性注意力,C11)。

归属要写对:QK-Norm 首提于 Henry 等人 2020 年的低资源翻译论文(ℓ₂ 归一化加可学习尺度) ;把它用到大规模的是 ViT-22B——在 8B 参数附近,注意力 logit 失控,权重塌成几乎 one-hot、熵接近 0,对 q、k 做 LayerNorm 之后才稳 。今天 Qwen3 / Qwen3.8、OLMo 2、Gemma 3、GLM-4.5 都用它。

soft-cap:在这一步直接压

Gemma 2 把注意力 logit 过一个 50·tanh(z/50)(最终层用 30),logit 永远不超过 50 。代价在上面「改进方案」的表里:被压扁的那几个 logit,tanh 的导数接近 0,梯度传不回去。下一代 Gemma 3 就把 soft-cap 换成了 QK-Norm 。

QK-Clip:超标之后回头缩权重

Kimi K2 的 QK-Clip(MuonClip 的一部分,技术报告 §2.1)
= min(1, /
👆 公式里带下划线的每一项都可以点——点哪项，讲哪项。
K2 的实测:最大 logit 很快涨到上限 100,大约 30% 的训练步之后才自己回落到稳定区间;15.5T token 全程 loss 平滑,没有 spike。
🔍 K2 用的是 MLA:QK-Clip 对不同分量怎么分?

MLA(C3 第 10 课)的 q、k 各有两部分:每个头自己的「内容」分量 qC、kC,和带 RoPE 的分量 qR、kR——其中 kR 是所有头共享的。K2 的做法:qC、kC 各乘 √γh;qR 乘 γh;共享的 kR 不动,免得一个头超标把所有头都缩了。出处:Kimi K2 技术报告 §2.1(arXiv 2507.20534)。

DeepSeek-V4 又走了一步:它的注意力结构允许在打分之前直接对 query 的每个头和唯一的压缩 KV 头做 RMSNorm,报告明说「因此在 Muon 优化器里不使用 QK-Clip」 ——K2 的 MuonClip 被一个归一化换掉了。

我们自己的模型上:QK-Norm 开 / 关

真训练日志 · 只改 QK-Norm 与学习率
项目实测 · QK-Norm 开 / 关:最大注意力 logit

读图三点:不开 QK-Norm 时,学习率 3e-2 的最大注意力 logit 终值是 6784,开了之后是 16.4;开 QK-Norm 那一组开局的最大 logit 就是 4.34,在上界 5.66 以内,后来超过 5.66 是因为可学习增益长大了;同一学习率下,开 QK-Norm 的最终 val 都不比关着差(3e-2 时 5.045 对 5.832),基线学习率下的差距在噪声以内。

2026 年各家怎么选

模型注意力 logit输出 logit出处
Qwen3 / Qwen3.8QK-Norm(q_norm、k_norm)—Qwen3 技术报告 §2;Qwen3.8 权重分项(R2)
OLMo 2QK-Norm(RMSNorm)z-loss 10⁻⁴OLMo 2 §3.3.2–3.3.3
PaLM—z-loss 10⁻⁴PaLM §5
Gemma 2 → Gemma 3soft-cap 50 → QK-NormGemma 2 最终层 soft-cap 30Gemma 2 §2;Gemma 3 §2
Kimi K2 / K3QK-Clip(τ = 100)—K2 §2.1;K3 沿用 K2 的权重裁剪(R3 §3.6)
DeepSeek-V4QKV RMSNorm,不用 QK-Clip—V4 §2.3.3、§2.4(本地 PDF)
Qwen3.8-Flash-Next自报没有用 qk-clip—Flash-Next 技术报告(R2 §1.5.2,自报)

共识是「一定要按住」,分歧在按法:从源头归一化(QK-Norm、QKV RMSNorm)成了多数,事后约束(QK-Clip)在 Muon 系里仍有人用。同一思路还在往 FFN 的激活上延伸(SiTU-GLU 软封顶、swiglu_limit 截断),C15 展开。

⚠

常见误解一:「÷√d 已经把 logit 的尺度定死了。」它只在初始化、q 与 k 独立时把方差拉回 1;训练中 q、k 变长、对齐,logit 照样长到上千(我们项目里就长到了几千)。C3 第 3 课那个 Dig 的伏笔在这里兑现。

常见误解二:「开了 QK-Norm,logit 就永远有界,不用再管。」上界是 max|gq|·max|gk|·√d,增益是可学习的;我们项目里开 QK-Norm 的最大 logit 就从 4 左右长到了十几。它把失控的速度压下来了,仪表仍然要看。

常见误解三:「z-loss 改变了模型的预测。」softmax 对整体平移不变,z-loss 只把 log Z 拉回 0,系数又只有 10⁻⁴;它补的是交叉熵在「整体平移」方向上缺的那股回复力。

常见误解四:「QK-Clip 是一种梯度裁剪。」它不碰梯度,也不改这一步的前向和反向;它在更新之后缩那个头的 Wq、Wk,约束的是权重的增长。

证据三问

问本课的回答
① 它解决什么问题?不降学习率,直接限制两处 softmax 的输入尺度:注意力 logit 与输出层的 log Z
② 数学上为什么可能有效?softmax 饱和时雅可比趋于 0(梯度传不回)、QK-Norm 的上界(柯西-施瓦茨)、交叉熵在整体平移方向上梯度为 0,都可证 ;「按住尺度就能减少 spike」是理论动机
③ 实验上为什么最后选了它?ViT-22B 在 8B 规模靠 QK-Norm 训稳;Wortsman 的小模型上 qk-layernorm 与 z-loss 把学习率敏感度大幅降低;Gemma 3 用 QK-Norm 替换 soft-cap;K2 用 QK-Clip 训完 15.5T;我们项目里 QK-Norm 把 logit 从几千压到十几

一句话

📌
logit 有两处会失控:注意力的 q·k 与输出层的 log Z;÷√d 只管得住初始化那一刻,QK-Norm 钉住 q、k 的长度,z-loss 把 log Z 拉回 0,QK-Clip 超标时回头缩权重——按住的是 logit 本身,不是梯度。
🚦

先停在这里:记住「对齐之后点积按 d 长」「QK-Norm 上界 = 增益 × √d」「z-loss 补的是平移方向」。
再深一层:读 K2 §2.1 与 Wortsman §3.1,展开上面 MLA 的 Dig。

下一个缺口

稳定手段都齐了:小模型上调好的学习率,放大到宽 20 倍的模型还能用吗? → 第 5 课:小模型上调好的学习率,大模型能用吗?

回扣:C6 第 4 课给了一条路——拟合超参缩放律;那一课说另一条路「换参数化,让学习率不随宽度变」(μP)在 C8 第 5 课兑现。μP 的论文里还有一句和本课直接相关:标准参数化下,一步更新之后注意力 logit 会随宽度变大。

你可能会问