PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Dense Adam States: Adaptive Log-Space Quantization for Memory-Efficient Optimizers

Yan Wang

arXiv:2608.22322v2cs.LG

TL;DR

Quantization methods designed for dense Adam states leave open how structured states transmit reconstruction error through different topologies and update operators. The paper introduces Adaptive Log-Space quantization with exact-zero preservation and state-specific precision, then evaluates it across four optimizer paths. Results support topology-aware choices, including competitive AdamW performance and distinct precision needs for CAME and Adafactor.

  • Problem

    The central gap is limited evidence on quantization beyond dense Adam moments, where factorization, confidence states, and projected statistics change error propagation.

  • Method

    The paper characterizes state heterogeneity, evaluates update-domain fidelity, and combines Adaptive Log-Space quantization with exact-zero preservation, independent momentum encoding, and state-specific precision.

  • Results

    72.90 perplexity is achieved by AdamW with AL8 second moments and uniform 8-bit momentum, versus 72.48 for FP32 AdamW and 73.54 for dynamic quantization.

  • Takeaways & Limitations

    Optimizer-state quantization should follow the state topology and update operator rather than apply one uniform precision policy.

  • Takeaways & Limitations

    The 100K-step topology experiment does not identify grouping-by-block-size interaction because G1 at B = 2048 was not evaluated.

Abstract

from arXiv · show

Low-precision optimizer-state methods are commonly designed and evaluated for dense Adam-style first and second moments. Memory-efficient optimizers depart from this setting: Adafactor factorizes second moments, CAME adds factored confidence states, and APOLLO maintains statistics in a projected gradient space. Consequently, an equal amount of state reconstruction error can induce different update errors depending on state topology and update semantics. We first characterize this heterogeneity in optimizer-state traces from language model pre-training. We then introduce Adaptive Log-Space (AL) quantization, a block-wise representation for non-negative states that adapts its nonzero range per block and enforces the exact-zero invariant $q = 0 \Leftrightarrow x = 0$. AL8 and AL16 are combined with independent signed-momentum encodings and state-specific precision choices rather than a single policy for every state. Across 96 runs totaling 214.7 GPU-hours, we evaluate dense, factored, confidence, and projected states in AdamW, Adafactor, CAME, and APOLLO paths. On a 20K-step TinyLlama-1.1B pre-training benchmark, an AdamW configuration with AL8 second moments and uniform 8-bit momentum reaches 72.90 perplexity, compared with 72.48 for FP32 AdamW and 73.54 for an 8-bit dynamic-quantization baseline, while reducing measured optimizer-state storage from 8392.7 to 2119.2 MiB. CAME exposes a different precision regime: promoting its non-negative states to AL16 recovers 86.16 perplexity versus 86.68 for the full-precision reference, whereas all-AL8 reaches 90.19. A 100K-step GPT-2 experiment further shows that topology-aware parameter protection reduces the late-loss gap of quantized Adafactor from +0.1185 to +0.0159 in the evaluated setup. These results support a state- and topology-aware view of optimizer quantization.

1 Introduction

Optimizer-state quantization must extend beyond dense Adam moments because memory-efficient optimizers change state topology and update semantics. The paper characterizes these differences, introduces Adaptive Log-Space quantization, and evaluates state-specific configurations in four optimizer paths.

  • Motivation: Optimizer-state quantization is motivated by training-memory costs, especially because FP32 Adam moments can require twice the parameter count in auxiliary values.Prior block-wise 8-bit methods show that this cost can be reduced while retaining useful training behavior.
  • Research gap: Adafactor, CAME, and APOLLO alter optimizer-state structure through factorized statistics, confidence states, or projected-space moments.These states differ in how stored perturbations propagate to parameter updates.
  • Approach: The paper studies state ranges, exact-zero patterns, update-domain fidelity, and end-to-end pre-training across four optimizer paths.The evaluation covers full-rank, factored, and confidence states during language-model pre-training.
  • Method: Adaptive Log-Space quantization fits each block’s positive values to a logarithmic interval while reserving code zero exclusively for exact zeros.AL8 and AL16 are combined with independent signed-momentum encoding and state-specific precision choices.

2 Background and Related Work

Prior work reduces optimizer memory through block-wise quantization, lower precision, factorization, projection, or dynamic allocation. This paper instead analyzes how representation choices interact jointly with the topology and semantics of already-structured optimizer states.

  • Quantized optimizer states: Block-wise dynamic quantization established practical 8-bit optimizer states, while later work explored 4-bit states and logarithmic quantization.Prior analyses identify second-moment zero-point errors as dangerous because inverse-square-root preconditioning can amplify them.
  • Quantized optimizer states: The paper’s exact-zero treatment complements prior work by preventing both small positive values from becoming zero and genuine zeros from becoming positive.This distinguishes its bidirectional zero invariant from earlier zero-point handling.
  • Related directions: Other methods address FP8 range expansion, temporal state adaptation, state resets, or dynamic precision allocation across layers, states, and steps.The present experiments explicitly isolate topology, block boundaries, and update-operator interactions.
  • Structured optimizers: Adafactor and CAME reduce state size through factorization, but errors in stored statistics can propagate across reconstructed matrices or affect confidence-guided momentum.CAME adds factored residual statistics beyond its factored second moment.
  • Structured optimizers: APOLLO maintains moments in a projected gradient space and derives scaling for the original full-rank gradient, providing a projected-state case for quantization analysis.The paper does not propose a new optimizer topology; it studies representation selection after topology has changed.

3 Optimizer States Beyond Dense Adam

Optimizer states differ in representation, topology, and update semantics, so reconstruction error is not a sufficient proxy for training impact. The paper therefore evaluates state heterogeneity and fidelity in the update domain across dense, factored, confidence, and projected cases.

  • State Topology and Update Semantics: AdamW stores dense parameter-shaped first- and second-moment states, whereas Adafactor reconstructs second moments from row and column statistics.A perturbation to one Adafactor statistic can affect an entire reconstructed row or column.
  • State Topology and Update Semantics: CAME adds factored confidence states that track residuals and modulate momentum, so its non-negative confidence and second-moment states enter through different operators.Their shared non-negativity does not imply identical numerical requirements.
  • State Topology and Update Semantics: APOLLO’s projected-space moments derive channel-wise or tensor-wise scaling factors applied to the original full-rank gradient.Quantization therefore perturbs a derived scaling operator rather than a directly reconstructed full-rank update.
  • State Topology and Update Semantics: Quantization error may act locally, propagate through factor reconstruction, alter confidence modulation, or perturb projected scaling depending on state topology and update semantics.These distinctions motivate treating representation, topology, and update semantics as interacting properties.
  • Quantization Error in the Update Domain: State reconstruction error does not directly measure the optimizer update produced from a quantized state.The paper uses update-domain error alongside reconstruction error, drift, loss trajectories, and final metrics.
  • Quantization Error in the Update Domain: State-space and update-space errors need not be proportional because inverse square roots, factor reconstruction, and projected scaling can change perturbation sensitivity.The update-domain diagnostic is intended as a common fidelity measure, not a substitute for end-to-end training quality.
  • State and Tensor Heterogeneity: 76.7 bits versus 3.5, 4.3, and 2.4 summarizes the factored Adafactor range gap between embedding second moments and attention, MLP, and normalization states.Embedding row statistics reach an 81.5-bit median block range.
  • State and Tensor Heterogeneity: 84.8% of AdamW embedding second-moment values are exactly zero at step 50, falling to 40.3% at step 10K, while other traced states are dense.CAME also shows state-type heterogeneity: embedding second-moment and confidence row ranges reach 84.3 and 53.3 bits.

4 Adaptive Log-Space Quantization

Adaptive Log-Space quantization represents non-negative optimizer states with block-local logarithmic ranges while reserving code zero exclusively for exact zeros. The design uses state-specific precision and separate signed momentum encoding to reflect differing state distributions and update roles.

  • Adaptive Nonzero Log-Space Encoding: AL fits each block’s positive values to a local logarithmic interval while reserving code zero exclusively for exact zeros.The representation enforces q = 0 ⇔ x = 0, so positive states cannot collapse to zero and dormant states do not acquire artificial positive history.
  • Exact-Zero Reservation: Exact-zero reservation preserves dormant optimizer history, which matters for sparse or partially updated parameters such as inactive embedding rows.Reconstructing a dormant second moment as a small positive value can alter the history used when that parameter is later activated.
  • State-Specific Precision: Non-negative exponential moving averages use AL8, AL16, or full precision, while signed first moments use independent uniform low-bit encodings.The paper distinguishes second-moment and confidence states from signed momentum rather than applying one quantization scheme to all states.
  • State-Specific Precision: UF8 shows lower reconstruction error and update drift than the tested Dettmers-style dynamic mapping at the same bit width for the evaluated momentum distributions.The authors present this as an empirical observation rather than a universal property of uniform quantization.
  • Block Granularity and Memory Cost: Smaller blocks improve range localization but increase metadata overhead, so block size controls a localization–metadata tradeoff.The storage cost combines the quantized code with shared FP32 metadata, and the paper evaluates this tradeoff empirically.

5 Controlled Analysis

Controlled diagnostics show that adaptive, state-aware quantization reduces range mismatch, preserves dormant zeros, and requires precision and block choices tailored to optimizer-state structure.

  • Adaptive Versus Fixed Log Ranges: 0.736% state drift with adaptive quantization outperforms fixed floors and the evaluated dynamic baseline in a 5K-step second-moment EMA.Fixed floors produce 25.254%–74.463% drift, while the dynamic baseline produces 10.032%.
  • Exact Zero and Sparse States: 0.122% unclipped update error preserves exact zeros, whereas mapping dormant entries to positive values produces error exceeding 100%.RMS clipping can mask the immediate contribution of reconstructed dormant entries.
  • Temporal Drift: 0.562% V drift and 0.280% update error at block size 256 are lower for AL8 than the dynamic baseline’s 7.282% and 3.436%.At block size 2048, AL8 remains lower at 0.736% drift and 0.357% update error versus 10.032% and 4.577%.
  • Block Granularity and Tensor Topology: 8.25 bits/element at block size 256 versus 8.031 bits/element at 2048 illustrates the metadata trade-off of smaller blocks.The corresponding 1.1B-element state estimates are 1081.8 and 1053.1 MiB.
  • Block Granularity and Tensor Topology: 0.577% to 0.814% drift across block sizes remains below the dynamic baseline’s increase from 8.243% to 12.259%.The numerical gap favors larger blocks as a memory-efficient default in this synthetic setting.
  • Precision as an Independent State Decision: Uniform 8-bit momentum is the default empirical choice, while sensitive non-negative states can require AL16 or FP32 because precision needs differ by state.Increasing a dense second moment from AL8 to AL16 roughly doubles storage, whereas factored confidence states are much smaller.

6 LLM Pre-training Evaluation

The pre-training evaluation compares quantized optimizer paths across TinyLlama and GPT-2, showing that precision, block size, and parameter grouping affect outcomes differently by optimizer.

  • Benchmark setup: 96 runs totaling 214.7 GPU-hours evaluate AdamW, CAME, Adafactor, and APOLLO across 20K-step and 100K-step language-model benchmarks.The core benchmark uses TinyLlama-1.1B on WikiText-103; the long-horizon benchmark uses GPT-2 124M.
  • AdamW: 72.90 perplexity is achieved by AdamW with AL8 second moments and UF8 momentum, versus 72.48 for FP32 AdamW and 73.54 for the evaluated dynamic 8-bit baseline.Measured optimizer-state storage falls from 8392.7 MiB for FP32 AdamW to 2119.2 MiB, while throughput rises from 2502 to 2960 tokens/s.
  • CAME: 86.16 perplexity with AL16 for both CAME non-negative states approaches the 86.68 full-precision reference, whereas all-AL8 reaches 90.19.Keeping confidence state C in FP32 narrows the gap with only 2.3 MiB additional measured state.
  • Adafactor and APOLLO: 78.15 versus 78.29 is the G1-versus-G0 comparison for Adafactor at B = 256, while B = 256 outperforms B = 2048 under both grouping policies.The factored reference uses only 3.6 MiB of measured optimizer state, making the comparison primarily about topology sensitivity.
  • Sensitivity: No method dominates the tested learning-rate range, with dynamic 8-bit AdamW and all-AL8 CAME collapsing at 10× while other listed configurations remain finite.The batch-size sweep supports the AdamW result across the tested range but shows a larger CAME precision gap.
  • Long-horizon training: +0.1185 versus +0.0159 is the Adafactor late-loss gap under G0 versus G1 at B = 256 in the 100K-step GPT-2 experiment.The long-horizon study lacks G1 at B = 2048, so it does not estimate a grouping–block-size interaction.

7 Implementation

The implementation provides multiple optimizer paths and independently configurable state encodings, while documenting reproducibility, validation, and scope limitations.

  • Implementation: One PyTorch optimizer interface exposes AdamW-style, Adafactor, CAME, and APOLLO update paths with AL8, AL16, and FP32 non-negative-state storage.Signed momentum supports uniform and dynamic low-bit encodings or FP32, and CAME confidence precision is configured independently.
  • Implementation: A JIT-compiled CUDA extension fuses dequantization, EMA updates, requantization, and parameter updates where the selected path permits it.A Python implementation serves as the semantic reference.
  • Validation: 171 PASS, 0 FAIL, 4 WARN, and 0 SKIP are reported by the complete smoke run.The CUDA warning is attributed to accumulation order in factored atomicAdd reductions; all-AL8 CAME shows differences around 1.25 × 10^-4 versus Python.
  • Limitations: The study covers one 1.1B architecture for 20K steps and one 124M architecture for 100K steps on WikiText-103 and does not establish transfer beyond this setting.End-to-end comparisons use seed 921, so repeated seeds would quantify variation in final losses.
  • Limitations: Optimizer-state bytes are not total training memory, and throughput reflects one implementation and GPU rather than other hardware or software stacks.Peak allocation excludes allocator fragmentation, and some values use matched 1K proxy runs.
  • Limitations: Synthetic diagnostic families cannot enumerate real optimizer-state distributions, and update-domain error does not guarantee final task quality.The diagnostics are complementary to, rather than substitutes for, language-model results.

8 Conclusion

The study concludes that optimizer quantization must account for state representation, topology, and update semantics rather than apply one uniform policy.

  • Conclusion: Quantization policies designed for dense Adam states do not transfer uniformly to factored, confidence, and projected optimizer states.The paper frames these differences in terms of how state perturbations affect downstream updates.
  • Conclusion: Adaptive Log-Space quantization provides a block-adaptive, exact-zero-preserving representation for non-negative states, but behavior still depends on precision, topology, and the downstream update operator.The conclusion is based on evaluations of dense, factored, confidence, and projected states.
  • Conclusion: No single quantizer is uniformly best across the evaluated paths; CAME needs greater precision care, while Adafactor benefits from sensitive-tensor protection and finer blocks.APOLLO was evaluated on a projected-state path but did not provide a throughput win in the reported implementation.

A Additional Results

The appendix records the fixed measurement conditions and distinguishes the main figures from supporting numeric results.

  • Additional results: Unless stated otherwise, language-model measurements use seed 921 and controlled probes use fixed random generators recorded in the executed analysis notebook.The appendix presents numeric results supporting the main figures without moving their visual narrative out of the paper.

A.1 Trace-Calibrated State Sensitivity

Trace-calibrated probes show that required precision depends jointly on a state’s traced range and consuming operator. Dyn8 embedding failures coincide with positive entries reconstructed as zero, but these probes do not establish universal update-error or training-quality rankings.

  • 32,768-element synthetic states are calibrated to tensor-wise medians and median block log2 ranges from step-10K CAME traces.The probes are controlled one-step operator tests rather than samples of the empirical block distribution.
  • Required precision depends jointly on traced state range and the operator consuming the state.
  • Dyn8 embedding failures coincide with many positive entries being reconstructed as zero.
  • One-step rankings do not establish that every real embedding update exceeds 100% error or predict final training quality.

A.2 Sensitivity Tables

The sensitivity tables report trajectory mean absolute errors against path-specific full-precision references across learning-rate and batch-size sweeps. They are diagnostic rather than repeated-seed uncertainty estimates, and CAME’s ordering changes with non-negative-state precision.

  • Trajectory mean absolute error is reported relative to the path-specific full-precision reference in the 10K learning-rate sweep.Collapsed runs are recorded at their observed failure step rather than assigned finite error.
  • The batch-size sweep reports corresponding trajectory mean absolute errors at each path’s nominal learning rate.
  • Sensitivity tables are diagnostic rather than repeated-seed uncertainty estimates.
  • CAME’s ordering changes with the precision assigned to its non-negative states, while a finite 10× AdamW result does not imply universal stability.

A.3 Artifact and Source Map

The reproducibility artifact maps figures, tables, traces, and environment claims to scripts, exports, reports, and ancillary artifacts. Reported memory quantities come from PyTorch and live CUDA-tensor counters, not operating-system resident-VRAM measurements.

  • Adafactor8Bit v0.4.3 at commit 4f544d8 anchors the optimizer implementation used for the experiments.The commit does not contain the manuscript or figure-generation sources.
  • The public artifact provides a claim-to-source map linking figures and benchmark tables to scripts, CSV exports, notebooks, and TensorBoard analysis.Figures 1–2 use matching scripts; Figure 3 uses fig03_trace_components.csv; Figure 4 uses checked CSV tables; Figure 5 and benchmark tables use tb_analysis_report.md.
  • Reference, CUDA, and environment claims use paired Markdown/JSON smoke-test and environment reports generated by reproducibility scripts.
  • Raw TensorBoard events and state-trace snapshots are ancillary artifacts, while machine-readable reports retain manuscript values.
  • Windows peak memory uses PyTorch allocation counters and optimizer-state storage uses live CUDA tensors, so neither measures total resident VRAM.
  • Tables 5–7 document trace-calibrated single-step update error and 10K-step learning-rate and batch-size sensitivity against full-precision references.Table 5 also reports Dyn8 positive-entry false-zero counts out of 32,768 elements.
Loading 2608.22322v2…