Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization
summary
The gist
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>.
In short
The paper addresses errors in quantizing AdamW optimizer states by focusing on rounding rules rather than just state space reconstruction. It proposes two 4-bit configurations, ZIP-SR and ZE-EDEN, which use specific stochastic rounding and calibration techniques. These methods successfully reduce the validation loss gap compared to standard 4-bit AdamW across pretraining and fine-tuning experiments.
Key concepts
- Rounding Space Redesign
- Instead of just rounding the state values in a mathematical space, this approach focuses on how adjacent reconstruction levels are chosen. The key insight is that the specific rounding rule critically determines whether zero-point failures occur when using a codebook with zero-inclusive or zero-excluding properties.
- ZIP-SR Configuration
- This configuration uses ZeroInclusive Preconditioner-space Stochastic Rounding (ZIP-SR). It combines a second moment codebook that includes zero with Update-SR, which assigns stochastic rounding probabilities based on preconditioners derived from the candidate reconstruction levels. This helps preserve the mean state value while controlling bias.
- ZE-EDEN Configuration
- This configuration employs ZeroExcluding EDEN calibration to handle distortion caused by a positive quantization floor in the second moment. It uses a guarded, EDEN-inspired calibration factor (cτ) to lower the decoded block scale and its positive floor when it overestimates small second-moment coordinates, helping restore suppressed tail information.
Terminology used across episodes
This episode discusses
- Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization · Paper Radio
- Gefen: Optimized Stochastic Optimizer · Paper Radio
- Evaluating Large Language Models Trained on Code
- Training Verifiers to Solve Math Word Problems
- Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention
- FlashOptim: Optimizers for Memory-Efficient Training
- Effective Quantization of Muon Optimizer States
- STQuant: Spatio-Temporal Adaptive Framework for Optimizer Quantization in Large Multimodal Model Training
- TorchAO: PyTorch-Native Training-to-Serving Model Optimization
- Understanding Quantization of Optimizer States in LLM Pre-training: Dynamics of State Staleness and Effectiveness of State Resets · Paper Radio
- Llama 2: Open Foundation and Fine-Tuned Chat Models
- Pushing the Limits of Low-Bit Optimizers: A Focus on EMA Dynamics
- Qwen3 Technical Report
- Instruction-Following Evaluation for Large Language Models
The paper
Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization · Read on arXiv
Hanyang Li, Shao Tang, Daniel Thomas Braithwaite, Gregory Dexter, Leonardo Neves, Aman Gupta, Hiroto Udagawa, Abhishek Shivanna, Daniel Silva
University of California, Berkeley
Transcript
Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.
Tom: Today's paper: "Rounding in Preconditioner Space".
Jane: 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>.
Tom: First, who's behind it and why it matters.
Paper summary: Tom: We started by looking at how quantizing AdamW’s optimizer states reduces persistent storage, but we found that errors propagate through the moment recurrences and those errors perturb subsequent adaptive updates >Page one--#pg1 <ref:2610.12444#pg1,quantizing AdamW’s optimizer states reduces persistent storage, but>.
Jane: The authors of "Rounding in Preconditioner Space: Redesigning four-bit AdamW Optimizer-State Quantization" propose a redesign focusing on rounding space, where the quantizer chooses between adjacent reconstruction levels >Page one--#pg1 <ref:2610.12444#pg1,Rounding in Preconditioner Space: Redesigning 4-bit AdamW Optimizer-State Quantization>.
Lu: They argue that conventional rounding in state space can still distort the resulting update even if it is accurate or unbiased there >Page four--#pg4 <ref:2610.12444#pg1>.
Meng: The key insight they present is that the rounding rule itself is critical to whether zero-point failure occurs with a zero-inclusive codebook >Page four--#pg4 <ref:2610.12444#pg1>.
Tom: They make two main contributions: first, characterizing how the rounding coordinate affects near-zero second-moment quantization at fixed reconstruction levels >Page two--#pg2 <ref:2610.12444#pg1>.
Jane: A local zero-cell analysis and a scalar quadratic construction show that this affects both preconditioner error and optimization dynamics >Page four--#pg4 <ref:2610.12444#pg1>.
Lu: This motivates Update-SR, which assigns stochastic-rounding probabilities using the preconditioners induced by the candidate reconstruction levels >Page two--#pg2 <ref:2610.12444#pg1>.
Meng: They also developed two specific four-bit AdamW configurations: ZIP-SR and ZE-EDEN >Page two--#pg2 <ref:2610.12444#pg1>.
Tom: ZIP-SR combines a zero-inclusive second moment codebook with Update-SR, while ZE-EDEN combines a zero-excluding codebook with EDEN block calibration >Page four--#pg4 <ref:2610.12444#pg1>.
Jane: Both methods use four-bit NormalFloat NF4 for the first moment and targeted stochastic rounding of the LM-head first moment during the final ten percent of training >Page three--#pg3.
Conclusion: Tom: So, looking at "Rounding in Preconditioner Space: Redesigning four-bit AdamW Optimizer-State Quantization," we see that optimizer state quantization should account for how stored values affect optimization, not just reconstruction error in state space >Page ten--#pg10.
Jane: For the second moment, this motivates Update-SR for the zero-inclusive route and EDEN calibration separately improves the zero-excluding route by mitigating distortion in the induced preconditioner distribution >Page ten--#pg10 <ref:2610.12444#pg3>.
Lu: Together with NF4 and SR for the LM-head first moment, these designs improve over TorchAO across pretraining and fine-tuning settings >Page ten--#pg10 <ref:2610.12444#pg3>.
Meng: The final validation loss differences in SFT show that ZE-EDEN attains the lower validation loss on both SFT models while ZIP-SR generally achieves smaller gaps in the pretraining comparisons >Page ten--#pg10 <ref:2610.12444#pg3>.
Tom: So, for anyone listening, this paper shows we need to look at how stored values affect optimization instead of just state-space reconstruction alone >Page ten--#pg10 <ref:2610.12444#pg3>.
More episodes
- 2610.09563-EvoSignal: LLM-Guided Evolutionary Design of Modular Traffic Signal Control Programs
- 2610.10135-Attention via Black-Box Vector Search
- 2610.09471-When Should an In-Context Learner Expand Its Hypothesis Space?
- 2610.12358-Subspace Uncertainty and Sharp Sampling Thresholds on the Boolean Cube
- 2610.12321-Spatial Pattern Formation from Multi-Agent Learning in Public Goods Dilemmas
- 2610.12340-Ambient Discrete Diffusion: Using the Wrong Data at the Right Time for Data Efficient Learning
- 2610.12349-SplitJEPA: Learning Invariant and Variant Latent Worlds without Reconstruction
- 2610.12379-Marformer: A Transformer for Predicting Missing Data Distributions
- 2610.12401-Learning Kilometer-Scale Weather Prediction with Global-Regional Alignment
- 2610.12397-Prospective Prediction of OOD Degradation from Source-Side Training Dynamics