Stable-SPAM: How to Train in 4-Bit More Stably than 16-Bit Adam (2502.17055v2)

Abstract: This paper comprehensively evaluates several recently proposed optimizers for 4-bit training, revealing that low-bit precision amplifies sensitivity to learning rates and often causes unstable gradient norms, leading to divergence at higher learning rates. Among these, SPAM, a recent optimizer featuring momentum reset and spike-aware gradient clipping, achieves the best performance across various bit levels, but struggles to stabilize gradient norms, requiring careful learning rate tuning. To address these limitations, we propose Stable-SPAM, which incorporates enhanced gradient normalization and clipping techniques. In particular, Stable-SPAM (1) adaptively updates the clipping threshold for spiked gradients by tracking their historical maxima; (2) normalizes the entire gradient matrix based on its historical $l_2$-norm statistics; and $(3)$ inherits momentum reset from SPAM to periodically reset the first and second moments of Adam, mitigating the accumulation of spiked gradients. Extensive experiments show that Stable-SPAM effectively stabilizes gradient norms in 4-bit LLM training, delivering superior performance compared to Adam and SPAM. Notably, our 4-bit LLaMA-1B model trained with Stable-SPAM outperforms the BF16 LLaMA-1B trained with Adam by up to $2$ perplexity. Furthermore, when both models are trained in 4-bit, Stable-SPAM achieves the same loss as Adam while requiring only about half the training steps. Code is available at https://github.com/TianjinYellow/StableSPAM.git.

An Analytical Review of "Stable-SPAM: How to Train in 4-Bit More Stably than 16-Bit Adam"

Abstract and Introduction

The paper "Stable-SPAM: How to Train in 4-Bit More Stably than 16-Bit Adam" addresses the escalating computational and memory demands in large-scale LLMs by embracing low-bit precision training. It investigates the challenges posed by the volatile nature of gradient norms and learning rates in 4-bit precision adaptive optimizers. The authors advance their proposed solution, Stable-SPAM, which incorporates enhanced gradient normalization and adaptive spike-aware clipping methods to counter these challenges.

Key Contributions

The primary contributions of the paper are the introduction and demonstration of the Stable-SPAM optimizer. Through detailed empirical analysis, Stable-SPAM shows improved stability and performance compared to both SPAM and Adam, especially in low-bit precision settings. The key techniques integrated into Stable-SPAM include:

1. Adaptive Spike-Aware Clipping (AdaClip): It dynamically adjusts clipping thresholds for spiked gradients by monitoring their historical maxima, differing from the fixed thresholds used previously.
2. Adaptive Gradient Norm (AdaGN): Employing historical $l_2$-norm statistics, it normalizes the entire gradient matrix, stabilizing gradient-induced fluctuations.
3. Momentum Reset (MoRet): Inherited from SPAM, it periodically resets the first and second moments, mitigating the accumulation of spiked gradients that can destabilize training with Adam.

Empirical Evaluation

The empirical evaluation section of the paper rigorously tests Stable-SPAM across multiple LLMs, focusing on the LLaMA architectures with varied parameter sizes, using both FP4 and INT4 quantization formats. The experiments reveal:

• Superior Stability and Performance: Stable-SPAM effectively mitigates the volatility of gradient norms and learning rates, achieving lower perplexity scores across different LLaMA model sizes compared to other advanced optimizers like SPAM and Adafactor.
• Efficiency in Training Steps: It consistently reaches performance levels comparable to traditional Adam using significantly fewer training steps when operating in a 4-bit precision setup.
• Broad Applicability: The paper explores the integration of AdaClip and AdaGN into other optimizers, demonstrating compatibility and performance gains with the likes of Lion and Adam-mini.

Discussion on Optimization Stability

The paper’s analysis highlights a pressing need in low-precision LLM training frameworks—addressing optimizer sensitivity to hyperparameter selection. Stable-SPAM’s design principles contribute significantly to the discourse on training stabilization, particularly through its dual focus on robust gradient normalization and adaptable spike-aware mechanisms. These components effectively counteract training instability-inducing factors such as loss and gradient norm spikes.

Implications and Future Directions

This work paves the way for more resource-efficient training of large-scale LLMs, offering a template for integrating adaptive and resilient optimization strategies in similar contexts. Theoretical implications extend to the refinement of gradient-based methods in deep learning, stressing the utility of historical gradient statistics and adaptive thresholding. Practically, these findings advocate for the deployment of Stable-SPAM in domains constrained by computational resources.

Looking forward, future research could explore the extension of Stable-SPAM strategies to other tensor computation contexts beyond LLMs, such as computer vision and reinforcement learning tasks in AI. Additionally, further exploration into hardware acceleration and integration, particularly where low precision is pivotal, would drive forward the debottlenecking of AI scalability in computing environments.

Conclusion

"Stable-SPAM: How to Train in 4-Bit More Stably than 16-Bit Adam" makes a significant contribution to low-bit optimization technology. By integrating adaptive gradient clipping and normalization strategies, Stable-SPAM provides a well-founded solution to the known instabilities in low-precision training regimes. The paper constructs a solid foundation on which subsequent research can build, ultimately advancing both theoretical understanding and practical applications of efficient large-scale model training.