Adam keeps two FP32 moments per parameter, twice the bytes of the FP32 weights, yet compressing them is risky: the update divides by the square root of the second moment, so small reconstruction errors can change the restored computation substantially. We study guarded moment compression, in which a coordinate's decoded moments are accepted only if they reproduce the native FP32 Adam update to within a tolerance τ, in units of the learning rate, for every gradient in a clipped domain; all other coordinates are stored exactly. The natural certificate, which ranges each update separately over gradient intervals, fails: on one million trained OPT-125M coordinates it escapes 98.9% of coordinates and the payload grows to 126.6% of raw FP32. Our key observation is that both updates consume the same gradient and parameter. Bounding their difference directly, with FP32 rounding-error allowances and a fixed mixed linear-logarithmic gradient partition, cuts escapes to 0.53% and the payload to 28.3% at τ=0.01. Across 120 complete-checkpoint restores of OPT-125M and ResNet-18, with four trained seeds and three training fractions per model, guarded 8-bit log codes shrink moment payloads by 69.5% and 71.9% relative to FP32, 40.8% and 45.4% relative to equally guarded BF16, and 29.3% and 37.6% relative to equally guarded blockwise 8-bit, and every restored first update stays within tolerance. Encoding and restoration together take 2.3-2.6× as long as guarded BF16 but less time than guarded blockwise 8-bit. Searching both moment codes jointly saves under 0.001% more bytes at 6% more codec time, so per-moment rounding suffices. The guarantee covers the first restored update under a pinned Torch operation model; we measure subsequent training empirically.
Figures & tables
Figure 1: Bound the difference, not the outputs: the shared-gradient bound with a mixed partition shrinks the guarded payload from 126.6% to 28.3% of raw FP32, below equally guarded BF16 and blockwise 8-bit. (a) Illustration with representative moments ( m=2×10−5 , v=10−9 , bias corrections ≈1 ) and decoded moments perturbed by +0.2% and −0.4% . Over the gradient interval I , both updates sweep a range far wider than τ , so bounding them separately gives 122τ ; the shared-gradient bound of Equation ( 9 ) gives 0.48τ , against a true maximum of 0.23τ (rounding allowances omitted). (b) Measured on one million trained OPT-125M coordinates at τ=0.01 with independent code selection. Every configuration spends the same 27.7% of raw FP32 bytes on codes and metadata; the bound decides how many coordinates escape, at 8 bytes each. Vertical lines mark equally guarded BF16 and blockwise 8-bit payloads on the same coordinates.
Figure 2: Guarded codec design. (a) The encoder proposes 8-bit codes per 256-coordinate block, decodes them exactly, and guards every coordinate with an ideal-real bound and the native shared-gradient bound; coordinates that fail are escaped to exact FP32. (b) First-level edges of the two fixed partitions on a log-magnitude axis. A uniform partition covers all ∣g∣<C/32 with one interval on each side of zero; the mixed partition adds base-2 logarithmic edges down to 2−40C . (c) Byte layout of the serialized payload for the one-million-coordinate subset. The header and SHA-256 checksum (879 bytes in total) are too small to show.
Representation
Bytes
% of FP32
Escapes
Raw FP32
8,000,000
100.000
—
Serialized FP32
8,000,595
100.007
0
Plain BF16
4,125,281
51.566
0
Guarded BF16
4,125,333
51.567
0
Plain blockwise 8-bit
2,158,768
26.985
0
Guarded blockwise 8-bit
3,177,969
39.725
127,396
Table 1: Complete moment payloads for the same one million trained OPT-125M coordinates ( τ=0.01 , block size 256), relative to 8,000,000 raw FP32 bytes. Guarded rows share the same ideal-real and native guards. Plain blockwise 8-bit is smallest but violates the contract.
Figure 3: Complete-checkpoint results. (a) Moment payload at each checkpoint fraction; small dots are the four trained seeds and large dots their mean. FP32 is 100%. Log codes use independent selection; joint selection differs by less than 0.002% and would overlap. (b) Moment pairs escaped to exact FP32. Guarded BF16 escapes at most 0.014% of OPT-125M coordinates and none of ResNet-18’s. (c) Payload versus combined encode-and-restore codec time relative to FP32, from the separate timing cohorts (seed 701, 50% checkpoint). Markers are geometric means over processes, whiskers span the process range, and hollow markers show joint selection.
OPT-125M
ResNet-18
Moment payload
Complete container
Moment payload
Complete container
Method
10%
50%
100%
10%
50%
100%
10%
50%
100%
10%
50%
100%
FP32
955.50
955.50
955.50
1993.32
1998.13
1997.27
85.31
85.31
85.31
165.78
166.02
166.44
Guarded BF16
492.81
492.81
492.80
1294.60
1299.44
1298.57
43.99
43.99
43.99
103.15
103.39
103.80
Guarded 8-bit
372.49
412.03
455.61
1218.29
1259.65
1309.16
32.51
41.87
42.77
95.42
108.61
109.82
Log, indep. (ours)
270.65
291.96
312.01
1021.20
1029.89
1056.70
23.73
24.01
24.26
78.92
79.33
80.17
Table 2: Complete-checkpoint sizes in MiB, averaged over four trained seeds at each training fraction. The moment payload includes serialization metadata and checksums; the complete container additionally holds the lossless model and the remaining checkpoint state. Payload and container are different denominators.
Figure 4: Restore fidelity and continuation for the 96 guarded restores. (a) Largest realized first-update error over all coordinates, per condition; every guarded condition stays below τ , and the log-code and blockwise guards accept errors up to nearly the whole budget. (b) Development-loss difference from the paired FP32 branch after 100 own-gradient updates, one marker per seed and checkpoint (note the different scales). Hollow markers show joint selection. Descriptive only; Table 6 in Appendix C lists means and ranges.
OPT-125M ( N=17 processes)
ResNet-18 ( N=50 processes)
Time (s)
Peak (GiB)
Time (s)
Peak (GiB)
Method
Encode
Restore
Total
Encode
Restore
Encode
Restore
Total
Encode
Restore
FP32
2.130
1.825
3.955
3.622
3.737
0.184
0.156
0.340
0.324
0.334
Guarded BF16
6.205
1.112
7.317
19.167
3.737
0.512
0.078
0.590
1.719
0.334
Guarded blockwise 8-bit
20.864
2.664
23.528
17.350
3.737
1.588
0.245
1.833
1.516
0.334
Log, independent (ours)
15.307
1.598
16.905
24.888
10.866
1.402
0.120
1.521
2.261
0.973
Table 3: Paired GPU codec timing at seed 701, 50% checkpoint: arithmetic means of process means over 17 OPT-125M and 50 ResNet-18 processes, each with 20 warm-ups and 100 measured rounds per method. Peaks include the resident model and moments.
Figure 5: Supporting study on 200 small logistic-optimization states with a binary64 guard. (a) Mean complete payload relative to raw FP32 moments, including metadata and escapes. (b) Largest new-probe error relative to the assigned tolerance; all formats also satisfy the 514 inherited probes. Zero tested violations are a numerical observation, not continuous native-operation certification.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Paired restoration and its measurement units. (a) OPT-125M and ResNet-18 each contribute four independently trained seeds; three checkpoint fractions are nested within each seed, and all five representations are paired within each checkpoint. (b) Every representation restores complete moment tensors, shares the first saved-gradient audit, and then continues with its own gradients from the same starting model, RNG state and data order. Ordered development evaluation is descriptive, not held-out quality equivalence. (c) The separate timing cohorts contain 17 OPT and 50 ResNet processes. Their synchronized codec intervals exclude the additional whole-payload SHA check, disk transfer, full-container serialization and later training.
Method
Bits
τ
Bytes / FP32
Max. error / τ
Independent
4
0.005
0.9735
0.9832
Independent
4
0.05
0.6979
0.9888
Independent
8
0.005
0.5987
0.9873
Independent
8
0.05
0.3819
0.9888
Joint
4
0.005
0.9740
0.9832
Joint
4
0.05
0.7071
0.9870
Appendix
Table 4: Complete serialized bytes divided by two raw FP32 moment arrays ( 8d ), with the largest new-probe error relative to τ . Each row retains all 200 states. Guarded rows use approximate FP32 escapes with an explicit exact-FP64 fallback. Both baseline formats pass every probe budget at both tolerances.
Bound
Independent codes
Joint codes
Separate ranges
0.43545
0.44392
Shared gradient, mixed partition
0.16056
0.16320
Appendix
Table 5: Warm encode medians in seconds on the same one million OPT-125M coordinates.
Method
OPT-125M ( ×10−4 )
ResNet-18 ( ×10−2 )
Guarded BF16
−0.58 [ −2.78 , 1.45 ]
0.59 [ −2.04 , 3.85 ]
Guarded 8-bit
−0.07 [ −1.91 , 0.93 ]
−0.10 [ −4.63 , 2.48 ]
Log, independent
−0.50 [ −3.31 , 0.73 ]
−0.69 [ −6.33 , 1.07 ]
Log, joint
−0.73 [ −2.96 , 0.38 ]
0.43 [ −0.26 , 1.96 ]
Appendix
Table 6: Development-loss difference from the paired FP32 branch after 100 own-gradient updates: mean and [min, max] over 4 seeds × 3 checkpoints. Descriptive only; positive values mean higher loss than FP32.
Optimizer-state quantization is commonly designed for Adam's dense, parameter-aligned first- and second-moment arrays. This abstraction breaks for memory-efficient optimizers, whose states may be factored, confidence-modulated, or maintained in a projected space, so similar reconstruction error can produce different update error. We formulate optimizer-state quantization as a joint problem over representation, topology, and update semantics. We then introduce Adaptive Log-Space (AL) quantization for non-negative states. AL fits each block's observed nonzero logarithmic interval and reserves a separate code for exact zero, enforcing q=0⇔x=0; signed momentum and state precision remain independently selectable. Controlled probes show that adaptive ranges reduce update error and temporal drift, exact-zero reservation preserves dormant states, and state topology constrains useful block granularity. End-to-end language-model training evaluates the resulting policy across dense, factored, confidence, and projected optimizer states. On TinyLlama-1.1B, AL8 with uniform 8-bit momentum reaches 72.90 perplexity versus 73.54 for bitsandbytes 8-bit AdamW, with comparable optimizer-state storage and higher throughput. CAME matches reference-level final perplexity across three seeds when its non-negative states use AL16, while a semantic grouping-and-protection policy closes most of quantized Adafactor's 100K-step late-loss gap. These results make state topology and update semantics first-class design constraints for optimizer quantization.
Adam is the default optimizer for training modern deep neural networks, yet its adaptive behavior remains poorly understood due to the complex interaction between its first- and second-moment exponential moving averages (EMAs). We study Adam in the tied-β regime, where the two EMA decay rates are equal, and show that its adaptive dynamics can be expressed through a transformed ratio with approximately scale-stable behavior. Empirically, this transformed ratio exhibits a stable, heavy-tailed distribution across tasks, model scales, and training stages, in contrast to the variability of raw moment magnitudes. This empirical stability has both practical and conceptual consequences. First, we derive a recurrence for the transformed ratio, yielding a reparameterization of Adam that replaces the second moment with a compressible state. Leveraging its stable distribution, we show that a fixed 4-bit codebook is sufficient in our experiments to store this state without auxiliary scaling, achieving performance competitive with full-precision Adam. Second, the transformed ratio view clarifies Adam's connection to sign-based methods: Adam reduces to sign-based momentum modulated by the transformed ratio, and replacing it with a constant recovers Signum as a limiting case. This perspective further provides a simple rule for transferring learning rates between the two methods. Together, these results suggest that tied-β Adam admits a simple and approximately stable ratio structure underlying its adaptive behavior and demonstrate its utility for both analysis and efficient implementation.
Yihe Zhou, Tongtian Zhu, Yingxiao Huo +4
Department of Computer Science, The University of Manchester, Manchester, UK · College of Computer Science and Technology, Zhejiang University, Hangzhou, China
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.
Hanyang Li, Shao Tang, Daniel Thomas Braithwaite +7