16 KiB
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 的语料和预处理:
- 按固定 parquet 行序读取
text; - 每行追加一个
\n; - UTF-8 编码,无 normalization、strip、去空行或大小写改写;
- byte vocabulary
0..255; - 每个窗口连续取 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×256targets 做算术平均,无 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 格都必须通过:
- 捕获数量严格等于 depth,shape 均为
[16,256,192]; - 全部 gradient 非
None、finite、storage 不别名; - 同输入把 loss 乘 2 后,每层
m_l比值在2±1e-5; - loss×2 前后 CV、normalized spectrum、quartile ratio 在
1e-6绝对容差内; - 全新进程重复 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
网站至少提供五个互相联动的视图:
- 论文 Figure 5 的“已知 / 未定义”拆解;
- activation gradient 绝对谱与 normalized spectrum;
- depth 16 / 32、三 seed、六时间点对比;
- output RMS 周期与 Block aggregation boundary;
- 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。