OKLS 把 KL-optimal Shampoo 带到语言模型训练
aryaman2020 · x · 2026-07-29
OKLS 想把 Shampoo 做成更接近全矩阵 AdaGrad 的训练优化器
这篇工作提出 Online KL Shampoo（OKLS），目标是在语言模型训练里，比传统的对角优化器更好地捕捉梯度坐标之间的相关性，同时又不至于像全矩阵 AdaGrad 那样带来过高状态开销。
- 方法核心是用 KL-optimal Kronecker factors，在行和列两个方向同时对白化矩阵梯度，让优化器更接近完整二阶几何。
- 论文认为它补上了 Muon 的一个缺口：Muon 能考虑相关性，但没有把历史信息纳入预条件。
- 最大工程难点是每一步都要得到“新鲜”的逆平方根预条件矩阵；作者用 Scaled CANS Coupled Newton–Schulz 把零陈旧预条件做成可行方案。
- 实现细节包括 10 次迭代、27 个 FP16 GEMM 和 FP32 累加。
- 结果上，OKLS 在 200M–1B 规模实验里达到 1.45× 的参数效率，同时保住了 98% 的训练吞吐。
「研究」频道最新
- Anthropic Claude 破解后量子密码，加速 AES 攻击数百倍 — petrusenko_max · 2026-07-29
- 技术报告解读：外部模型初始化视觉编码器并非良策 — SeunghyunSEO7 · 2026-07-29
- AI 药物发现流程已覆盖小分子和纳米抗体实验验证 — const_reborn · 2026-07-29
- Jon Durbin 用每小时不到 10 美元预训练 20B MoE — const_reborn · 2026-07-29
- 商业 LLM 预训练再多，为什么下棋还是不行 — alex_peys · 2026-07-29
- 微软发布 Mage-VL：编码原生流式多模态模型，推理提速 3.5 倍 — pmttyji · 2026-07-29