Files
llm-atlas/research/K3_ATTNRES_GRADIENT_DEFINITION_AUDIT.md

9.1 KiB
Raw Permalink Blame History

Kimi K3 第五轮前置审计:Attention Residuals 的“梯度更均匀”到底指什么

审计日期:2026-07-30(Asia/Shanghai)
官方仓库:MoonshotAI/Attention-Residuals@85e22310fe5ee860b4a023de312d791de8a5a5e6
官方 PDF SHA-256:e5831b0db1347606453b5176b0142115a18887b6a9c2e1d05a266d4805a26b2f
结论性质:一手工件审计,不是新的实验结果

0. 先说结论

K3 Round 04 的反结果——Block AttnRes 的核心参数梯度 RMS 跨层 CV 更高——不能直接 反驳 Attention Residuals 论文 Figure 5(c) 所说的“梯度分布更均匀”,因为两边很可能测的 不是同一个对象:

Round 04:
每个 Transformer block 内所有核心参数梯度拼接后的 RMS
∇θL,θ = attention + MLP + 两个输入 norm 的参数

论文 Figure 5(c):
图题只写 “Each transformer block's gradient magnitude”
结合 Figure 5(b) 的 block output magnitude 和正文,最自然的操作化是
每个 block 输出 activation 的梯度 ∂L/∂h_l

但“最自然”不等于“官方已经明确定义”。论文和当前官方仓库都没有给出足以唯一重建 Figure 5(c) 的测量合同,也没有发布训练代码。因此,下一轮不会把自己的 activation-gradient 定义冒充成论文原始实现,而会把它命名为:

与 Figure 5 叙述对齐的一种公开、冻结、可复现的 operationalization

这一区分很重要:参数梯度回答“这一层的权重此刻收到多大更新信号”,activation 梯度回答 “损失对这一深度的表征有多敏感”。二者相关,但不会因为链式法则而自动同方向。


1. 官方工件实际提供了什么

官方仓库在固定 revision 下只包含:

  • README.md;
  • Attention_Residuals.pdf;
  • 论文图片资产;
  • citation 与外部入口。

仓库不包含:

  • 模型或 residual mixer 的可执行实现;
  • Figure 5 的统计脚本;
  • 训练配置、日志或 checkpoint;
  • gradient hook、norm 和 reduction 定义;
  • 用于复画 Figure 5 的原始数组。

因此,本审计能固定论文的文字、公式、图与模型尺度,不能从官方代码恢复一个不存在的 隐藏测量合同。

2. Figure 5 能确认的事实

官方 training_dynamics.png 和 PDF Figure 5 有三个并列面板:

面板 图题 横轴
(a) Validation Loss training step
(b) Output Magnitude Layer / Transformer Block Index
(c) Gradient Magnitude ×10⁻⁵ Layer / Transformer Block Index

Figure 5 caption 对 (b) 与 (c) 的完整对象描述分别是:

Each transformer block's output magnitude at the end of training.
Each transformer block's gradient magnitude.

正文明确表达了两个方向:

  1. Baseline 的 hidden-state magnitude 随深度单调增长;Block AttnRes 把增长约束在局部 residual block 内,形成周期性的深度图案。
  2. Baseline 的早期层梯度“不成比例地大”;Block AttnRes 的可学习 softmax 权重产生了 “明显更均匀”的梯度分布。

图中最终大模型约有 27 个 Transformer blocks。论文的最终模型描述也是 27 个 Transformer blocks / 54 个 residual layers;Block AttnRes 每 6 个 residual layers 聚合一次,共 9 个 聚合块,外加 embedding 形成 10 个跨块来源。

3. Figure 5 不能确认的事项

下面每一项都会改变曲线,却没有在论文或官方仓库中被唯一指定:

未定义项 至少两种合理解释
梯度对象 block 输出 activation 梯度;block 参数梯度;分支输出梯度
block 输出位置 attention+MLP 后;只在 MLP 后;进入下一层 norm 前;聚合块边界后
norm L2 norm;RMS;mean absolute value;每 token norm 后再平均
reduction batch/token/channel 联合;先按 token 再按 batch;只取末 token
loss token mean;sample mean;未归一化 sum;带或不带 mask
采样 一个 batch;多 batch 平均;训练流中的 moving average
时间点 “训练结束”单点;末段平均;某个 checkpoint
数值阶段 AMP 缩放前/后;gradient clipping 前/后;BF16 或 FP32
运行模式 train 或 eval;dropout 是否开启
归一化 绝对值;再除全层均值;再除 Baseline

Figure 5(c) 的纵轴是绝对 magnitude 标度,并不等价于 CV。只报告 CV 还会丢失两个信息:

  • 全部层梯度是否一起缩小或放大;
  • 不均匀来自“早层系统性偏大”,还是某一个中间/末端尖峰。

所以 Round 05 必须同时公开绝对曲线、按层均值归一化曲线、CV 与前后深度分位比。

4. 为什么 activation gradient 是合理推断,但仍只是推断

把 Figure 5(c) 操作化为 ∂L/∂h_l 有三条证据:

  1. 它与 Figure 5(b) 的 “transformer block output magnitude” 在横轴和叙述上成对;
  2. “早期层的梯度”在表示传播语境中通常可由对 block output 保留梯度直接比较;
  3. 参数张量的大小和类型在 attention 与 MLP 间差异很大,若把参数拼接,论文通常需要说明 聚合口径,否则 “each transformer block” 不是天然的单一标量。

但也有无法排除的替代解释:

  • 论文作者可能测 block 参数梯度;
  • 可能测 residual branch output 而非完整 block output;
  • 可能先对每个 token 做 L2 norm,再跨 token 平均;
  • 可能在内部训练系统中有未公开的统一 telemetry 定义。

因此,网站和审计只说“与论文叙述对齐的公开定义”,不说“论文就是这样算的”,也不把 数值和 Figure 5 纵轴直接对齐。

5. Round 04 与 Round 05 的对象对照

维度 Round 04 已测对象 Round 05 主对象
数学对象 ∇θ_l L ∂L/∂h_l
l 的单位 Transformer block Transformer block
张量内容 block 的核心参数 block 的 post-MLP output activation
聚合 参数元素联合 RMS batch×time×channel 联合 RMS
是否含 AttnRes 参数 否 不适用;梯度穿过 mixer
时间 final diagnostic batch 全部预注册 diagnostic steps
目的 权重更新信号是否均匀 表征深度的反向信号是否均匀

Round 04 的参数梯度结果不会被改名、删去或用新指标覆盖。它仍是一个有效反结果,只是不能 代表论文未定义清楚的 Figure 5(c)。

6. 两种结构怎样取得真正对齐的 16 / 32 个位置

Grok Headless 被用作一次对抗式方法审阅,不作为事实来源。它正确指出了 non-leaf tensor、 alias、AMP、clip 时点和只看 CV 的风险;但它也提出了一个不适用于本实现的担忧: “Block AttnRes 只有约 8 个 block 输出,无法与 Baseline 的 16 / 32 层对齐”。

这里要区分两种 block:

Transformer block:
attention + MLP;depth=16 时始终有 16 个,depth=32 时始终有 32 个

AttnRes aggregation group:
把若干 residual sublayers 的 partial sum 保存为一个跨组 source;
两种深度都约为 8 组

Round 05 在每个 Transformer block 的 MLP 分支完成后取 h_l。所以 Baseline 与 Block 都有完全相同的 l=1..depth:

架构 Round 05 的 h_l shape
Baseline 第 l 个 attention residual 与 MLP residual 都完成后的 hidden state [B,T,C]
Block 第 l 个 MLP branch 加入后、可能保存并清空 aggregation partial 之前的 partial output [B,T,C]

Block 的 h_l 是局部 residual partial,而不是跨组 source 列表;这正对应论文 Figure 5(b) 所描述的“增长被限制在每个 block 内”的周期性图案。组边界前取值也避免把 reset 后的零张量 错误当成 Transformer block output。

7. 实验实现必须通过的梯度测量闸门

正式训练前,四个结构格(2 个深度 × 2 个 residual graph)都必须证明:

  1. 所有 h_l 都是不同的、非别名的捕获对象,数量严格等于 Transformer depth;
  2. retain_grad() 或 hook 后所有梯度非 None、finite,shape 与 activation 完全一致;
  3. diagnostic loss 使用固定输入、固定 token-mean CE、eval()、无 optimizer step;
  4. backward 发生在 parameter gradient clip 之前,且不经过 GradScaler;
  5. 将同一 diagnostic loss 精确乘 2 后,每层 activation-gradient RMS 也乘 2;
  6. 乘 2 前后的 CV、归一化曲线与深度分位比在数值容差内不变;
  7. 同配置全新进程重复运行,冻结诊断字段 exact。

若任一项失败,正式 8,000-step grid 不得开始。

8. 本审计带来的研究决策

下一轮不再把一个宽度 192、深度 16、训练 2,000 step 的参数梯度 CV 与论文最终模型图强行 放在同一条结论线上,而是:

  • 增加 depth 32;
  • 把训练预算扩为 8,000 step;
  • Baseline / Block 使用相同的 Transformer block index;
  • 在 6 个固定时点测 post-MLP activation gradient;
  • 绝对标度、归一化形状、CV、前后四分位失衡一起公开;
  • 参数梯度作为次要指标保留;
  • 预注册“支持 / 混合 / 不支持”规则,并公开全部反结果。

精确协议见 research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md。