4-bit AdamW优化器量化新思路：在预条件子空间做舍入

为什么4-bit优化器状态量化值得重新审视

AdamW 的优化器状态——一阶矩和二阶矩——通常以 32-bit 浮点形式常驻显存，是训练大模型时不可忽视的持久存储开销。把它压缩到 4-bit 看起来是一条直接的降本路径，但量化误差会沿矩的递推关系逐步传播，扰动后续的自适应更新。

真正的问题不在于「能不能压到 4-bit」，而在于在哪个坐标里决定舍入。已有做法通常在状态空间（state space）中量化：直接对一阶矩、二阶矩的数值做舍入，再用还原后的状态去计算更新量。这篇工作提出一个被忽略的视角——舍入空间（rounding space），即量化器在相邻重建电平之间做选择时所处的坐标系。对二阶矩而言，选择在状态空间还是预条件子空间（preconditioner space）中舍入，会带来定性不同的优化动力学。

AdamW 与优化器状态：AdamW 是 Adam 优化器的解耦权重衰减变体，为每个可训练参数维护两份指数移动平均——一阶矩（梯度的均值估计，类似动量）与二阶矩（梯度平方的均值估计，用于自适应缩放步长）。对于拥有数十亿参数的模型，这两份状态与参数本身同量级，通常以 fp32 存储，因此在混合精度训练中，优化器状态往往比模型权重和激活更占显存。将状态压到 4-bit 意味着每个数值仅有 16 个可能的取值（重建电平），必须借助码本（codebook）在浮点范围内非均匀地分配这些电平。量化器在每个数值上做的核心决策是：在相邻两个重建电平之间选哪一个，这个选择过程就是「舍入」。理解舍入发生在哪个坐标系，正是本文的出发点。

预条件子（preconditioner）：在 AdamW 中，参数更新量约等于「一阶矩 /（二阶矩的平方根 + ε）」。这里二阶矩平方根的倒数扮演了预条件子的角色——它为每个坐标自适应地缩放梯度，使得梯度变化剧烈的方向步长变小、平缓的方向步长变大。由于该映射包含平方根与取倒数两步非线性，当二阶矩接近零时，其微小扰动会被显著放大。这意味着在二阶矩的数值空间中「保真」并不等于在预条件子空间中「保真」，两者的误差度量并不等价。

状态空间舍入的盲点

论文对二阶矩中靠近零的量化单元做了局部分析，得到一个反直觉的结论：小的平均状态误差，并不意味着下一步有小的平均预条件子误差。

AdamW 的更新依赖二阶矩的平方根倒数，也就是预条件子。靠近零的区域恰恰是这一映射最敏感的地方——分母趋近于零时，状态的微小偏差会被放大成预条件子的显著偏差。如果量化器在状态空间里追求「重建值尽量接近原值」，它优化的目标与真正影响更新的量之间存在错位。

作者进一步构造了一个一维二次型问题，验证在状态空间舍入与预条件子空间舍入下，优化轨迹会出现定性差异。这为「换个坐标做舍入」提供了超出经验层面的动机。

随机舍入（stochastic rounding）：与「四舍五入」这种确定性舍入不同，随机舍入以与小数部分成正比的概率在两个相邻重建电平之间选择，从而在期望意义上无偏。在低比特量化中它被广泛用于缓解累积偏差：即使单个数值的舍入误差不可忽略，大量此类误差在期望上会相互抵消，使得矩的递推不会被系统性地拉偏。本文的关键改动在于，把这一概率的计算从状态空间迁移到预条件子空间，让「无偏性」作用于真正影响更新的量，而非中间状态。

两条互补的技术路线

基于上述分析，论文给出两种方案。

ZIP-SR（Zero-Inclusive Preconditioner-space Stochastic Rounding，零包含预条件子空间随机舍入）：在二阶矩码本中保留零，并把随机舍入的概率计算放到预条件子空间里完成。保留零这一条很关键——它避免了正数量化下限带来的系统性偏移。

ZE-EDEN（Zero-Excluding EDEN calibration，零排除 EDEN 校准）：走相反方向，使用排除零的二阶矩码本，再对量化后的二阶矩块做重缩放，以缓解正量化下限造成的预条件子失真。

两条路线的共同点是：都用 4-bit NormalFloat（NF4）处理一阶矩，并在训练最后 10% 阶段对 LM-head 的一阶矩施加有针对性的随机舍入。

实验结果

在 GPT 风格与 Llama 风格的预训练任务上，作者覆盖了从 130M 到 2.7B 参数的多个模型规模。与 TorchAO 的 4-bit AdamW 相比，ZIP-SR 与 ZE-EDEN 在每一个评测规模上都缩小了相对 32-bit AdamW 的平均验证损失差距，最大差距缩减幅度达到 70%。

在全参数监督微调场景中，两种方法都取得了比 TorchAO 更低的验证损失，同时在下游任务上保持接近 32-bit AdamW 的表现。这一点值得留意：量化方案的收益不只在预训练阶段，微调同样受益。

几点观察

这项工作最有价值的部分是把「量化误差」重新表述为「在哪个空间度量误差」。当被量化的量要经过非线性变换（如平方根倒数）才进入更新时，状态空间的保真度不是正确的优化目标。类似的思路或许能迁移到其他带非线性耦合的优化器状态上。

另一个信号是零的处理方式。两种方案在「码本是否包含零」上给出了相反选择，却都能改善结果——说明关键不在零本身，而在于是否对预条件子失真做了显式补偿。论文未在摘要中给出两种方案之间的直接对比结论，实际选择可能取决于具体模型规模与训练阶段。