Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization

The gist Quantizing AdamW’s optimizer states reduces persistent storage, but quantization errors propagate through the moment recurrences and perturb subsequent adaptive updates<ref:2610.12444#pg2>.

Redesigning Rounding Space

The paper proposes redesigning 4-bit optimizer-state quantization by shifting the perspective to rounding space, where the quantizer chooses between adjacent reconstruction levels<ref:2610.12444#pg3>. This approach is motivated because conventional rounding in state space can still distort the resulting update even if it is accurate or unbiased in state space<ref:2610.12444#pg4>. The key insight is that the rounding rule is critical to whether zero-point failure occurs with a zero-inclusive codebook <ref:2610.12444#pg4>.

ZIP-SR and ZE-EDEN Configurations

The authors develop two specific 4-bit AdamW configurations: ZeroInclusive Preconditioner-space Stochastic Rounding (ZIP-SR) and ZeroExcluding EDEN calibration (ZE-EDEN)<ref:2610.12444#pg3>. ZIP-SR combines a zero-inclusive second-moment codebook with Update-SR, while ZE-EDEN combines a zero-excluding codebook with EDEN block calibration<ref:2610.12444#pg4>. Both methods use 4-bit NormalFloat (NF4) for the first moment and targeted stochastic rounding of the LM-head first moment during the final 10% of training<ref:2610.12444#pg3>. The goal is to reduce TorchAO 4-bit AdamW’s mean validation-loss gap to 32-bit AdamW at every evaluated model size, with the largest reported gap reduction reaching 70% in pretraining experiments<ref:2610.12444#pg3>.

ZIP-SR Dynamics

ZIP-SR is characterized by Update-SR, which assigns stochastic-rounding (SR) probabilities using preconditioners induced by the candidate reconstruction levels<ref:2610.12444#pg4>. Proposition 1(a) shows that for adjacent levels 0 ≤ a < b in a nonnegative decoded codebook Cv, State-SR preserves the mean state value, whereas Update-SR preserves the mean of the current step preconditioner<ref:2610.12444#pg4>. The analysis demonstrates that Update-SR strictly reduces the magnitude of the bias while keeping the conditional mean preconditioner at or below its unquantized value <ref:2610.12444#pg4>.

ZE-EDEN Dynamics

ZE-EDEN uses a zero-excluding second-moment codebook and EDEN block calibration to mitigate the distortion caused by the positive quantization floor<ref:2610.12444#pg4>. The guarded EDEN-inspired calibration uses a factor cτ which lowers the decoded block scale and its positive floor when it overestimates small second-moment coordinates, potentially restoring part of the tail suppressed by uncalibrated quantization<ref:2610.12444#pg4>. Figure 2 shows that EDEN partially restores the large-preconditioner tail in both the embedding and LM-head states compared to Dyn4-NZ<ref:2610.12444#pg7>.

First Moment Quantization

The first moment is quantized with NF4, using RTN by default and switching the LM-head first moment to SR during the final 10% of training<ref:2610.12444#pg6>. In GPT-style 2.7B pretraining, NF4-RTN produces a predominantly negative radial bias in the LM-head first moment, whereas NF4-SR remains centered near zero (Figure 4b)<ref:2610.12444#pg8>. This suggests that the failure is associated with deterministic LM-head rounding rather than 4-bit LM-head quantization itself <ref:2610.12444#pg8>.

Experimental Results

Across pretraining and supervised fine-tuning (SFT) experiments, both ZIP-SR and ZE-EDEN achieve lower validation loss than TorchAO while remaining close to 32-bit AdamW on downstream tasks<ref:2610.12444#pg3>. Table 3 shows that the mean gap reduction for step 3b (ZIP-SR) is a larger reduction of 32.5% compared to step 3a (ZE-EDEN) when comparing against the parent configuration<ref:2610.12444#pg9>. The final results show that both proposed optimizers attain lower mean validation loss than TorchAO 4-bit AdamW in SFT settings<ref:2610.12444#pg10>.

Conclusion

This work shows that optimizer-state quantization should account for how stored values affect optimization, rather than state-space reconstruction alone<ref:2610.12444#pg10>. For the second-moment, this perspective motivates Update-SR for the zero-inclusive route, while EDEN calibration separately improves the zero-excluding route by mitigating distortion in the induced preconditioner distribution<ref:2610.12444#pg10>. Together with NF4 and SR for the LM-head first-moment, these designs improve over TorchAO across pretraining and fine-tuning settings<ref:2610.12444#pg10>. The final validation loss differences in SFT show that ZE-EDEN attains the lower validation loss on both SFT models, whereas ZIP-SR generally achieves smaller gaps in the pretraining comparisons<ref:2610.12444#pg10>.

Limitations and Future Work

The theory compares conditional one-step preconditioner biases and optimization dynamics on a toy model, but does not establish general end-to-end convergence for quantized AdamW; extending these results to broader settings remains future work<ref:2610.12444#pg11>. Targeted LMhead SR alleviates the observed late-training regression, but the cause of the directional first-moment bias and its possible relationship to learning-rate decay remain unresolved<ref:2610.12444#pg11>. Further directions include extending our methodology below 4-bit optimizer-state quantization and investigating whether it applies to other adaptive optimizers<ref:2610.12444#pg11>.

AI Use Statement

Building on the authors’ initial ideas to compare the optimizer state and preconditioner errors of different rounding rules, we used ChatGPT (GPT-6 and GPT-5.6 Sol) to help formulate the statements of Propositions 1 and 2 for theoretical justification and polish the proof in multiple rounds of discussion<ref:2610.12444#pg11>. We used ChatGPT and Claude to give feedback on the experimental design and to polish the writing and organization<ref:2610.12444#pg11>. The authors verified all AI-assisted proofs and text and take full responsibility for the manuscript<ref:2610.12444#pg11>.

References

The paper cites several works including Kingma & Ba (2015) on Adam, Dettmers et al. (2022) on 8-bit optimizers, and various studies on dynamics of quantized optimizer states<ref:2610.12444#pg4>.

Quantization Background

The paper defines several 4-bit formats including FP4, Dyn4-NZ, and NF4, noting that NF4 is constructed from quantiles of a standard normal distribution with an exact zero included<ref:2610.12444#pg16>. The zero-excluding linear mapping for the second moment is denoted as TLin4-NZ<ref:2610.12444#pg16>. The rounding rules are defined, including Round-to-Nearest (RTN) and Stochastic rounding (SR)<ref:2610.12444#pg16>.

Experimental Details

All pretraining and supervised fine-tuning (SFT) runs use AdamW with specific hyperparameters, including BF16 mixed-precision computation with FP32 master parameters<ref:2610.12444#pg18>. The eligibility criterion for moment tensors is identical across the 4-bit methods, and ZE-EDEN uses the denominator safeguard τ = 10−12 and stores the effective scale cτ s by absorbing its calibration factor into the block scale<ref:2610.12444#pg18>. The final validation loss differences in SFT show that ZE-EDEN attains the lower validation loss on both SFT models, whereas ZIP-SR generally achieves smaller gaps in the pretraining comparisons<ref:2610.12444#pg10>.

Model Architectures

The paper uses GPT-style models at 162M, 405M, 834M, and 2.

1. Bold header: Zero-Inclusive Preconditioner-space Stochastic Rounding (ZIP-SR) for Optimizer States

This configuration retains zero in the second-moment codebook and computes stochastic-rounding probabilities in preconditioner space, which is shown to be a superior route compared to the zero-excluding approach when considering optimization dynamics. This leads to a reported reduction of 43.2% lower validation loss on GPT-style 834M pretraining compared to TorchAO.

2. Bold header: Zero-Excluding EDEN Calibration (ZE-EDEN) for Robustness

This method utilizes a zero-excluding second-moment codebook and rescales the quantized second-moment block to mitigate the preconditioner distortion caused by the positive quantization floor. This configuration is effective at reducing the mean validation-loss gap to 32-bit AdamW across evaluated model sizes.

3. Bold header: Targeted First Moment Stochastic Rounding for Stability

The use of targeted stochastic rounding of the LM-head first moment during the final 10% of training is motivated by addressing instability, as NF4-RTN produces a predominantly negative radial bias in the LM-head first moment, which is a key factor in observed late-training regression.

4. Bold header: Dynamic Codebook Selection Based on First Moment Distribution

The paper demonstrates that the choice of first-moment codebook—where FP4 and SDyn4 allocate several levels to narrow cells near zero, each receiving little probability mass, whereas NF4 spreads its levels across the high-density region—significantly impacts reconstruction error. This suggests dynamically selecting a codebook based on empirical moment distributions can improve quantization quality.

5. Bold header: Rounding Space Separation for Training Stability

The results show that changing the rounding space avoids the observed instability without removing zero; state-space stochastic rounding alone does not resolve it. Specifically, Update-SR remains locally contractive and contrasts with State-RTN which can develop a stable nonzero two-cycle on scalar quadratic objectives.

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 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 (ZIP-SR), which retains zero in the second-moment codebook and computes stochastic-rounding probabilities in preconditioner space. As a complementary route, Zero-Excluding EDEN calibration (ZE-EDEN) uses a zero-excluding second-moment codebook and rescales the quantized second-moment block to mitigate the preconditioner distortion caused by the positive quantization floor. Both configurations use 4-bit NormalFloat (NF4) for the first moment, with targeted stochastic rounding of the LM-head first moment during the final 10% of training. Across GPT- and Llama-style pretraining experiments ranging from 130M to 2.7B parameters, both methods reduce TorchAO 4-bit AdamW's mean validation-loss gap to 32-bit AdamW at every evaluated model size, with the largest reported gap reduction reaching 70%. In full-parameter supervised fine-tuning, both recipes achieve lower validation loss than TorchAO while remaining close to 32-bit AdamW on downstream tasks.