wuyang
199 lines
9.1 KiB
Markdown
199 lines
9.1 KiB
Markdown
# 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) 所说的“梯度分布更均匀”，因为两边很可能测的
|
||
不是同一个对象：
|
||
|
||
```text
|
||
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) 的完整对象描述分别是：
|
||
|
||
```text
|
||
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：
|
||
|
||
```text
|
||
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`。
|