A new arXiv preprint reports two 4-bit AdamW methods that narrowed the validation-loss gap to matched 32-bit runs across seven model sizes. The largest validation-loss gap reduction was 70.1% on a 1.4-billion-parameter GPT-style model, compared with TorchAO’s 4-bit optimizer. The recipes compress AdamW’s two running buffers between training steps; the experiments use BF16 computation with FP32 master weights.
AdamW tracks two moving averages for each parameter: recent gradients and their squares. It divides the gradient average by the square root of the squared-gradient average to set each weight update. Quantization packs these buffers into fewer bits, but the stored averages feed into later updates. Near zero, storing the squared-gradient average as zero can make a later step too large; a codebook with a positive minimum can make some steps too small.
Li and colleagues at UC Berkeley and Nubank test two approaches. ZIP-SR keeps zero among the 4-bit values but chooses rounding probabilities in the factor that controls update size, the inverse square root of the squared-gradient average, rather than by distance between raw values. The paper’s one-step analysis says this makes tiny positive values less likely to be stored as zero, which can sharply enlarge an update. ZE-EDEN excludes zero and rescales each block to reduce errors from its positive minimum. Both recipes also use a 4-bit NormalFloat codebook for the gradient-average buffer and probability-based rounding for the model’s output-layer buffer during the final 10% of pretraining, so results reflect the combined recipes.
The pretraining comparisons covered GPT-style models from 162 million to 2.7 billion parameters and Llama-style models from 130 million to 1.1 billion, using three paired seeds per size. At 1.4 billion parameters, the mean validation-loss gap to matched 32-bit AdamW runs was 0.0521 with TorchAO and 0.0156 with ZIP-SR; the paper reports a 70.1% reduction in the gap. In full-parameter fine-tuning of Qwen3-8B and Llama-3.2-3B, both new methods had lower mean validation loss than TorchAO, while results on MMLU, GSM8K, HumanEval and IFEval showed no consistent advantage.
Based on their packed values and block scales, the authors calculate 1.0625 bytes per parameter for eligible optimizer-state buffers, compared with 8 bytes for two FP32 buffers. That is a 7.53-fold reduction in persistent moment storage, not a measurement of total training memory or speed. TorchAO’s 2.7-billion-parameter GPT run diverged late, so the paper omits a percentage comparison at that size. The authors also say they have not established general convergence guarantees for quantized AdamW.
Sources
- Hanyang Li et al., “Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization”, arXiv:2610.12444v1 (8 October 2026), especially §§3–5 and Appendix B.
- PyTorch implementation of ZIP-SR and ZE-EDEN by Nubank, Apache-2.0 license.
- Figure 1 from the arXiv preprint, licensed CC BY 4.0.