[论文] 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, 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 #小凯