A Convergence Analysis of Adaptive Optimizers under Floating-point Quantization

AI summary

Overview

• Research area: Machine learning theory — convergence analysis of adaptive optimizers (Adam, Muon) when their computation is carried out in low-precision floating-point formats.
• Technical level: Advanced. The paper is a mathematical optimization-theory paper built on convergence proofs, stochastic gradient assumptions and a quantization error model, aimed at readers comfortable with convergence-rate notation such as 𝒪(T^{-1/4}).
• One-sentence scope: The paper builds the first analytical framework that models quantization of weights, gradients, momentum and second-moment estimates together, and proves convergence rates for Adam and Muon that explain why low-precision LLM training still works.

What This Paper Is About

Modern large-scale training runs on low-precision formats such as BF16 and FP8 to save memory and compute, yet existing convergence theories for adaptive optimizers assume every quantity is stored and computed exactly. The authors ask why low-precision training with adaptive optimizers (Adam and Muon) still converges, and derive convergence rates that explicitly depend on how much relative error each quantized component — weights, gradients, first moment, second moment — introduces.

Key Contributions

1. A hardware-aware analytical framework for quantized adaptive training. The paper models a master-worker pipeline in which full-precision weights are quantized and sent to workers, gradient computation happens on quantized weights, gradients are quantized on the way back, and the optimizer states (momentum and second moment) are updated and re-quantized for storage. Quantization is described by a relative error model — for any value x, |x^Q − x| ≤ q|x| with q = Θ(2^{-M}), where M is the mantissa length of the target format — which the authors argue matches real floating-point behavior, unlike unbiased-quantization or error-feedback assumptions used previously.

2. First convergence guarantees for quantized Adam and quantized Muon. Theorem 4.5 covers Adam and Theorem 4.6 covers Muon on smooth non-convex objectives under standard stochastic gradient assumptions, with results written explicitly in terms of the relative errors q_W (weights), q_G (gradients), q_M (first moment) and q_V (second moment).

3. A precise account of which quantized components matter for which optimizer. The analysis shows Adam is particularly sensitive to weight and second-moment quantization because β₂ is typically set close to 1, while Muon tolerates weaker error conditions because its SVD-based sign operator does not amplify errors through an inverse square root of accumulated gradient variance.

4. Numerical experiments supporting the theory. The authors report experiments on synthetic, image and LLM benchmarks, with the synthetic Rosenbrock setup described in the main text (m = 50, n = 100, T = 10,000 iterations, learning rate η = 5 × 10^{-4}, mantissa bit-lengths chosen from M = {4, 8, 16, 24, 32, 52}).

Main Findings

• Both optimizers keep full-precision rates if precision increases only logarithmically with iterations. Under the stated hyperparameter choices, quantized Adam reaches 𝔼[‖∇F(W_τ)‖_F] = Õ(T^{-1/4}) and quantized Muon reaches (1/T)Σ𝔼[‖∇F(W_t)‖_F] = 𝒪(T^{-1/4}), matching their full-precision counterparts, provided the mantissa length satisfies M = Ω(log T) so that the relative errors decay sufficiently fast.

• Adam's error budget is quantified component by component. Theorem 4.5 requires q_G = 𝒪(1/T) and q_M = 𝒪(1/T) for gradients and first moment, but q_V = 𝒪(1/T²) for the second moment and q_W = 𝒪(1/T²) for weights, with η = Θ(1/√T) and 1 − β₂ = Θ(1/T).

• The second moment is the most delicate part of Adam. Because the second-moment estimate is fed through an inverse square root in the update, accumulated errors in V_t are non-linearly amplified, so q_V must be smaller than q_M. The paper states this gives a rigorous explanation for the empirical observation that second moments need higher precision than first moments in low-bit training.

• The strict weight-precision requirement is a worst-case artifact. The paper explains that the 𝒪(1/T²) condition on q_W comes from accounting for potentially unbounded growth of weight magnitudes over training; in practice, where weight norms often stay bounded, it says the requirement could be relaxed to 𝒪(1/T).

• Muon needs only square-root-order accuracy. Theorem 4.6 requires q_G = q_W = q_M = 𝒪(T^{-1/2}), with 1 − β = Θ(T^{-1/2}), η = Θ(T^{-3/4}) and batch size B = 1. This is a sharp contrast with Adam's at least 𝒪(T^{-1}) and in some cases 𝒪(T^{-2}) requirements.

• Muon's robustness has a structural explanation. The SVD-based sign operator avoids the amplification of quantization errors by the inverse square root of historical gradient variances, which the authors say explains why Muon adapts more efficiently to low precision than Adam, corroborating empirical observations in prior work.

• Sanity checks are consistent with prior theory. Setting q_W = q_G = q_M = 0 recovers the exact convergence rate 𝒪(1/T^{1/4}) reported for full-precision Muon up to constant factors, and the Adam rate matches established full-precision rates for smooth non-convex optimization.

• Experiments show graceful degradation. On the Rosenbrock function with Adam and Muon, larger mantissa bit-lengths produce smaller converged gradient norms over 10,000 iterations, and near full-precision performance is observed at moderate mantissa lengths. The paper reports additional experiments on image and LLM benchmarks, but the detailed setup of those experiments is not included in the provided content.

Methodology in Plain English

The authors start from a concrete picture of how low-precision training actually runs on hardware. A master copy of the model keeps weights in high precision; it sends quantized weights to workers, the workers compute gradients on those quantized weights, quantize the gradients, and send them back; the master dequantizes, updates the quantized optimizer states (momentum and second moment for Adam, momentum for Muon), and stores them again in low precision.

Rather than assuming the quantizer is unbiased (a common theoretical shortcut) or storing a per-parameter error-feedback term (memory-intensive and not used in large-scale LLM training), they assume only that quantization error is relative: quantizing x produces an error at most q|x|, where q shrinks exponentially with the number of mantissa bits. This is the property that makes floating-point formats like BF16 and FP8 behave the way they do.

With that error model in place, they run a standard smooth non-convex stochastic optimization analysis, but track how the quantization errors of each component propagate and accumulate across iterations — including the recursion through the optimizer states. The result is a bound whose leading terms look like a full-precision bound plus explicit correction terms driven by q_W, q_G, q_M and q_V. Choosing the learning rate, β₂ (or β) and the required precision then yields rates and precision requirements. They test the predictions on the Rosenbrock function with mantissa lengths from 4 to 52 bits, and report additional experiments on image and LLM benchmarks.

Why This Matters

• It closes a theory-practice gap. Previous analyses of quantized training mostly covered SGD with quantized gradients, or Adam with error feedback, or stochastic rounding analyses that quantize only the final weight update and assume full-precision gradients and optimizer states. None modeled quantization of the optimizer states, which is exactly what memory-saving low-bit training systems compress.
• It gives a principled precision budget. Instead of treating "use BF16" as a rule of thumb, the theorems say how much relative error each component can tolerate, and make it explicit that Adam's second moment and weights are the most demanding components.
• It explains why hardware formats work. Because the required relative error shrinks only polynomially in T while the mantissa length M controls error as Θ(2^{-M}), a mantissa length growing logarithmically in the number of iterations is enough to preserve full-precision convergence rates.

Real-world applications:

• Low-precision LLM pre-training. Memory and bandwidth savings from quantizing gradients, weights and optimizer states for models trained on trillion-token workloads.
• FP8 training on modern accelerators. The paper situates its analysis around 8-bit floating point being a practical datatype on NVIDIA's Hopper GPU architecture and in mixed-precision systems.
• Optimizer-state compression. Optimizer states can consume as much memory as the model parameters, motivating methods that store momentum and second moments in low precision and decompress just in time for the weight update.
• Choosing optimizers and precision per component for specific training pipelines. The Adam-versus-Muon contrast supports deciding where to spend precision bits and which adaptive optimizer is more forgiving at low bit-widths.

Industry relevance: the analysis targets the exact configuration used in state-of-the-art large-scale training frameworks (BF16, FP8, and more aggressive 4-bit efforts), so its guidance is directly actionable for teams deciding how to allocate precision across gradients, weights and optimizer states.

Future Directions

• Extending to (L₀, L₁)-smoothness. The paper explicitly states that a more general smoothness condition, in which the smoothness constant depends on the gradient norm and which has been used in recent Adam analyses, is left as an interesting direction for future work.
• Tightening the weight-precision requirement. Since the 𝒪(1/T²) condition on q_W stems from the proof allowing unbounded weight growth, deriving a bound under bounded weight norms — or empirically verifying when 𝒪(1/T) suffices — is a natural follow-up.
• Broadening the optimizer and format coverage. The framework is instantiated for Adam and Muon; it could be applied to other adaptive methods and to more aggressive formats (the related work mentions FP4 efforts), as well as to integer quantization, which the paper contrasts with floating-point formats.
• Understanding the effects the framework does not model. The analysis assumes no underflow or overflow, so the interaction between quantization error and clipping, scaling or range violations remains an open question for very low-precision regimes.

Target Audience

• Optimization and learning-theory researchers who want convergence guarantees for adaptive methods that include quantization, rather than assuming exact arithmetic.
• Machine learning systems and efficiency engineers working on mixed-precision, BF16, FP8 or lower-bit training pipelines who need to decide which components can be compressed and how much.
• Practitioners training or fine-tuning large models who want a theoretical reason for preferring one adaptive optimizer over another when precision is limited.
• Students and researchers entering the low-precision training literature who need a compact statement of what is already known about quantized SGD and quantized Adam, and what this paper adds.

Authors’ abstract

The rapid scaling of large language models (LLMs) has made low-precision training essential for reducing memory, improving efficiency, and enabling larger models and datasets. Existing convergence theories for adaptive optimizers, however, assume all components are exact and neglect hardware-aware quantization, leaving open the question of why low-precision training remains effective. We introduce the first theoretical framework for analyzing the convergence of adaptive optimizers, including Adam and Muon, under floating-point quantization of gradients, weights, and optimizer states (e.g., moment estimates). Within this framework, we derive convergence rates on smooth non-convex objectives under standard stochastic gradient assumptions, explicitly characterizing how quantization errors from different components affect convergence. We show that both algorithms retain rates close to their full-precision counterparts provided mantissa length scales only logarithmically with the number of iterations. Our analysis further reveals that Adam is highly sensitive to weights and second-moment quantization due to its reliance on $β_2 \to 1$, while Muon requires weaker error control and is thus potentially more robust. These results narrow the gap between empirical success and theoretical understanding of low-precision training methods. Numerical experiments on synthetic and real-world data corroborate our theory.