New Method for Compressing Adam Optimizer State: Shared Gradient Bounds Safeguard Update Accuracy

Title

New Method for Compressing Adam Optimizer State: Shared Gradient Bounds Safeguard Update Accuracy

Meta Description

This article introduces a new method for compressing Adam optimizer state: guarded moment compression via shared gradient bounds, ensuring decoded moments can reproduce FP32 updates. Experiments show ~70% reduction in moment payload on OPT-125M and ResNet-18, with first updates staying within tolerance.

Article Content

When training large models, the Adam optimizer maintains two FP32 moments (first-order and second-order momentum) for each parameter, doubling the storage overhead relative to the FP32 weights themselves. Compressing these moments seems like a way to save significant memory and storage, but it is risky: Adam's update rule divides by the square root of the second moment, so small reconstruction errors can be amplified, significantly altering the restored computation. A recent study proposes "Guarded Moment Compression," which uses shared gradient bounds to ensure that decoded compressed moments can reproduce native FP32 updates, offering a new approach to optimizer state compression.

Why Is Compressing Adam Moments Dangerous?

Each parameter in the Adam optimizer corresponds to two FP32 moments: the first moment estimates the mean of gradients, and the second moment estimates the mean of squared gradients. During updates, the parameter update equals the learning rate times the first moment divided by the square root of the second moment (plus a small constant). This division structure means that small errors in the second moment directly cause large deviations in the update. More specifically, if the second moment is underestimated, the update is amplified; if overestimated, the update is shrunk. In extreme cases, compression errors can cause training to diverge.

Existing compression methods (such as block quantization and BF16 storage) typically focus only on the storage precision of moments, ignoring whether the compressed updates remain consistent with the original updates. This study proposes a stricter goal: for each coordinate, accept the compression only if the decoded moments can reproduce the native FP32 Adam update (within a tolerance τ, measured in learning rate units) over a given gradient domain; otherwise, store that coordinate exactly.

The first and second moments in Adam are essentially exponential moving averages (EMAs) of gradients and their squares, with update formulas: m_t = β₁·m_{t-1} + (1-β₁)·g_t, v_t = β₂·v_{t-1} + (1-β₂)·g_t², where g_t is the current gradient and β₁ and β₂ are decay rates (typically 0.9 and 0.999). Because moment estimates are biased toward zero early in training, Adam also introduces bias correction: m̂_t = m_t/(1-β₁^t), v̂_t = v_t/(1-β₂^t), and the final parameter update is θ_t = θ_{t-1} - η·m̂_t/(√v̂_t + ε). This structure places the square root of the second moment v_t directly in the denominator, so any compression-induced error in v_t is nonlinearly amplified—for example, if v_t is underestimated by 10%, then √v_t is underestimated by about 5%, and the update is amplified by about 5.3%. Furthermore, moments are typically stored at the same precision as model parameters (FP32), and in models with billions of parameters, moment memory can be twice that of the weights, becoming a key training bottleneck. Therefore, compressing moments is not only about storage efficiency but also requires ensuring numerical consistency of update behavior.

Why Does the Natural Certificate Fail?

Intuitively, one could compute a certificate for each update individually: let the gradient vary within an interval and check whether the range of the update stays within tolerance. However, the researchers found that this "natural certificate" fails severely on 1 million trained OPT-125M coordinates: 98.9% of coordinates fail validation, causing the payload to grow to 126.6% of the original FP32. The reason is that when each update is considered separately, the gradient interval is over-amplified because the two updates actually share the same gradient.

Shared Gradient Bounds: A Key Observation

The researchers' core observation is that the updates of the first and second moments consume the same gradient and the same parameter. Rather than bounding the two updates separately, it is better to directly bound the difference between them. Specifically, they use the FP32 rounding-error allowance and a fixed hybrid linear-log gradient partition to directly compute a bound on the difference between the two updates. The effect is striking: at τ=0.01, the fraction of escaping coordinates drops from 98.9% to 0.53%, and the payload drops from 126.6% to 28.3%.

Encoding and Recovery Process

The guarded compression pipeline has two stages: during encoding, each moment is encoded with 8-bit logarithmic coding, and the decoded update is checked to see whether it is within tolerance; if not, the original FP32 value is stored. During recovery, the moments are decoded and the update is computed. The entire process guarantees that the first update after recovery matches the original update within tolerance. Note that this guarantee covers the first recovered update under a fixed Torch operation model; subsequent training behavior is measured experimentally.

The root cause of the "natural certificate" failure is that it treats the updates of the first and second moments independently, whereas in reality the two updates share the same gradient sample. In Adam, the updates of the first and second moments depend on the same gradient g_t, but the natural certificate treats g_t as an arbitrary value within an interval and computes the possible ranges of the two updates separately, causing the interval to be over-conservatively enlarged. The shared gradient bound exploits this correlation: it directly considers the difference between the two updates (i.e., the difference between the first-moment update and the second-moment update) as a whole and introduces the FP32 rounding-error allowance and a hybrid linear-log gradient partition to tighten the bound. The hybrid linear-log partition divides the gradient domain into linear intervals (for small gradients) and logarithmic intervals (for large gradients) to more precisely cover the actual gradient distribution. Through this joint bounding, the certificate can more accurately identify which coordinates are safe to compress, greatly reducing the escape rate.

8-bit logarithmic encoding is a compression format that converts floating-point numbers into 8-bit logarithmic values, typically including a sign bit, exponent bits, and mantissa bits. For example, a configuration of 1 sign bit, 4 exponent bits, and 3 mantissa bits can represent a wide dynamic range with fewer bits, but with relatively low precision. In guarded compression, each moment is independently encoded with 8-bit logarithmic coding, immediately decoded, and the update is computed and compared with the original FP32 update. If the error exceeds the tolerance τ (in units of learning rate), the moments for that coordinate are reverted to exact FP32 storage. This "encode-verify-fallback" mechanism ensures compression safety for each coordinate. During recovery, simply decoding the stored moments (8-bit or FP32) and performing one Adam update step guarantees a deviation from the original update within τ. This process does not depend on future gradients, so the certificate is statically computable.

Experimental Validation: Compression Rate and Accuracy

The researchers conducted extensive experiments on OPT-125M and ResNet-18, including 120 full checkpoint restorations, using four training seeds and three training fractions per model. Results show that guarded 8-bit logarithmic encoding reduces moment payload by 69.5% and 71.9% relative to FP32, by 40.8% and 45.4% relative to equally guarded BF16, and by 29.3% and 37.6% relative to equally guarded block-wise 8-bit. All restored first updates remain within tolerance.

The total encoding and recovery time is 2.3-2.6 times longer than guarded BF16 but shorter than guarded block-wise 8-bit. In addition, jointly searching the encodings of both moments saves less than 0.001% in bit rate but increases codec time by 6%, so per-moment rounding is already sufficient.

Implications for AI Training

This study provides a framework with accuracy guarantees for optimizer state compression. As model sizes grow, optimizer state becomes one of the memory bottlenecks, and guarded compression can significantly reduce moment storage overhead while maintaining training stability. However, the current guarantee covers only the first recovered update; the stability of subsequent training (e.g., multi-step updates) still relies on empirical validation. Future work may need to extend the guard to multi-step updates or explore more efficient encoding schemes.

For practitioners, the core value of this method is that compression is no longer about blindly pursuing fewer bits, but rather about using update tolerance as a constraint to ensure that compressed training aligns with native training from the very first step. This provides a quantifiable basis for safely compressing Adam state.

One-line Summary

Guarded moment compression uses shared gradient bounds to ensure that the first update after Adam state compression matches FP32, greatly reducing storage overhead.

Paragraph Summary

This article introduces a method for compressing Adam optimizer state called "Guarded Moment Compression." Adam maintains two FP32 moments per parameter, doubling storage overhead relative to weights, but compression is risky because updates divide by the square root of the second moment, amplifying small errors. Existing methods focus only on storage precision, whereas this study requires that compressed moments can reproduce native FP32 updates (within tolerance τ) over the gradient domain. The researchers found that a "natural certificate" computing a certificate for each update separately fails severely, with 98.9% of coordinates failing validation; by exploiting the observation that the first and second moments share the same gradient and directly bounding the difference between the two updates, the escape rate drops to 0.53% and the payload to 28.3% (τ=0.01). Experiments show that guarded 8-bit logarithmic encoding reduces moment payload by 69.5% and 71.9% on OPT-125M and ResNet-18, respectively, and all restored first updates are within tolerance. This method provides an accuracy-guaranteed framework for optimizer state compression.