research: preregister AttnRes gradient scale study

2026-07-30 07:46:07 +08:00
parent 40f6a1d631
commit 091a05f0a0
2 changed files with 651 additions and 0 deletions
@@ -0,0 +1,453 @@
+# Kimi K3 第五轮：Attention Residuals 梯度定义与深度扩展实验协议
+
+> 协议 ID：`llm-atlas-k3-attnres-gradient-scale-v1`  
+> 冻结日期：2026-07-30（Asia/Shanghai）  
+> 状态：正式输出前预注册  
+> 前置定义审计：`research/K3_ATTNRES_GRADIENT_DEFINITION_AUDIT.md`
+
+## 0. 目标与一句话研究问题
+
+Round 04 在缩小模型中得到两个同时成立的结果：
+
+- Full / Block AttnRes 的 2,000-step 验证 BPC 都优于 Baseline；
+- 以“每个 Transformer block 的核心**参数**梯度 RMS”定义时，跨深度 CV 反而更高。
+
+论文 Figure 5(c) 没有公开足以唯一恢复的梯度测量合同。Round 05 不猜作者的隐藏代码，而是
+冻结一个可复现、与 Figure 5 的 output/gradient 并列叙述对齐的 activation-gradient 定义，
+再问：
+
+> 当深度从 16 增至 32、训练预算从 2,000 增至 8,000 step 时，Block AttnRes 是否比
+> PreNorm Baseline 更一致地降低 post-MLP block-output activation gradient 的跨深度失衡？
+
+这不是 K3 checkpoint forward，也不是论文 Figure 5 数值复画。
+
+## 1. 一手来源与不可补写的空白
+
+| 工件 | 固定 revision / checksum | 用途 |
+|---|---|---|
+| Attention Residuals GitHub | `85e22310fe5ee860b4a023de312d791de8a5a5e6` | 公式、Figure 5 / 8、模型尺度 |
+| `Attention_Residuals.pdf` | SHA-256 `e5831b0d…a26b2f` | 论文一手图文 |
+| WikiText-2 raw | `Salesforce/wikitext@b08601e04326c79dfdd32d625aee71d232d685c3` | 固定公开训练语料 |
+| Round 04 protocol | `llm-atlas-k3-attnres-reduced-v1` | 公共主干与数据合同来源 |
+
+官方仓库没有模型实现、训练脚本、checkpoint 或 Figure 5 原始数据。以下字段不能归因给论文：
+
+- Figure 5 的确切 gradient tensor；
+- norm / reduction；
+- diagnostic batch；
+- AMP / clipping 时点；
+- 单点还是时间平均。
+
+Grok Headless 只进行一次对抗式方法检查；其建议和错误都在前置审计中公开，不是事实来源。
+
+## 2. 设计总览
+
+```text
+2 个深度：16 / 32 Transformer blocks
+× 2 个 residual graph：PreNorm Baseline / Block AttnRes
+× 3 个冻结 seed
+× 8,000 training steps
+× 32 windows/step
+× 256 target bytes/window
+= 12 个正式训练格
+= 786,432,000 target bytes
+```
+
+只比较 Baseline 与 Block，因为论文 Figure 5 的训练动力学面板也是这两个结构的直接对照。
+Round 04 的 Full AttnRes 结果保持公开，但本轮不增加一个与主问题无关的 6-run 分支。
+
+## 3. 数据合同
+
+继承 Round 04 的语料和预处理：
+
+1. 按固定 parquet 行序读取 `text`；
+2. 每行追加一个 `\n`；
+3. UTF-8 编码，无 normalization、strip、去空行或大小写改写；
+4. byte vocabulary `0..255`；
+5. 每个窗口连续取 257 bytes，前 256 预测后 256。
+
+固定拼接后 split：
+
+| split | bytes | SHA-256 |
+|---|---:|---|
+| train | 10,951,563 | `0ca7d3e74dbe44564ea5942b85232f1bbcb525c9cd481cd5d28a87ee90e7e9b4` |
+| validation | 1,148,008 | `a42356f6a8ff1d25daf25ec9db49e10a537c265581b61c74604bb63231dee719` |
+| test | 1,292,014 | `bfe9eb16ab9987fb88bde4ea9a30a00f2a45db01dfc14bad78d05325789c4f12` |
+
+训练第 `step`、第 `row` 的 window 起点：
+
+```text
+z = first 8 bytes of SHA256(
+  protocol_id + "\0train-window\0" + seed + "\0" + step + "\0" + row
+)
+start = uint64_be(z) mod (len(train_bytes) - 257)
+```
+
+同一 seed 的 4 个结构格逐 step / row 使用完全相同的 token tensor。validation 64 windows、
+diagnostic 16 windows，分别由标签 `validation-window` / `diagnostic-window` 与固定 index
+生成，对全部结构与 seed 相同。
+
+正式运行前 manifest 必须记录 parquet hash、split bytes/hash、全部 `3×8,000×32=768,000`
+唯一训练窗口起点的 schedule hash、validation tensor hash 与 diagnostic tensor hash。
+
+## 4. 模型合同
+
+### 4.1 两个深度共享的结构
+
+| 项 | 固定值 |
+|---|---:|
+| vocabulary | 256 bytes |
+| context | 256 |
+| `d_model` | 192 |
+| heads / head dimension | 6 / 32 |
+| `d_ff` | 768 |
+| dropout | 0 |
+| positional embedding | learned absolute，256 × 192 |
+| norm | RMSNorm，`eps=1e-6` |
+| attention | causal MHA；score 以 FP32 softmax |
+| MLP | bias-free SwiGLU，`192→768`, `192→768`, `768→192` |
+| embedding / readout | tied；final RMSNorm 后乘 token embedding |
+
+全部 bias-free linear 与 embedding 初始化为 `N(0,0.02)`。attention output projection 与 MLP
+down projection 的标准差为 `0.02 / sqrt(2×depth)`；普通 RMSNorm 为 1。
+
+### 4.2 深度与 Block AttnRes 聚合
+
+| Transformer depth | residual sublayers | aggregation groups | sublayers/group | Transformer blocks/group |
+|---:|---:|---:|---:|---:|
+| 16 | 32 | 8 | 4 | 2 |
+| 32 | 64 | 8 | 8 | 4 |
+
+Baseline 子层为：
+
+```text
+h ← h + f(RMSNorm(h))
+```
+
+Block AttnRes：
+
+- embedding 永远是 source 0；
+- 对已经完成的 aggregation-group sums 做跨组 softmax mixture；
+- 组内 attention / MLP branch output 累加到 `partial`；
+- 达到组边界时，保存完整 `partial` 为新 source，再开始下一组；
+- output mixer 聚合 embedding + 8 个完整 group sums。
+
+每个子层的 pseudo-query 为 `d_model` 向量，严格 zero-init；每个 source 的 key RMSNorm weight
+严格 one-init。Baseline / Block 不要求总参数量或 residual-mixer FLOPs 相等，但同 seed /
+depth 的 token/position embedding、attention、MLP、input norm、final norm 和 tied readout
+必须逐 tensor SHA-256 exact。
+
+## 5. 训练合同
+
+| 项 | 固定值 |
+|---|---:|
+| seeds | `2026073001, 2026073002, 2026073003` |
+| formal steps | 8,000 |
+| batch | 32 |
+| context | 256 |
+| target bytes / run | 65,536,000 |
+| optimizer | AdamW |
+| betas / epsilon | `(0.9,0.95)` / `1e-8` |
+| peak / min LR | `3e-4` / `3e-5` |
+| warmup | 400 steps，linear |
+| decay | cosine，step 400→8,000 |
+| weight decay | `0.1` if `ndim>=2`，否则 `0` |
+| parameter grad clip | global norm `1.0` |
+| compute | BF16 autocast；FP32 optimizer state |
+| compile | off / eager |
+| device | one RTX 5090 |
+| deterministic | deterministic algorithms；`CUBLAS_WORKSPACE_CONFIG=:4096:8` |
+
+验证和诊断都发生在：
+
+```text
+step 0, 100, 500, 2,000, 4,000, 8,000
+```
+
+验证固定 64 windows，8 windows/eval batch，报告 token-mean CE nats 与
+`bits_per_byte = CE / ln(2)`。诊断固定 16 windows，一次性输入，不改变 optimizer state。
+
+计时合同：
+
+- 前 20 个 training step 不计入；
+- step 21–8,000 每步前后 CUDA synchronize；
+- step 20 后 reset peak memory；
+- 报告 mean / median / p95 step ms、peak allocated / reserved；
+- timing、wall clock、hostname、GPU temperature 不进入 exact replay 字段。
+
+## 6. 主梯度对象：精确到代码位置
+
+### 6.1 `h_l` 定义
+
+对 `l=1..depth`，统一在第 `l` 个 Transformer block 的 MLP branch 完成后捕获：
+
+| 架构 | 精确定义 |
+|---|---|
+| Baseline | attention residual 与 MLP residual 都完成后的 hidden state |
+| Block | MLP branch 已加入、本 aggregation partial 可能保存/reset **之前**的 partial |
+
+所有 `h_l.shape = [16,256,192]`。实现必须为 diagnostic forward 返回独立引用列表，不允许捕获
+reset 后的零张量，不允许把 8 个 aggregation sources 当成 16 / 32 个 Transformer outputs。
+
+### 6.2 diagnostic loss 与 gradient magnitude
+
+```text
+model.eval()
+logits, h[1..depth] = forward(fixed_diagnostic_x, capture=true)
+L = mean(cross_entropy(logits.float(), fixed_diagnostic_y))
+backward(L)          # 不使用 GradScaler，不执行 optimizer.step
+m_l = sqrt(mean(float32(h_l.grad)² over batch×time×channel))
+```
+
+规则：
+
+- forward 仍使用与训练一致的 BF16 autocast；
+- logits 在 FP32 中计算 CE；
+- loss 对全部 `16×256` targets 做算术平均，无 mask、无 label smoothing；
+- backward 前 model / optimizer gradients 清零；
+- activation gradient 在任何 parameter clipping 之前读取；
+- diagnostic 不消耗训练数据，不进入 optimizer，不改变学习率或模型状态。
+
+### 6.3 同位置 output magnitude
+
+同一批 `h_l` 计算：
+
+```text
+o_l = sqrt(mean(float32(h_l)² over batch×time×channel))
+```
+
+它用于显示 Baseline 单调累积与 Block 组内周期，而不作为主 confirmatory endpoint。
+
+## 7. 主指标与预注册判据
+
+每个 diagnostic step 保存完整 `m_1..m_depth`。以下统计由未四舍五入的 float64 数组计算。
+
+### 7.1 绝对尺度
+
+```text
+mean_grad = mean_l(m_l)
+```
+
+它防止“曲线更平只是全部梯度趋近于零”被 CV 隐藏。绝对尺度不设置优劣阈值，只公开。
+
+### 7.2 归一化谱与 CV
+
+```text
+n_l = m_l / mean_grad
+CV = population_std_l(m_l) / mean_grad
+```
+
+使用 population standard deviation（`ddof=0`）。网站必须同时显示 `m_l` 与 `n_l`，不能只显示
+CV 排名。
+
+### 7.3 前后四分位失衡
+
+```text
+q = depth / 4
+Q_first = mean(m_1 .. m_q)
+Q_last  = mean(m_(depth-q+1) .. m_depth)
+imbalance = abs(ln(Q_first / Q_last))
+```
+
+depth 16 时各取 4 层；depth 32 时各取 8 层。`imbalance=0` 才表示首尾一致；这个定义不会把
+“早层偏大”和“晚层偏大”错误地都解释为越小越好。原始有符号比 `Q_first/Q_last` 仍公开。
+
+### 7.4 seed 内配对对比
+
+只在最终 step 8,000 做 confirmatory verdict：
+
+```text
+relative_CV_reduction
+  = (CV_baseline - CV_block) / CV_baseline
+
+relative_imbalance_reduction
+  = (imbalance_baseline - imbalance_block) / imbalance_baseline
+```
+
+若 Baseline imbalance 精确为 0，则该 seed 的 relative imbalance reduction 定义为不可计算，
+该 depth 自动不能得到“联合支持”；仍公开绝对差。
+
+对每个 depth 分别判定：
+
+- **joint directional support at this depth**：三个 seed 的 CV reduction 都 `>0`，其均值
+  `>=20%`；同时三个 seed 的 imbalance reduction 都 `>0`，其均值 `>=20%`。
+- **joint directional concern at this depth**：三个 seed 的两项 reduction 都 `<0`，且两项
+  平均相对恶化都 `>=20%`。
+- 其他：**mixed / inconclusive at this depth**。
+
+总判定：
+
+- 两个 depth 都 support：**scale-consistent directional support in this operationalization**；
+- 两个 depth 都 concern：**scale-consistent directional concern in this operationalization**；
+- 其他：**depth-dependent or inconclusive**。
+
+不计算 p-value、population CI，不把 3 seeds 称为统计证明。
+
+## 8. 必须公开的次要指标
+
+### 8.1 全时间轨迹
+
+六个预注册时点的以下数据必须全部公开，不能选择“最好看”的 checkpoint：
+
+- validation BPC；
+- absolute activation-gradient spectrum；
+- normalized activation-gradient spectrum；
+- activation CV；
+- first/last quartile ratio 与 imbalance；
+- output RMS spectrum；
+- mean activation-gradient scale。
+
+### 8.2 参数梯度
+
+延续 Round 04 定义：每个 Transformer block 的 attention、MLP 与两个 input norm 的参数梯度
+拼接后计算：
+
+```text
+parameter_grad_rms[l] = sqrt(sum(g²) / total_parameter_elements)
+```
+
+不含 embedding、final norm、LM head 与 AttnRes mixer 参数；clip 前读取。报告完整谱、CV 和
+前后四分位失衡，但它们不进入 Round 05 主判定。
+
+### 8.3 mixer / 成本
+
+Block 同报：
+
+- 各子层 softmax mixture 的 source-depth 分布；
+- output mixer 分布；
+- entropy 与 embedding / latest-complete-group mass；
+- step time、peak allocated/reserved、参数量。
+
+这些用于解释机制与成本，不改变 confirmatory verdict。
+
+## 9. 运行前闸门
+
+### 9.1 数据闸门
+
+- split bytes/hash 与 Round 04 exact；
+- 新 protocol 的 768,000-window schedule hash 落盘；
+- validation / diagnostic tensor hash 落盘；
+- 同 seed 四结构的至少 step 0 / 1 / 7,999 输入 tensor hash exact。
+
+### 9.2 公共权重闸门
+
+每个 depth / seed 的 Baseline 与 Block 公共参数逐 tensor exact；输出：
+
+- 公共参数 tensor 数；
+- 公共参数 element 数；
+- name / shape / dtype / bytes 联合 hash；
+- value bytes 联合 hash。
+
+### 9.3 activation-gradient 闸门
+
+四个 depth×architecture 格都必须通过：
+
+1. 捕获数量严格等于 depth，shape 均为 `[16,256,192]`；
+2. 全部 gradient 非 `None`、finite、storage 不别名；
+3. 同输入把 loss 乘 2 后，每层 `m_l` 比值在 `2±1e-5`；
+4. loss×2 前后 CV、normalized spectrum、quartile ratio 在 `1e-6` 绝对容差内；
+5. 全新进程重复 smoke 的冻结字段 exact。
+
+### 9.4 smoke
+
+四个格都运行 20 training steps；冻结字段包括：
+
+- protocol、architecture、depth、seed、device、dtype；
+- data/schedule/tensor hashes；
+- public-weight hashes；
+- step 0 / 20 loss 与 validation；
+- 全部预注册 diagnostic 数组；
+- finite / alias / loss-scale checks。
+
+smoke 不能写入 formal 目录。
+
+## 10. 正式执行与独立 replay
+
+12 个 formal grid 必须各自在全新进程中执行。目录键为：
+
+```text
+depth-{16|32}/{baseline|block}/seed-{2026073001|2026073002|2026073003}
+```
+
+完成后预注册 replay：
+
+```text
+depth-32 / block / seed-2026073001
+```
+
+replay 再从初始化训练完整 8,000 steps，不加载 formal checkpoint。比较时排除：
+
+- wall time / step-time samples；
+- peak memory；
+- process ID / hostname；
+- GPU 温度与驱动层瞬时字段；
+- 文件路径和生成时间。
+
+必须 exact 的字段：
+
+- 数据与公共权重 hashes；
+- 全部 validation CE/BPC；
+- 全部 activation/output/parameter gradient 数组；
+- mixer 数组与 summary；
+- final model-state tensor hash；
+- optimizer-state tensor hash；
+- training-loss checkpoint 数组。
+
+若 replay 不 exact，停止聚合并公开失败，不挑选另一 seed 替代。
+
+## 11. 公开工件
+
+正式结果完成后仓库必须包含：
+
+```text
+experiments/k3/attnres_gradient/
+  README.md
+  build_dataset.py
+  manifest.json
+  train.py
+  analyze.py
+  results/raw/*.json
+  results/compact.json
+  reproduction.json
+
+research/
+  K3_ATTNRES_GRADIENT_DEFINITION_AUDIT.md
+  K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md
+  K3_ATTNRES_GRADIENT_SCALE_AUDIT.md
+```
+
+网站至少提供五个互相联动的视图：
+
+1. 论文 Figure 5 的“已知 / 未定义”拆解；
+2. activation gradient 绝对谱与 normalized spectrum；
+3. depth 16 / 32、三 seed、六时间点对比；
+4. output RMS 周期与 Block aggregation boundary；
+5. activation gradient / parameter gradient 并排，以及判定、成本、哈希和声明边界。
+
+图中必须能切换到所有负结果；不能只放均值、只放 final 或隐藏某个 seed。
+
+## 12. 允许与禁止的结论
+
+若达到支持条件，允许写：
+
+> 在这个公开定义、两种缩小深度和 8,000-step byte-LM 合同中，Block AttnRes 方向一致地
+> 降低了 post-MLP output activation gradient 的跨深度 CV 与首尾四分位失衡。
+
+无论结果怎样，都禁止写：
+
+- 复现了论文 Figure 5 的数值；
+- 证明了论文未公开实现采用相同梯度定义；
+- 证明了 Kimi K3 的真实梯度更健康；
+- 证明 AttnRes 解决梯度消失、梯度爆炸或训练稳定性的全部问题；
+- 从 3 seeds 推导总体显著性；
+- 从 depth 16 / 32 外推至 48B、1T+400B Token 或 K3 2.8T 参数；
+- 隐藏 activation 与 parameter gradient 方向不一致的结果。
+
+## 13. 变更纪律
+
+本文件提交后：
+
+- 允许修复使实现符合本协议的 bug；
+- 允许补充日志、注释、可视化和不改变数值的导出；
+- 不允许看过 formal 结果后修改主指标、阈值、diagnostic step、seed、训练预算或 replay 格；
+- 任何不得不改变实验合同的事项必须先停止、写入审计、升级 protocol ID，再重新执行全部 grid。