4 比特 AdamW 的舍入空间之争：为什么要在预条件子空间里做随机舍入

AdamW 是当今大模型训练的默认优化器，但它的代价常被低估。除了权重和梯度，每个参数还要保存一阶矩与二阶矩两份状态，按 32 位存储时，这部分常驻内存与权重本身相当甚至更多。把优化器状态压到 4 比特，是削减训练显存最直接的办法之一。TorchAO 已经提供了 4 比特 AdamW 的实现，但它与 32 位 AdamW 之间仍有可测的验证损失差距。Hanyang Li 等五位作者在 2026 年 10 月 8 日发布于 arXiv cs.LG 的这篇论文，没有去调码本的形状或块的大小，而是追问一个更基础的问题：量化器在相邻两个重建电平之间做选择时，究竟应该在哪个坐标里舍入。 作者称这个坐标为“舍入空间”。问题的根源在于误差不是一次性的。量化误差会沿着矩的递推传播，扰动之后每一步的自适应更新。对一阶矩而言，更新量对状态近似线性，在状态空间里舍入问题不大。二阶矩则不同：Adam 的更新与二阶矩平方根的倒数成正比，这个倒数就是预条件子。论文对紧邻零的那个量化单元做了局部分析，结论是：状态的平均误差很小，并不意味着下一步预条件子的平均误差也很小。状态稍微偏离零，倒数函数就被急剧放大，误差在预条件子一侧被成倍拉高。作者还构造了一个一维二次问题，显示在状态空间舍入与在预条件子空间舍入，会产生性质上不同的优化动力学，而不只是数值上的差异。 这些观察催生了第一个方案 ZIP-SR，全称“保留零的预条件子空间随机舍入”。它在二阶矩码本里保留零这个电平，并且在预条件子空间而不是状态空间里计算随机舍入的概率。随机舍入的好处是期望无偏，把确定性的系统误差变成零均值噪声；而选对空间，则保证无偏性落在真正决定更新幅度的那个量上。第二个方案 ZE-EDEN 走互补路线：它使用不含零的二阶矩码本，因此量化值存在一个正的下限，避免了接近零时倒数发散。代价是这个正下限本身会扭曲预条件子，于是论文用 EDEN 式校准，对量化后的二阶矩块重新缩放，来抵消这种畸变。两种配置的一阶矩都使用 4 比特 NormalFloat，即 NF4，并且在训练的最后 10% 阶段，对语言模型输出头的一阶矩施加有针对性的随机舍入。

实验覆盖了 GPT 与 Llama 风格的预训练，模型规模从 1.3 亿到 27 亿参数。结果是，两种方法在每一个被评估的规模上，都缩小了 TorchAO 4 比特 AdamW 相对 32 位 AdamW 的平均验证损失差距，最大的一次缩小达到 70%。在全参数监督微调中，两种配方的验证损失都低于 TorchAO，同时在下游任务上与 32 位 AdamW 保持接近。值得注意的是，“每个规模都改善”比单点的大幅领先更有说服力，因为它说明收益来自机制，而不是某次运行的运气。 把这些机制放回工程语境，可以看到几层含义。第一，4 比特状态量化的难点从来不只是比特数少，而是误差如何被后续计算放大。以往的做法默认把状态当作要逼近的对象，用均方误差衡量好坏；这篇论文提醒我们，真正被消费的是由状态派生出来的预条件子，评价与舍入都应当围绕它展开。第二，随机舍入与确定性舍入的取舍在这里有了新的依据：在状态空间里无偏，并不等于在预条件子空间里无偏，所以同样是随机舍入，换一个坐标就是另一种估计器。第三，两条路线给出了一个有趣的对照：一条保留零并改变概率的计算方式，另一条排除零并靠重新缩放修正下限。它们殊途同归，都在处理同一个病灶，即零附近那个量化单元里的预条件子失真。对工程团队而言，这意味着可以按自己的码本约束和内核实现选择更容易落地的一条。此外，对输出头一阶矩在训练末段做针对性随机舍入，也提示了一个经验：并非所有张量对量化同样敏感，把有限的精度预算花在最敏感的位置，往往比整体均匀降低误差更划算。这也解释了为何论文不满足于单一配方，而是同时给出两条可互相印证的路线，让读者能够判断收益究竟来自哪一环。

我们的判断是，这篇论文最大的价值在于把优化器状态量化的设计空间重新划分了一次。过去的讨论多集中在码本怎么设计、块怎么切；这里提出的是第三根轴，即舍入发生在哪个坐标。这个视角很可能迁移到其他对状态做非线性变换的优化器。同时也要保持克制：摘要没有给出显存节省的具体数字和训练吞吐的变化，评估规模止于 27 亿参数，更大模型与更长训练是否维持同样的收益，还需要独立复现来回答。对正在为训练显存发愁的团队，合理的做法是把 ZIP-SR 与 ZE-EDEN 当作值得在自己的规模上对照 TorchAO 做小范围验证的候选，而不是直接替换生产配置。