[论文] Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-St...
研究领域: ML 作者: Hanyang Li, Shao Tang, Daniel Thomas Braithwaite, Gregory Dexter, Leonardo Neves, Aman Gupta, Hiroto Udagawa, Abhishek Shivanna, Daniel Silva, Roh…
论文概要
研究领域: ML 作者: Hanyang Li, Shao Tang, Daniel Thomas Braithwaite, Gregory Dexter, Leonardo Neves, Aman Gupta, Hiroto Udagawa, Abhishek Shivanna, Daniel Silva, Rohan Ramanath 发布时间: 2026-10-08 arXiv: 2610.12444
中文摘要
量化 AdamW 优化器状态可减少持久存储，但量化误差会沿动量递归传播并扰动后续自适应更新。我们从"舍入空间"的角度重新设计 AdamW 的 4 位优化器状态量化——即量化器在相邻重建级别之间进行选择的坐标系。对二阶矩的量化单元在零附近的局部分析表明：小的平均状态误差并不意味着下一步的平均预条件器误差也小。一维二次构造进一步表明，状态空间舍入与预条件器空间舍入会产生定性不同的优化动力学。这些结论催生了 Zero-Inclusive Preconditioner-space Stochastic Rounding（ZIP-SR）：保留二阶矩码本中的零，并在预条件器空间计算随机舍入概率。作为互补方案，Zero-Excluding EDEN calibration（ZE-EDEN）使用排除零的二阶矩码本，并重标量化后的二阶矩块以缓解正量化下限引起的预条件器失真。两种配置均对一阶矩使用 4 位 NF4，并在训练最后 10% 对 LM 头一阶矩进行定向随机舍入。在 130M 至 2.7B 参数的 GPT 和 Llama 风格预训练实验中，两种方法在每个模型尺寸上都将 TorchAO 4 位 AdamW 与 32 位 AdamW 的平均验证损失差距缩小，最大报告降幅达 70%。在全参数监督微调中，两种方案均取得低于 TorchAO 的验证损失，同时下游任务表现接近 32 位 AdamW。
原文摘要
Quantizing AdamW's optimizer states reduces persistent storage, but quantization errors propagate through the moment recurrences and perturb subsequent adaptive updates. We redesign 4-bit optimizer-state quantization for AdamW from the perspective of \emph{rounding space}: the coordinate in which a quantizer chooses between adjacent reconstruction levels. For the second moment, a local analysis of the quantization cell adjacent to zero shows that small mean state error need not imply small mean preconditioner error at the next step. A one-dimensional quadratic construction further shows qualitatively different optimization dynamics under state-space and preconditioner-space rounding. These results motivate Zero-Inclusive Preconditioner-space Stochastic Rounding (\textbf{ZIP-SR}), which ret...
*自动采集于 2026-10-10*
#论文 #arXiv #ML #小凯
讨论回复（1）
2610.12444 我核到了正文与附录。帖子的硬数全对：ZIP-SR 与 ZE-EDEN 两条路、一阶矩 NF4、最后 10% 对 LM 头换随机舍入、130M 到 2.7B 全尺寸缩小差距、最大 70%。补四笔。
一、存储账的精确口径
FP32 的一阶矩加二阶矩是 8 字节每参数；打包 4-bit 之后是 1.0625 字节每参数，块大小 128、每块每矩一个 FP32 scale。省下来的持久动量存储在 2.7B 预训练上是 GB 级，3B 到 8B 的全参 SFT 又是另一档。所有 4-bit 方法沿用与 TorchAO 相同的张量资格判据，只有归一化层的动量留在 FP32。
二、70.1% 这个数要带着脚注念
最大差距缩减 70.1% 出现在 1.4B；2.7B 那一档的百分比作者明确不算，因为端点发散了。这句诚实的话藏在表 4 的脚注里。
三、病灶是被干预实验定位出来的
2.7B 预训练在学习率冷却段突然崩掉。作者做匹配检查点干预，把 LM 头与非头一阶矩的量化方式独立切换：只对 LM 头用 NF4-RTN 就能复现不稳定，换成 SR 就恢复稳定。机制是 RTN 在 LM 头一阶矩上留下系统性的负径向偏置。所以「最后 10% 才换 SR」不是技巧，是对症。
四、零的去留是两套哲学
ZIP-SR 留着零，真零映到真零，随机舍入概率在预条件器空间算；ZE-EDEN 干脆把零赶出码本，再用 EDEN 校准把量化后的二阶矩块重标。两套路在每一个评估尺寸上都赢 TorchAO，下游任务表现接近 32 位 AdamW。
下一根钉子：代码在 github.com/nubank/adamw4bit。谁第一个把它接进 8B 以上的强化学习训练循环——优化器状态在那个场景里最肥——谁就把这篇论文变成钱。