PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Hanyang Li, Shao Tang, Daniel Thomas Braithwaite, Gregory Dexter, Leonardo Neves, Aman Gupta, Hiroto Udagawa, Abhishek Shivanna, Daniel Silva, Rohan Ramanath

arXiv:2610.12444v1cs.LG

TL;DR

Quantizing AdamW’s optimizer states reduces storage, but state-space rounding can still distort adaptive updates because the second moment is transformed before scaling them. The paper presents ZIP-SR and ZE-EDEN, two 4-bit configurations with different second-moment strategies. Both improve on TorchAO across the evaluated pretraining and fine-tuning settings.

  • Problem

    State-space accuracy or unbiasedness can still distort AdamW’s updates because the second moment is transformed nonlinearly before it scales them.

    Source: §1 · p. 1

  • Method

    ZIP-SR uses a zero-inclusive second-moment codebook with preconditioner-space stochastic rounding, while ZE-EDEN uses a zero-excluding codebook with EDEN block calibration.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    Both configurations reduce TorchAO 4-bit AdamW’s mean validation-loss gap to 32-bit AdamW across evaluated pretraining settings and attain lower mean SFT validation loss.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The conclusion is that optimizer-state quantization should account for how stored values affect optimization, not only state-space reconstruction.

    Source: §6 · p. 9

    The theory compares conditional one-step preconditioner biases and toy-model dynamics, but does not establish general end-to-end convergence for quantized AdamW.

    Source: §6 · p. 10

Abstract

from arXiv · show

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 \emph{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 (\textbf{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 (\textbf{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 \textbf{130M} to \textbf{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 \textbf{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.

1 INTRODUCTION

AdamW’s FP32 moment buffers add substantial memory overhead, motivating optimizer-state compression; at 4 bits, the rounding coordinate matters because state-space accuracy can still distort updates. The paper develops two 4-bit configurations that improve on TorchAO across the evaluated pretraining settings and achieve lower mean SFT validation loss.

1 INTRODUCTION

  • FP32 first- and second-moment buffers require 8 bytes per parameter, motivating research on optimizer-state compression.Prior 4-bit work identified zero-point failure when a small positive second moment is mapped to zero.Source: §1 · p. 1
  • Preconditioner-space rounding preserves a zero-inclusive second-moment codebook while reducing harmful zero selections that state-space rounding can permit.AdamW transforms the second moment nonlinearly before it scales the update, so state-space accuracy or unbiasedness can still distort the resulting update.Source: §1 · pp. 1–2
  • Both ZIP-SR and ZE-EDEN reduce TorchAO 4-bit AdamW’s mean validation-loss gap to 32-bit AdamW across evaluated pretraining settings and attain lower mean SFT validation loss.Both use NF4 for the first moment and targeted LM-head stochastic rounding.Source: §1 · p. 2
  • ZIP-SR uses preconditioner-induced stochastic-rounding probabilities, while ZE-EDEN pairs a zero-excluding codebook with EDEN block calibration.The two configurations use different second-moment approaches while sharing NF4 for the first moment.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries define how quantized moment states enter AdamW’s recurrence and update, then situate this work among low-precision optimizer-state methods and other memory-efficient optimizers.

2.1 TERMINOLOGY AND NOTATION

  • Quantized states are stored between steps; the next working moments use previously decoded states, while the parameter update uses the current working moments before storage.The notation defines AdamW’s moment and update quantities, with elementwise squares, square roots, and divisions.Source: §2.1 · p. 2

2.2 RELATED WORK

  • Low-bit optimizer-state work spans 4-bit quantization, geometric and logarithmic encodings, precision allocation, and transformed or companded representations.The cited examples include COAT, FlashOptim, and Full-Stack FP4.Source: §2.2 · pp. 2–3
  • Other studies analyze recurrent quantization effects, including signal swamping, variance amplification, state staleness, persistent drift, and convergence guarantees under floating-point quantization.These topics are associated with SOLO (Xu et al., 2025), Topollai & Choromanska (2026), BReD, and Tang et al. (2026), respectively.Source: §2.2 · p. 3
  • Memory-efficient alternatives also reduce optimizer memory by structuring stored statistics, sharing adaptive rates, or maintaining optimizer information in lower-dimensional spaces.Examples include Adafactor, SM3, Adam-mini, GaLore, and APOLLO.Source: §2.2 · p. 3

3 SECOND-MOMENT QUANTIZATION

The section examines how second-moment rounding affects next-step preconditioners and optimization dynamics. It develops a zero-inclusive preconditioner-space approach and a zero-excluding calibration approach to address distinct near-zero distortions.

3 SECOND-MOMENT QUANTIZATION

  • Selecting zero near zero can produce an excessively large next-step preconditioner, whereas excluding zero imposes a positive floor that can suppress large preconditioners.ZIP-SR retains zero but computes stochastic-rounding probabilities in preconditioner space; ZE-EDEN instead rescales a zero-excluding block.Source: §3 · p. 3, §3.1 · p. 4

3.1 ZIP-SR: ZERO-INCLUSIVE PRECONDITIONER-SPACE STOCHASTIC ROUNDING

  • On a scalar quadratic, State-RTN approaches a nonzero two-cycle while Update-RTN converges geometrically; near the optimum, State-SR is locally expansive and Update-SR contractive.These are distinct trajectories under the zero-inclusive codebook.Source: §3.1 · pp. 4–5
  • With FP32 first moments and the same zero-inclusive Dyn4 second-moment codebook, State-RTN/SR develop large training-loss gaps while Update-RTN/SR avoid this deterioration.The controlled comparison attributes the observed stability difference to changing the rounding space without removing zero.Source: §3.1 · p. 5, Figure 1 · p. 5

3.2 ZE-EDEN: ZERO-EXCLUDING EDEN CALIBRATION

  • ZE-EDEN calibrates the decoded second-moment block to mitigate positive-floor distortion; lowering its scale can permit larger preconditioner entries.The calibration is guarded by τ = 10^-12 in the normalized inner product.Source: §3.2 · p. 5
  • 0.0049 to 0.0038: EDEN lowers Dyn4-NZ's final validation-loss gap with FP32 first moments while partially restoring the large-preconditioner tail.The tail restoration is reported for embedding and LM-head states.Source: §3.2 · p. 5, Figure 2 · p. 6
  • With NF4 or SDyn4 first moments, the uncalibrated Dyn4-NZ validation-loss gap rises to 0.0207 or 0.0312, motivating a first-moment study.These values are reported for the respective first-moment settings.Source: §3.2 · p. 6

4 FIRST-MOMENT QUANTIZATION

The section compares first-moment codebooks and rounding rules, finding NF4 performs best in the reported 834M comparison and that stochastic rounding of the LM-head moment addresses a late-training instability in GPT-style 2.7B pretraining.

4.1 CODEBOOK CHOICE

  • NF4 had the lowest mean squared reconstruction error at 1.056 × 10−3, versus 1.739 × 10−3 for FP4 and 2.495 × 10−3 for SDyn4.In the sampled distribution, FP4 and SDyn4 allocate several levels to low-mass cells near zero, while NF4 spreads levels across the high-density region.Source: §4.1 · p. 6
  • Under ZE-EDEN, final gaps were 0.0289 (SDyn4), 0.0272 (FP4), and 0.0183 (NF4); under ZIP-SR, they were 0.0239, 0.0214, and 0.0140, respectively.These are single-seed GPT-style 834M runs with RTN for the first moment throughout.Source: §4.1 · p. 6

4.2 ROUNDING RULE FOR THE FIRST MOMENT

  • In GPT-style 2.7B cooldown, NF4-RTN departed sharply from 32-bit AdamW, whereas NF4-SR on the LM-head first moment remained stable.The smaller GPT-style pretraining experiments showed no comparable late-stage degradation with NF4-RTN.Source: §4.2 · pp. 6–7
  • Matched checkpoint interventions localized the instability to LM-head NF4-RTN; replacing it with SR restored stability, while RTN produced predominantly negative radial bias.The bias under NF4-SR remained centered near zero; persistent inward RTN bias may accumulate through the first-moment recurrence and help explain subsequent loss instability.Source: §4.2 · p. 7
  • The final methods use NF4-RTN for the first moment by default, switching only the LM-head first moment to NF4-SR during the final 10% of pretraining.This policy follows the observed LM-head instability under RTN and stable trajectory under SR.Source: §4.2 · p. 7

5 EXPERIMENTS

Across evaluated pretraining model sizes, ZIP-SR and ZE-EDEN reduce TorchAO's mean validation-loss gap to 32-bit AdamW; in full-parameter SFT, both achieve lower validation loss than TorchAO while downstream scores remain close to 32-bit.

5 EXPERIMENTS

  • ZIP-SR combines zero-inclusive Dyn4 with Update-SR, while ZE-EDEN uses zero-excluding Dyn4-NZ with EDEN calibration; both share NF4 first moments and late LM-head SR.The LM-head first moment switches from NF4-RTN to NF4-SR during the final 10% of training.Source: §5 · pp. 7–8, Table 2 · p. 8
  • All three 4-bit configurations store eligible AdamW moments at 1.0625 bytes per parameter, a 7.53× compression versus 8 bytes per parameter for FP32 moments.The codebook and rounding changes add no per-parameter state; ZE-EDEN absorbs its calibration factor into the existing block scale.Source: §5 · p. 8

5.1 COMPONENT ATTRIBUTION

  • In GPT-style 834M component attribution, final gaps were 28.2% lower for the EDEN branch and 43.2% lower for the Dyn4 Update-SR branch than TorchAO.The latter changes both second-moment format and rounding rule and uses a single run, so it does not isolate Update-SR's effect.Source: §5.1 · p. 8, Table 3 · p. 8

5.2 PRETRAINING

  • Across every model size, ZE-EDEN and ZIP-SR reduce TorchAO’s mean paired 4-bit pretraining gap relative to matched 32-bit runs; Table 4 uses three seeds per row.Source: §5.2 · p. 9, Table 4 · p. 9
  • At the 95%-of-training-tokens checkpoint on GPT-style 2.7B, gaps were +0.0754 (TorchAO), +0.0471 (ZE-EDEN), and +0.0286 (ZIP-SR).These checkpoints precede TorchAO's terminal regression; the paper does not report a percentage reduction from its divergent endpoint.Source: §5.2 · p. 9

5.3 SUPERVISED FINE-TUNING

  • In full-parameter SFT, both methods achieved lower mean validation loss than TorchAO, while downstream scores remained close to 32-bit AdamW without a consistent task advantage.ZE-EDEN's mean paired validation-loss gap was essentially zero on Qwen3-8B and slightly negative on Llama-3.2-3B.Source: §5.3 · p. 9, Table 5 · p. 9

6 CONCLUSION

Conclusion frames second-moment quantization around its effect on optimization: Update-SR supports the zero-inclusive route, while EDEN calibration mitigates positive-floor distortion for zero-excluding codebooks. With NF4 and targeted LM-head first-moment stochastic rounding, the designs improve over TorchAO across pretraining and fine-tuning, while the theory remains limited to conditional one-step bias and toy-model dynamics.

6 CONCLUSION

  • The paper argues that second-moment quantization should account for how stored values affect optimization, not only state-space reconstruction.Update-SR motivates the zero-inclusive route; EDEN calibration mitigates distortion in the induced preconditioner distribution for the zero-excluding route.Source: §6 · pp. 9–10
  • Update-SR and EDEN calibration, combined with NF4 and targeted LM-head stochastic rounding, improve over TorchAO across pretraining and fine-tuning settings.The theory compares conditional one-step preconditioner biases and toy-model optimization dynamics, not general end-to-end convergence.Source: §6 · p. 10

A.1 4-BIT FORMATS AND DECODED CODEBOOKS

  • Block quantization uses a guarded absmax scale; ZE-EDEN sets cτ=0 for all-zero blocks so their decoded values remain zero.Without calibration, a zero-excluding codebook reconstructs an all-zero block at its positive decoded floor.Source: §A.1 · p. 14
  • Dyn4-NZ has a decoded floor of 0.00325s, while the zero-excluding Lin4-NZ mapping used in Li et al. (2023) has floor s/16.The paper distinguishes these two positive codebooks; Li et al. (2023) proposed and used Lin4-NZ in their final 4-bit AdamW configuration.Source: §A.1 · p. 15

A.2 ROUND-TO-NEAREST AND STOCHASTIC ROUNDING

  • State-RTN and State-SR round in the increasing state coordinate, whereas Update-RTN and Update-SR use the decreasing preconditioner coordinate and require a nonnegative second-moment codebook.Inputs below a zero-excluding codebook's floor saturate to that floor.Source: §A.2 · p. 15

B.1 SHARED SETTINGS

  • The 4-bit configurations use 128-element blocks with packed indices and one FP32 scale per block; ineligible moment tensors remain FP32.NF4-RTN is the default first-moment setting, with only the LM-head first moment switching to NF4-SR after 90% of training.Source: §B.1 · p. 16

B.2 PRETRAINING

The appendices specify the pretraining and supervised fine-tuning setups and outline additional diagnostics. Pretraining spans GPT- and Llama-style models, while SFT evaluates two base models on Tulu-3.

B.2 PRETRAINING

  • Pretraining covers GPT-style models from 162M to 2.7B parameters and Llama-style models from 130M to 1.1B, using FineWeb-Edu and three paired seeds.Runs use 20 tokens per parameter, an effective batch size of 256 sequences, and a warmup-stable-decay schedule; the main validation split contains 50,000 sequences.Source: §B.2 · pp. 16–17

B.3 SUPERVISED FINE-TUNING

  • Full-parameter SFT uses Qwen3-8B-Base and Llama-3.2-3B on Tulu-3, with a fixed 50,000-example training subset and one epoch.Evaluation uses five paired seeds and reports validation loss, MMLU, GSM8K, HumanEval, and strict IFEval.Source: §B.3 · p. 17

C ADDITIONAL DIAGNOSTICS

  • Additional diagnostics cover the GPT-style 834M component grid, a cosine-schedule ablation, EDEN calibration, one-step update errors, and 2.7B training dynamics.The diagnostics also include validation trajectories for the complete optimizer configurations.Source: C ADDITIONAL DIAGNOSTICS · p. 18

C.1 COMPONENT ATTRIBUTION ON GPT-STYLE 834M PRETRAINING

The GPT-style 834M component grid compares validation-loss gaps to 32-bit AdamW across first- and second-moment settings. The reported changes vary with the combination and some endpoints are single runs.

C.1 COMPONENT ATTRIBUTION ON GPT-STYLE 834M PRETRAINING

  • About 7 × 10−3 lower displayed gap followed replacing Dyn4-NZ State-RTN with Dyn4 Update-SR under either first-moment codebook.The Dyn4/Update-SR endpoints are single runs, compared with three-seed Dyn4-NZ means.Source: §C.1 · p. 18, Table 8 · p. 18
  • EDEN lowered the displayed gap by 0.39 × 10−3 under SDyn4 and by 3.05 × 10−3 after switching to NF4.These reductions differ by the first-moment codebook and the order of changes.Source: §C.1 · p. 18
  • Table 8 varies the first moment by row and the second moment by column; parenthesized entries are three-seed mean and sample standard deviation, while others are single runs.The grid uses RTN throughout for the first moment.Source: Table 8 · p. 18

C.2 SENSITIVITY TO THE LEARNING-RATE SCHEDULE

The cosine-decay comparison tests the complete recipes on GPT-style 834M. In seed 42, both recipes have a smaller gap to 32-bit AdamW than TorchAO, matching the reported ordering under WSD.

C.2 SENSITIVITY TO THE LEARNING-RATE SCHEDULE

  • The cosine comparison uses the complete ZE-EDEN and ZIP-SR configurations, including the final-10% LM-head switch to NF4-SR.Source: §C.2 · p. 18
  • The GPT-style 834M cosine-decay table reports final validation loss at update 31,810 for seed 42 and shows gaps to 32-bit AdamW directly.The test checks whether the ordering persists under cosine decay; the WSD component runs do not provide matched schedule comparisons for the complete recipes.Source: Table 9 · p. 18, §C.2 · p. 18
  • With cosine decay, both recipes had a smaller gap to 32-bit AdamW than TorchAO in seed 42, as in the WSD pretraining experiments.The comparison uses complete configurations that include the LM-head first-moment switch to NF4-SR during the final 10% of training.Source: §C.2 · p. 18

C.3 EDEN CALIBRATION FOR ZERO-EXCLUDING SECOND-MOMENT QUANTIZATION

The section compares Lin4-NZ second-moment quantization with and without EDEN calibration across three GPT-style 834M checkpoints. EDEN partially restores the distorted preconditioner distribution and modestly reduces one-step update error.

C.3 EDEN CALIBRATION FOR ZERO-EXCLUDING SECOND-MOMENT QUANTIZATION

  • The replay starts from 32-bit AdamW checkpoints at 25%, 50%, and 75%; within each codebook, raw and calibrated branches share indices and gradients, differing only in effective decoded block scales.The comparison uses GPT-style 834M pretraining and measures the first resumed update.Source: §C.3 · p. 19
  • At all three checkpoints, uncalibrated Lin4-NZ produces floor-induced peaks and suppresses the large-preconditioner tail, while EDEN partially restores that tail.Residual distributional distortion remains after calibration.Source: §C.3 · p. 19, Figure 5 · p. 19
  • EDEN reduces Lin4-NZ relative update error at 25%, 50%, and 75% of training, from 30.99% to 30.58%, 28.80% to 28.34%, and 27.87% to 27.38%, respectively.Dyn4-NZ also improves slightly at each checkpoint and has lower one-step replay error than Lin4-NZ.Source: §C.3 · p. 19

C.4 GPT-STYLE 2.7B TRAINING DYNAMICS AND LM-HEAD DIAGNOSTICS

Matched GPT-style 2.7B continuations localize the late-training instability to deterministic rounding of the LM-head first moment. During cooldown, ZE-EDEN and ZIP-SR remain close to 32-bit AdamW, while TorchAO regresses sharply.

C.4 GPT-STYLE 2.7B TRAINING DYNAMICS AND LM-HEAD DIAGNOSTICS

  • The matched continuations preserve weights, optimizer state, data order, and learning-rate schedule while varying the stated first-moment storage rule.Each branch keeps the second moment fixed to Dyn4 Update-SR with an FP32 current read.Source: §C.4 · p. 19, Table 10 · p. 19
  • Replacing NF4-RTN with NF4-SR for the LM-head first moment restores the original trajectory; applying NF4-RTN only to that state reproduces the instability.Keeping the LM-head first moment in FP32 also prevents the loss increase even when other first moments use NF4-RTN.Source: §C.4 · p. 19
  • Near the end of cooldown, the NF4-RTN LM-head pre-clipping gradient norm reaches 2.044, while NF4-SR remains below 0.388.The norms are initially comparable and separate rapidly near the end; Figure 6 marks the first NF4-RTN global clipping at step 102,500 (53.74B tokens).Source: §C.4 · p. 20, Figure 6 · p. 20
  • NF4-SR has larger per-write and accumulated errors despite avoiding inward bias and gradient growth, while NF4-RTN performs better on all three magnitude diagnostics yet becomes unstable.Table 11 reports means over 91 evaluations; lower values are better.Source: §C.4 · p. 20, Table 11 · p. 20
  • During cooldown, ZE-EDEN and ZIP-SR remain close to 32-bit AdamW, while TorchAO undergoes a sharp terminal validation-loss regression.Figure 7 shows one seed-42 run per configuration; the LM-head first-moment switch to stochastic rounding occurs at 90% of training.Source: §C.4 · p. 20, Figure 7 · p. 21

D.1 ONE-STEP STATE AND PRECONDITIONER ERRORS

The analysis shows that rounding error in second-moment state space can produce large next-step preconditioner error near zero. A scalar quadratic construction further yields different dynamics under state-space and preconditioner-space rounding.

D.1 ONE-STEP STATE AND PRECONDITIONER ERRORS

  • When the zero-write probability is Θ(1), next-step preconditioner error is +Θ(ϵ^-1); bounded error requires p0=O(ϵ).This follows under the stated small-gradient assumption, with r(0)=Θ(ϵ^-1) while r(v) and r(b) remain bounded.Source: §D.1 · p. 21
  • State-RTN has state error ∆v=-v and State-SR has ∆v=0, yet both produce ∆r=+Θ(ϵ^-1) in the zero-adjacent cell.Thus zero mean state error under State-SR does not ensure small next-step preconditioner error.Source: §D.1 · p. 22
  • For sufficiently small ϵ, Update-RTN selects the positive endpoint, giving ∆v=b-v and ∆r=-Θ(1).The comparison fixes a positive input v and decoded floor b.Source: §D.1 · p. 22

D.2 ROUNDING-SPACE SEPARATION ON A SCALAR QUADRATIC

  • On a common initialization interval, State-RTN converges to a stable nonzero two-cycle, whereas Update-RTN converges geometrically to zero.The remark explains that State-SR selects the zero branch with probability approaching one, while Update-SR makes it sufficiently rare for positive-branch contraction to dominate.Source: §D.2 · pp. 22–23

Related papers

Loading 2610.12444v1…