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

Paper: 2610.12444
Authors: Hanyang Li, Shao Tang, Daniel Thomas Braithwaite, Gregory Dexter, Leonardo Neves
Categories: cs.LG

The Gap

In modern large language model pre-training, optimizer states are the single largest consumer of persistent accelerator memory: standard 32-bit AdamW requires 8 bytes per parameter (4 bytes for first moment $m_t$, 4 bytes for second moment $v_t$), dwarfing 16-bit model weights. While 8-bit AdamW is widely accepted, pushing optimizer states down to 4-bit has notoriously caused measurable validation loss degradation and training instability.

Existing 4-bit quantization solutions (such as TorchAO and bitsandbytes) treat optimizer state compression as a generic float quantization problem in state space: they construct codebooks to minimize the reconstruction error $|v_t - \hat{v}_t|$. However, this formulation ignores the core mathematical mechanism of AdamW: parameter updates are scaled by the preconditioner:

$P(v_t) = \frac{1}{\sqrt{v_t} + \epsilon}$

Because the derivative of $P(v)$ approaches infinity as $v \to 0$, even a minuscule state reconstruction error near zero causes an exponential distortion in the effective learning rate. Traditional round-to-nearest in state space creates severe step-size bias, degrading convergence.

   THE ASYMMETRY OF STATE-SPACE VS PRECONDITIONER-SPACE ROUNDING

   Second Moment Value: v_t close to 0 (e.g., 1e-5)
            |
            v
   Quantizer reconstructs: hat{v}_t = 4e-5  (State Error = 3e-5: "Looks small!")
            |
            v
   Calculated Preconditioner:
   True P(v_t)        = 1 / sqrt(1e-5) = 316.2
   Quantized P(hat{v}) = 1 / sqrt(4e-5) = 158.1
   ===> EFFECTIVE STEP SIZE CUT IN HALF (-50% ERROR)!
            |
            v
   SOLUTION: ROUND DIRECTLY IN PRECONDITIONER SPACE
   - Compute rounding thresholds and stochastic probabilities on P(v)
   - Eliminates near-zero distortion and restores 32-bit convergence dynamics

The Increment

One sentence: By redesigning 4-bit optimizer quantization to calculate rounding thresholds and stochastic probabilities directly in preconditioner space rather than state space, ZIP-SR and ZE-EDEN close up to 70% of the performance gap between 4-bit and 32-bit AdamW across diverse pre-training scales.

Core Mechanism

The authors formulate optimizer quantization through the lens of coordinate choice for rounding decisions:

1. ZIP-SR (Zero-Inclusive Preconditioner-space Stochastic Rounding):
   • Retains exact zero in the second-moment codebook to prevent artificial gradient floors.
   • For an unquantized value $v \in [q_i, q_{i+1}]$, standard stochastic rounding uses state distance $p = (v - q_i) / (q_{i+1} - q_i)$.
   • ZIP-SR instead calculates transition probabilities in the transformed preconditioner coordinate: $p_{\text{ZIP-SR}} = \frac{P(v) - P(q_i)}{P(q_{i+1}) - P(q_i)}$
   • This guarantees that the expected preconditioner update matches the full-precision operator: $\mathbb{E}[P(\hat{v})] = P(v)$.
2. ZE-EDEN (Zero-Excluding EDEN Calibration):
   • Employs a zero-excluding second-moment codebook with dynamic block-level scaling.
   • Calibrates the positive quantization floor to neutralize preconditioner distortion without blowing up numerical precision.
3. Targeted First-Moment Stabilization:
   • First moments ($m_t$) are compressed using 4-bit NormalFloat (NF4).
   • In the final 10% of the training cosine schedule, targeted stochastic rounding is enabled specifically for the language modeling head (LM-head), preventing gradient stagnation in high-entropy final token logits.

   ZIP-SR QUANTIZATION FLOW IN PRECONDITIONER COORDINATES

   Unquantized Second Moment: v_t
                 |
                 v
   Map to Preconditioner Space: P(v_t) = 1 / (sqrt(v_t) + eps)
                 |
                 v
   Locate Adjacent Quantization Points: [P(q_i), P(q_{i+1})]
                 |
                 v
   Compute Stochastic Probability:
   p = (P(v_t) - P(q_i)) / (P(q_{i+1}) - P(q_i))
                 |
                 +-----------------------+
                 |                       |
                 v (prob: 1 - p)         v (prob: p)
           Select Level q_i        Select Level q_{i+1}
                 |                       |
                 +-----------+-----------+
                             |
                             v
               Store 4-bit Quantized Index
            (Expected Preconditioner Unbiased!)

Key Concepts

• Rounding Space: The mathematical domain (state space vs function space) in which rounding decisions or stochastic probabilities are evaluated during quantization.
• Preconditioner Distortion: The asymmetric magnification of error that occurs when inverting square-root second-moment estimates near the origin.
• Unbiased Preconditioner Expectation: The property where the expected value of the quantized preconditioner over random stochastic rounding seeds equals the exact unquantized preconditioner.

Framework Shift

Before (State-Space 4-bit AdamW - TorchAO / bitsandbytes):
  Quantization Objective: Minimize |v - hat{v}|.
  Rounding Metric: Linear distance in moment value.
  Failure Mode: Extreme learning rate suppression near zero gradients; validation loss gap.

After (Preconditioner-Space 4-bit AdamW - ZIP-SR / ZE-EDEN):
  Quantization Objective: Minimize |P(v) - P(hat{v})|.
  Rounding Metric: Distance evaluated along P(v) = 1 / (sqrt(v) + eps).
  Performance: Closes up to 70% of the 32-bit gap; identical memory savings.

Expert Assessment

Problem choice: High-impact systems optimization. Optimizer memory is the primary reason engineers resort to ZeRO-3 offloading or activation recomputation; practical 4-bit AdamW saves gigabytes of VRAM per GPU.

Method maturity: Grounded in rigorous first-principles analysis of the non-linear curvature of $1/\sqrt{x}$, providing both an elegant stochastic rounding recipe and a calibrated deterministic alternative.

Experimental integrity: Tested extensively on GPT and Llama pre-training from 130M up to 2.7B parameters, as well as full-parameter SFT tasks, showing consistent superiority over TorchAO across all benchmarks.

Writing quality: Mathematically precise, clearly isolating the analytical source of quantization degradation.

Verdict: strong accept — A foundational systems paper that turns 4-bit AdamW from an experimental compromise into an industrial-ready training baseline.

Takeaways

• Never round optimizer states in raw magnitude space; always round in the coordinate space of the downstream functional operator.
• Use ZIP-SR to retain zero in second-moment codebooks and make expected preconditioner updates mathematically unbiased.
• Apply targeted stochastic rounding to the LM-head during the final 10% of cosine decay to safeguard final model convergence.