Files
llm-atlas/research/K3_ATTNRES_GRADIENT_SCALE_AUDIT.md
2026-07-30 09:49:53 +08:00

17 KiB
Raw Permalink Blame History

Kimi K3 第五轮:Attention Residuals 梯度定义与深度扩展实验审计

协议:llm-atlas-k3-attnres-gradient-scale-v1
前置定义审计:research/K3_ATTNRES_GRADIENT_DEFINITION_AUDIT.md
预注册协议:research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md
数据清单:experiments/k3/attnres_gradient/manifest.json
执行日期:2026-07-30
设备:NVIDIA GeForce RTX 5090;PyTorch 2.11.0+cu128

0. 先说结论

这一轮本来想澄清一个看似矛盾的问题:

论文 Figure 5:
Block AttnRes 的梯度沿深度“明显更均匀”

本站 Round 04:
Block AttnRes 的核心参数梯度 RMS 跨层 CV 反而更高

一手工件审计先确认:论文没有公开 Figure 5(c) 的确切 gradient tensor、norm、reduction、 diagnostic batch、AMP / clip 时点或统计代码。因此,Round 05 没有假装恢复作者的隐藏实现, 而是冻结一个与 Figure 5 的 output/gradient 并列叙述对齐、可复现的定义:

h_l:
第 l 个 Transformer block 完成 attention + MLP 后的 FP32 residual output

m_l:
sqrt(mean((∂L / ∂h_l)² over batch × time × channel))

L:
固定 16 × 256 targets 的 token-mean cross entropy

结果不是一句“是”或“不是”,而是一个更有信息量的分解:

  1. Block 确实大幅缓解了早层整体偏大。
    Baseline 首四分位梯度平均是末四分位的 3.01–3.90×;Block 把它改到 0.55–1.40×。预注册的首尾失衡指标在 6 / 6 个 depth×seed 配对中都改善, depth-16 平均改善 61.0%,depth-32 平均改善 72.0%。

  2. 但 Block 没有让整条深度谱更平。
    它在中后段形成了局部尖峰,所以 population CV 在 6 / 6 个配对中都恶化: depth-16 平均相对恶化 10.3%,depth-32 平均相对恶化 60.0%。

  3. Block 的绝对 activation-gradient 平均尺度更小。
    它只有 Baseline 的 57.4%(depth-16)和 54.4%(depth-32)。这意味着“首尾更接近” 不能自动解释为所有层都获得更强更新信号。

  4. 参数梯度仍与 Round 04 同方向。
    核心参数梯度 CV 从 0.416→0.683(depth-16),从 0.397→0.772 (depth-32);Block 更不均匀。

  5. 验证 BPC 在 6 / 6 配对中都更低,但不是同算力优势。
    平均改善 0.00894 BPC(depth-16)与 0.00987 BPC(depth-32);Block 实际 step time 是 Baseline 的约 2.55–2.60×,peak allocated memory 约 2.15–2.16×。

按看结果前冻结的联合判据,CV 与首尾失衡必须同时改善才算 support;必须同时恶化才算 concern。这里二者方向相反,所以:

depth 预注册判定
16 mixed / inconclusive at this depth
32 mixed / inconclusive at this depth
总判定 depth-dependent or inconclusive

最准确的中文总结是:

在这套公开 operationalization 中,Block AttnRes 把“早层系统性偏大”变成了“首尾更接近、 但中后段有局部尖峰”的另一种梯度分布。它修正了一类失衡,却没有降低全层离散度。

这不复现论文 Figure 5 的数值,也不反驳一个没有公开测量合同的隐藏实现。


1. 为什么要另开一轮,而不是改写 Round 04

Round 04 的梯度对象是:

每个 Transformer block 的 attention、MLP 与两个 input norm
所有核心参数梯度拼接后的 RMS

这回答“该层权重收到多大更新信号”。Round 05 的对象是 ∂L/∂h_l,回答“损失对该深度 表征有多敏感”。链式法则把两者联系起来,但不会保证跨层形状同方向。

因此:

  • Round 04 参数梯度反结果继续有效;
  • Round 05 不把它改名为 activation gradient;
  • 两种对象在网站并排展示;
  • 任何方向冲突都保留,而不是选择更像论文的一种。

官方工件边界见前置定义审计。固定 revision 为:

MoonshotAI/Attention-Residuals
85e22310fe5ee860b4a023de312d791de8a5a5e6

Attention_Residuals.pdf SHA-256
e5831b0db1347606453b5176b0142115a18887b6a9c2e1d05a266d4805a26b2f

官方仓库没有可执行训练代码或 Figure 5 原始数组。

2. 实验规模与配对合同

正式网格:

2 depths
× 2 residual graphs
× 3 seeds
× 8,000 steps
× 32 windows
× 256 target bytes
= 12 independent runs
= 786,432,000 formal target bytes
项 depth-16 depth-32
Transformer blocks 16 32
residual sublayers 32 64
AttnRes aggregation groups 8 8
sublayers / group 4 8
Transformer blocks / group 2 4
width / heads / FFN 192 / 6 / 768 192 / 6 / 768

每个 depth / seed 的 Baseline 与 Block:

  • 使用逐 tensor exact 的公共主干初始化;
  • 使用逐 step / row exact 的 byte windows;
  • 使用相同 optimizer、LR schedule、batch、context 与 target-byte budget;
  • 分支线性计算走 BF16 autocast;
  • residual accumulator 与被测 h_l 都为 FP32;
  • 每个结构在全新进程中从零训练。

二者不匹配:

  • mixer 参数;
  • mixer FLOPs;
  • step wall time;
  • activation memory。

所以 BPC 只能叫“同 token / step 预算对比”,不能叫“同算力优势”。

3. 数据与日程

对象 固定值
dataset Salesforce/wikitext
revision b08601e04326c79dfdd32d625aee71d232d685c3
variant wikitext-2-raw-v1
train bytes 10,951,563
train SHA-256 0ca7d3e7…e9b4
validation bytes 1,148,008
validation SHA-256 a42356f6…e719
formal schedule cells 768,000
schedule SHA-256 5041e09b…f4e
validation tensor SHA-256 f459316f…338
diagnostic tensor SHA-256 21117e31…716

固定诊断时点:

0, 100, 500, 2,000, 4,000, 8,000

六个时点的全部 activation gradient、output RMS、parameter gradient、mixer 权重与验证 BPC 都进入 raw JSON;没有只挑“最好看”的 checkpoint。

4. 正式训练前的故障与修订

4.1 第一次 smoke 的标量序列化错误

首个 depth-16 Baseline 20-step smoke 已完成数值计算,但在写 JSON 前失败:

RuntimeError:
self.dim() cannot be 0 to view Float as Byte

原因是 optimizer 的 step 是 0 维 tensor,hash helper 直接把它 view(torch.uint8)。修复为:

tensor.reshape(-1).view(torch.uint8)

当时:

  • 没有 formal 运行;
  • 没有输出 JSON;
  • 没有可供选择的正式结果。

4.2 smoke 发现被测 residual dtype 不一致

第一版 smoke 通过了 finite / loss×2 / replay 闸门,但检查 capture metadata 时发现:

Baseline post-MLP h_l:FP32
Block aggregation partial:BF16

原因:

  • Baseline 把 BF16 branch 加到 FP32 embedding/residual stream;
  • Block 每组第一个 partial 直接引用 BF16 branch output。

这会把数值精度差异混进结构比较。正式训练前,协议与实现补充为:

BF16 branch output → 显式转 FP32 → residual partial 累加

随后 4 个 depth×architecture 格的 smoke 全部从头重跑两次。正式输出是在这次修订之后才开始。

4.3 一次未启动模型的 zsh 调度错误

首个正式 Baseline 完成后,批处理脚本用 Bash 式标量切分处理 zsh 字符串,第一行就退出:

argument --architecture: invalid choice: ''

runner 没有启动,也没有创建新结果文件。调度改为显式 : 分隔数组后继续。这是 orchestration 故障,不是模型运行失败,但仍在审计时间线中保留。

5. 运行前与复现闸门

5.1 activation-gradient 测量闸门

4 / 4 个 depth×architecture 格都通过:

  • 捕获数量严格等于 16 / 32;
  • shape 严格为 [16,256,192];
  • dtype 全部为 FP32;
  • gradient 全部 present、finite;
  • capture storage 全部不别名;
  • diagnostic loss×2 后每层 gradient RMS 精确×2;
  • CV、归一化谱、首尾比不变。

5.2 两次独立 smoke

每个格都在两个全新进程中训练 20 steps。排除 timing 后的冻结字段:

格 compare SHA-256
depth-16 Baseline 525cdcd9…bd2a
depth-16 Block e4330a0d…90cc
depth-32 Baseline 8ee0dc37…8e56
depth-32 Block 289073b1…34c

4 / 4 exact。

5.3 完整 formal replay

预注册格:

depth-32 / Block / seed-2026073001

从初始化重新训练完整 8,000 steps,不加载 formal checkpoint。排除 run_kind、timing、 memory 与进程元数据后的全部冻结字段:

formal compare SHA-256
46300a452840a9dc6a5180efe7949cf3d81cc4471ed4a942da407e9343064817

replay compare SHA-256
46300a452840a9dc6a5180efe7949cf3d81cc4471ed4a942da407e9343064817

最终状态:

对象 formal replay
model state 3f0b97ec…2f59 3f0b97ec…2f59
optimizer state ed03e6fb…4637 ed03e6fb…4637

字段级 exact。

6. 主结果:CV 与首尾失衡为什么方向相反

6.1 depth-16

seed Base CV Block CV 相对 CV 改善 Base imbalance Block imbalance 相对 imbalance 改善
2026073001 0.41719 0.46768 −12.1% 1.29458 0.38821 +70.0%
2026073002 0.43744 0.45175 −3.3% 1.35991 0.58958 +56.6%
2026073003 0.41927 0.48391 −15.4% 1.29826 0.56693 +56.3%
均值 0.42463 0.46778 −10.3% 1.31759 0.51491 +61.0%

这里“相对 CV 改善”为负,表示恶化。

原始首/末四分位比:

Baseline:3.65×, 3.90×, 3.66×
Block:   0.68×, 0.55×, 0.57×

Block 不只是把早层优势降到 1;它在三个 seed 中都略微“过冲”,变成末四分位平均更大。 但 abs(log(first/last)) 仍比 Baseline 更接近 0,所以失衡改善。

三 seed 平均 normalized activation-gradient 的最高点:

Baseline:layer 3 = 1.50× mean;layer 4 = 1.42×
Block:   layer 11 = 2.09× mean;layer 13 = 1.95×

Baseline 是宽而平滑的早层隆起;Block 是更局部的中后段尖峰。CV 对尖峰敏感,所以升高。

6.2 depth-32

seed Base CV Block CV 相对 CV 改善 Base imbalance Block imbalance 相对 imbalance 改善
2026073001 0.36109 0.64203 −77.8% 1.10030 0.20781 +81.1%
2026073002 0.37162 0.72997 −96.4% 1.14087 0.43810 +61.6%
2026073003 0.40320 0.42685 −5.9% 1.25169 0.33417 +73.3%
均值 0.37864 0.59962 −60.0% 1.16429 0.32669 +72.0%

原始首/末四分位比:

Baseline:3.01×, 3.13×, 3.50×
Block:   0.81×, 0.65×, 1.40×

三 seed 平均 normalized spectrum 的 Block 峰值:

layer 21 = 3.04× mean
layer 22 = 2.41×
layer 23 = 1.91×
layer 25 = 1.77×

depth-32 每个 AttnRes aggregation group 含 4 个 Transformer blocks。21–24 是第 6 组, 25–28 是第 7 组。尖峰集中在这两个中后段组附近,是数据中直接可见的结构;但仅凭本实验 不能断言 pseudo-query、某个 source 或组边界是唯一因果。

6.3 seed-3 的中期反例

depth-32 seed-3 在 step 2,000:

Baseline CV 0.39965
Block CV    0.34914

此时 Block 更平;到 step 8,000 才变成:

Baseline CV 0.40320
Block CV    0.42685

前两个 seed 在 step 2,000 已明显恶化,seed-3 没有。网站必须保留 seed switch,不能用最终 均值倒写成“三个 seed 从头到尾都一样”。

7. 绝对梯度尺度:更平不等于更强

最终 activation-gradient mean:

depth Baseline Block Block / Baseline
16 约 2.02×10⁻⁴ 约 1.16×10⁻⁴ 0.574×
32 约 1.44×10⁻⁴ 约 0.78×10⁻⁴ 0.544×

所以 Block 的首尾比更接近 1,并不是因为它把晚层全部抬高到 Baseline 早层的强度。更接近的 描述是:

整体尺度下降,早层系统性高值被削弱,同时某些中后段位置相对全层均值形成尖峰。

这也是只看 normalized curve 或只看 CV 都不够的原因。

8. 参数梯度没有翻转 Round 04

最终核心参数梯度 CV:

depth Baseline mean Block mean Block / Baseline
16 0.41596 0.68287 1.64×
32 0.39661 0.77176 1.95×

6 / 6 个配对中 Block 都更高。Round 04 的反结果不是在把梯度对象改成 activation 后自动消失; 两种梯度对象在本轮 final endpoint 都显示更高的跨层 CV。

但 activation gradient 又显示首尾失衡大幅改善,这说明“均匀”至少要拆成:

首尾是否平衡
全层是否有尖峰
绝对尺度多大
参数更新信号是否平衡

一个标量不能代替全部。

9. Output RMS:最接近论文叙述的正向结果

三 seed 最终 post-MLP output RMS 的最后/第一层比:

depth Baseline Block
16 4.59× 1.17×
32 6.08× 1.89×

Baseline output magnitude 随深度明显累积;Block 把增长限制在 aggregation group 内并产生 周期性 reset。这个缩小实验的 output-RMS 方向与论文 Figure 5(b) 的叙述一致。

仍不能把数值直接叠到论文图上:

  • 模型宽度、深度和数据不同;
  • 训练 Token 相差巨大;
  • 论文图的确切 output norm / reduction 也未完整公开;
  • 本实验的 Block group 是 8 组固定设计。

10. 验证 BPC 与真实成本

最终 BPC:

depth-16

seed Baseline Block Block − Base
2026073001 1.74880 1.73739 −0.01141
2026073002 1.73683 1.73118 −0.00565
2026073003 1.73637 1.72661 −0.00976
均值 1.74067 1.73173 −0.00894

depth-32

seed Baseline Block Block − Base
2026073001 1.71790 1.71235 −0.00554
2026073002 1.72698 1.70932 −0.01765
2026073003 1.70951 1.70310 −0.00642
均值 1.71813 1.70826 −0.00987

6 / 6 为负。这是有价值的次要方向,但 Round 05 没有为 BPC 再预注册一个新的 support 阈值, 因此不追加事后显著性结论。

真实成本:

depth Base mean ms Block mean ms time ratio Base peak alloc Block peak alloc memory ratio
16 21.47 54.71 2.55× 3.04 GB 6.54 GB 2.15×
32 42.11 109.38 2.60× 5.94 GB 12.82 GB 2.16×

这个简单 eager 实现没有论文训练系统的 kernel、并行或工程优化;成本数值不应外推到 K3。 但它足以说明本站的 BPC 对比不是同 wall time / FLOPs。

11. 预注册判定为何是 mixed

支持需要:

CV:3 / 3 seeds 改善,平均相对改善 ≥20%
AND
imbalance:3 / 3 seeds 改善,平均相对改善 ≥20%

concern 需要两项都以相同规则恶化。

实际:

depth-16:
CV 3 / 3 恶化,平均 10.3%
imbalance 3 / 3 改善,平均 61.0%

depth-32:
CV 3 / 3 恶化,平均 60.0%
imbalance 3 / 3 改善,平均 72.0%

两个指标相反,所以两个 depth 都是 mixed / inconclusive at this depth。这不是“数据没规律”, 而是预注册的“更均匀”概念被实验拆成了两个方向相反的组成部分。

12. 开放工件与校验哈希

工件 SHA-256
manifest 080afb17…6371
protocol f772629b…ce22
definition audit 79221c56…cc6
runner 04ae69e1…800f
analyzer 017d38d9…2fc8
aggregate canonical 69be133c…7b51
compact canonical 8cdb7180…044f
reproduction canonical addb2e59…f68a

公开目录包含:

  • 12 个完整 formal raw JSON;
  • 8 个两套 smoke raw JSON;
  • 1 个完整 replay raw JSON;
  • 完整 aggregate;
  • 网站 compact payload;
  • manifest、runner、analyzer 与 reproduction 清单;
  • 前置定义审计、本协议和本结果审计。

experiments/k3/attnres_gradient/reproduction.json 记录每个 raw 文件 SHA-256。

13. 允许和禁止的结论

允许:

在本轮公开定义下,Block AttnRes 一致缓解了首/末深度四分位失衡,但在中后段形成局部 梯度尖峰,导致全层 CV 一致升高;因此“梯度更均匀”必须拆成多个指标解释。

同 token / step 预算下,Block 的最终验证 BPC 在 6 / 6 个配对中更低,但实际运行成本 约为 Baseline 的 2.6× step time 与 2.2× peak allocated memory。

禁止:

  • “复现了论文 Figure 5(c)”;
  • “论文的梯度结论是错的”;
  • “已测到 Kimi K3 checkpoint 的真实梯度”;
  • “Block 解决了梯度消失 / 爆炸”;
  • “CV 更高就代表训练一定更不稳定”;
  • “BPC 改善是同 FLOPs / wall time 优势”;
  • 从 3 seeds 推导总体显著性;
  • 从 depth 16 / 32 外推到 48B、1T+400B Token 或 K3 2.8T 参数。

14. 下一步最值得问什么

这轮已经把“梯度”从一个模糊词拆成了可复查对象。下一个有价值的问题不是再换一个漂亮 汇总指标,而是追踪局部尖峰从哪里来:

  1. 分开捕获 pre-attention 与 pre-MLP residual positions;
  2. 把 layer 21–25 的 activation-gradient 与 mixer source weights 同步对齐;
  3. 比较 aggregation-group boundary 前后;
  4. 在不改变 formal 结果的前提下,对相同 raw gradient tensor 做多种公开 reduction sensitivity analysis;
  5. 若官方之后发布 Figure 5 telemetry 代码,再按其定义单独开新 protocol。

这些属于后续轮次,不能倒写进本轮预注册结论。