预条件子空间的舍入：重新设计 4-bit AdamW 优
来自 arXiv · 期号 2026-10-09
摘要
AdamW 优化器的状态量做 4-bit 量化时，误差会沿着动量递归累积，扭曲后续的自适应更新。 作者从「舍入空间」的角度重新审视这个问题：量化器在相邻重建值之间做选择时，用的是哪个坐标系。对二阶矩靠近零的量化单元做局部分析后发现，状态量的均值误差小，并不代表下一步的预条件子误差也小。 他们提出两种方法：ZIP-SR 在二阶矩码本里保留零值，在预条件子空间里计算随机舍入概率；ZE-EDEN 则用不含零的码本，再对量化后的二阶矩块做缩放，补偿正量化底噪带来的预条件子畸变。两者都用 NF4 量化一阶矩，并在训练最后 10% 对 LM head 的一阶矩做定向随机舍入。 在 130M 到 2.7B 参数的 GPT 和 Llama 预训练实验中，两种方法都比 TorchAO 的 4-bit 基线更接近 32-bit AdamW 的验证损失，最大差距缩小了 70%。全参数微调时，两者验证损失也低于 TorchAO，下游任务表现接近 32-bit AdamW。
深读
导读 这篇论文从「舍入空间」这一新视角重新审视了 AdamW 优化器状态 4-bit 量化的误差来源，发现二阶矩的量化误差在预条件子空间中会被放大。作者据此提出了两种改进方法，在 130M 到 2.7B 参数的预训练实验中，将 4-bit 量化与 32-bit 全精度之间的验证损失差距缩小了最多 70%。对于需要在大模型训练中节省显存、同时尽量保持训练质量的从业者，这是一个值得关注的优化方向。 关键信息 - 论文提出的两种方法——ZIP-SR（保留零值的预条件子空间随机舍入）和 ZE-EDEN（不含零值的 EDEN 校准）——均使用 NF4 量化一阶矩，并在训练最后 10% 对 LM head 的一阶矩做定向随机舍入。 - 在 GPT 和 Llama 两种架构的预训练实验中，模型规模从 130M 到 2.7B 参数，两种方法在所有规模上都比 TorchAO 的 4-bit 基线更接近 32-bit AdamW 的验证损失。 - 最大的验证损失差距缩减达到 70%。 - 在全参数监督微调中，两种方法在下游任务上的验证损失均低于 TorchAO 基线，且接近 32-bit AdamW。 逐节讲解 问题：4-bit 量化误差在动量递归中累积 AdamW 优化器需要存储一阶矩（动量）和二阶矩（自适应学习率），量化这些状态量可以减少显存占用。但论文指出，量化误差并非独立存在，它会沿着动量的递归更新不断累积，进而扭曲后续的自适应更新。作者从「舍入空间」的角度重新审视这个问题：量化器在相邻重建值之间做选择时，用的是哪个坐标系。 核心发现：状态误差小不等于预条件子误差小 论文对二阶矩靠近零的量化单元做了局部分析。关键结论是：二阶矩状态量的均值误差很小，并不代表下一步的预条件子（即自适应学习率的分母部分）误差也小。作者还通过一个一维二次型构造，展示了在状态空间和预条件子空间中做舍入会导致截然不同的优化动态。这意味着，传统的以状态误差最小化为目标的量化方法，可能恰恰在最重要的地方引入了大误差。 两种解决方案：ZIP-SR 与 ZE-EDEN 论文提出两条互补的技术路线： - ZIP-SR（Zero-Inclusive Preconditioner-space Stochastic Rounding）：在二阶矩的码本中保留零值，并在预条件子空间（而非状态空间）中计算随机舍入的概率。这样量化器在做选择时，直接以对下一步预条件子的影响为优化目标。 - ZE-EDEN（Zero-Excluding EDEN calibration）：使用不含零值的二阶矩码本，然后对量化后的二阶矩块做缩放，以补偿正量化底噪（即量化后所有值都大于零）带来的预条件子畸变。 两种方法都共用一阶矩的 NF4 量化方案，并在训练最后 10% 对 LM head 的一阶矩做定向随机舍入，以保护对最终损失影响最大的参…