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

arXiv: https://arxiv.org/abs/2610.12444 · HTML: https://mlsanity.com/paper/2610.12444/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 13 of its 41 references that are in the map.
- First posted on arXiv 8 Oct 2026.

## Builds on (most cited first, 13 of 13)

- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Instruction-Following Evaluation for Large Language Models](https://mlsanity.com/paper/2311.07911/) (2023)
- [TorchAO: PyTorch-Native Training-to-Serving Model Optimization](https://mlsanity.com/paper/2507.16099/) (2025)
- [Effective Quantization of Muon Optimizer States](https://mlsanity.com/paper/2509.23106/) (2025)
- [Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention](https://mlsanity.com/paper/2607.04422/) (2026)
- [Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets](https://mlsanity.com/paper/2603.16731/) (2026)
- [Pushing the Limits of Low-Bit Optimizers: A Focus on EMA Dynamics](https://mlsanity.com/paper/2505.00347/) (2025)
- [FlashOptim: Optimizers for Memory-Efficient Training](https://mlsanity.com/paper/2602.23349/) (2026)
- [STQuant: Spatio-Temporal Adaptive Framework for Optimizer Quantization in Large Multimodal Model Training](https://mlsanity.com/paper/2604.06836/) (2026)
- [Gefen: Optimized Stochastic Optimizer](https://mlsanity.com/paper/2606.13894/) (2026)

## Abstract (arXiv)

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.