Files
llm-atlas/research/K3_ATTNRES_GRADIENT_SCALE_PROTOCOL.md
T
2026-07-30 07:53:12 +08:00

16 KiB
Raw Blame History

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. 设计总览

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 起点:

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。

数值精度进一步固定为:attention / MLP 线性分支受 BF16 autocast;embedding、Baseline hidden residual stream 与 Block aggregation partial 都以 FP32 累加。也就是说,Block 每个 BF16 branch output 在进入 partial 前显式转为 FP32。这样两种结构被捕获的 h_l 都是 FP32, 不会把 residual accumulator 精度差异混进梯度形状对比。

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 子层为:

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

验证和诊断都发生在:

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]、dtype 为 FP32。实现必须为 diagnostic forward 返回独立 引用列表,不允许捕获 reset 后的零张量,不允许把 8 个 aggregation sources 当成 16 / 32 个 Transformer outputs。

6.2 diagnostic loss 与 gradient magnitude

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 计算:

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 绝对尺度

mean_grad = mean_l(m_l)

它防止“曲线更平只是全部梯度趋近于零”被 CV 隐藏。绝对尺度不设置优劣阈值,只公开。

7.2 归一化谱与 CV

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 前后四分位失衡

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:

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 的参数梯度 拼接后计算:

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 必须各自在全新进程中执行。目录键为:

depth-{16|32}/{baseline|block}/seed-{2026073001|2026073002|2026073003}

完成后预注册 replay:

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. 公开工件

正式结果完成后仓库必须包含:

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。