A New Approach to 4-bit AdamW Optimizer Quantization: Rounding in Preconditioner Space

Why 4-bit Optimizer-State Quantization Deserves a Second Look

AdamW's optimizer states—the first moment and second moment—typically reside in GPU memory in 32-bit floating point, and they represent a persistent storage cost that cannot be ignored when training large models. Compressing them to 4-bit looks like a straightforward cost-reduction path, but quantization error propagates step by step along the moment recurrence, perturbing subsequent adaptive updates.

The real question isn't "can we compress to 4-bit," but rather in which coordinate system we decide to round. Existing approaches typically quantize in state space: they directly round the values of the first and second moments, then use the restored states to compute the update. This work proposes an overlooked perspective—the rounding space, meaning the coordinate system in which the quantizer chooses between adjacent reconstruction levels. For the second moment, choosing whether to round in state space or in preconditioner space leads to qualitatively different optimization dynamics.

AdamW and optimizer states: AdamW is the decoupled weight decay variant of the Adam optimizer, maintaining two exponential moving averages for each trainable parameter—the first moment (a mean estimate of gradients, akin to momentum) and the second moment (a mean estimate of squared gradients, used to adaptively scale the step size). For models with billions of parameters, these two states are of the same order of magnitude as the parameters themselves, usually stored in fp32, so in mixed-precision training the optimizer states often consume more GPU memory than the model weights and activations. Compressing the states to 4-bit means each value has only 16 possible values (reconstruction levels), which requires a codebook to distribute these levels non-uniformly across the floating-point range. The core decision the quantizer makes for each value is: which of two adjacent reconstruction levels to choose. This selection process is "rounding." Understanding in which coordinate system rounding occurs is precisely the starting point of this paper.

Preconditioner: In AdamW, the parameter update is approximately "first moment / (square root of second moment + ε)." Here the reciprocal of the square root of the second moment plays the role of a preconditioner—it adaptively scales the gradient for each coordinate, so that directions where the gradient changes sharply get smaller step sizes and smooth directions get larger step sizes. Because this mapping involves two nonlinear steps (square root and reciprocal), when the second moment is near zero, tiny perturbations get significantly amplified. This means that "fidelity" in the numerical space of the second moment is not equivalent to "fidelity" in preconditioner space; the two error metrics are not the same.

The Blind Spot of State-Space Rounding

The paper performs a local analysis of quantized units of the second moment near zero and reaches a counterintuitive conclusion: a small average state error does not imply a small average preconditioner error at the next step.

AdamW's update relies on the reciprocal square root of the second moment, i.e., the preconditioner. The region near zero is exactly where this mapping is most sensitive—as the denominator approaches zero, a tiny deviation in the state gets amplified into a significant deviation in the preconditioner. If the quantizer pursues "reconstructed values as close as possible to the original" in state space, there is a mismatch between the objective it optimizes and the quantity that actually affects the update.

The authors further construct a one-dimensional quadratic problem to verify that optimization trajectories diverge qualitatively between state-space rounding and preconditioner-space rounding. This provides motivation beyond the empirical level for "rounding in a different coordinate system."

Stochastic rounding: Unlike deterministic rounding such as "round half to even," stochastic rounding chooses between two adjacent reconstruction levels with probability proportional to the fractional part, thereby being unbiased in expectation. It is widely used in low-bit quantization to mitigate accumulated bias: even if the rounding error of a single value is non-negligible, many such errors cancel out in expectation, preventing the moment recurrence from being systematically pulled off course. The key change in this paper is to move the computation of this probability from state space to preconditioner space, so that "unbiasedness" acts on the quantity that actually affects the update rather than on the intermediate state.

Two Complementary Technical Approaches

Based on the above analysis, the paper proposes two schemes.

ZIP-SR (Zero-Inclusive Preconditioner-space Stochastic Rounding): Keeps zero in the second-moment codebook and computes the stochastic rounding probability in preconditioner space. Keeping zero is crucial—it avoids the systematic offset caused by a positive quantization lower bound.

ZE-EDEN (Zero-Excluding EDEN calibration): Takes the opposite direction, using a second-moment codebook that excludes zero, then rescaling the quantized second-moment blocks to mitigate preconditioner distortion caused by a positive quantization lower bound.

What both approaches share: they both use 4-bit NormalFloat (NF4) for the first moment, and they apply targeted stochastic rounding to the first moment of the LM-head during the last 10% of training.

Experimental Results

On GPT-style and Llama-style pretraining tasks, the authors cover multiple model scales from 130M to 2.7B parameters. Compared with TorchAO's 4-bit AdamW, ZIP-SR and ZE-EDEN reduce the average validation loss gap relative to 32-bit AdamW at every evaluated scale, with a maximum gap reduction of 70%.

In full-parameter supervised fine-tuning scenarios, both methods achieve lower validation loss than TorchAO while maintaining performance close to 32-bit AdamW on downstream tasks. This is worth noting: the benefits of the quantization scheme are not limited to pretraining—fine-tuning benefits as well.

A Few Observations

The most valuable part of this work is that it reframes "quantization error" as "in which space the error is measured." When the quantity being quantized must pass through a nonlinear transformation (such as reciprocal square root) before entering the update, fidelity in state space is not the correct optimization objective. Similar ideas might transfer to other optimizer states with nonlinear coupling.

Another signal is how zero is handled. The two schemes make opposite choices regarding "whether the codebook includes zero," yet both improve results—indicating that the key lies not in zero itself, but in whether explicit compensation is made for preconditioner distortion. The paper does not provide a direct comparison between the two schemes in the abstract; the actual choice may depend on the specific model scale and training stage.