Stable-SPAM: 16비트 Adam보다 더 안정적으로 4비트에서 학습하는 방법

Stable-SPAM: How to Train in 4-Bit More Stably than 16-Bit Adam

February 24, 2025

저자: Tianjin Huang, Haotian Hu, Zhenyu Zhang, Gaojie Jin, Xiang Li, Li Shen, Tianlong Chen, Lu Liu, Qingsong Wen, Zhangyang Wang, Shiwei Liu

cs.AI

초록

본 논문은 4비트 학습을 위해 최근 제안된 여러 옵티마이저를 종합적으로 평가하며, 저비트 정밀도가 학습률에 대한 민감도를 증폭시키고 종종 불안정한 그래디언트 노름을 유발하여 높은 학습률에서 발산을 일으키는 것을 밝혀냈습니다. 이 중에서도 모멘텀 리셋과 스파이크 인식 그래디언트 클리핑을 특징으로 하는 최신 옵티마이저인 SPAM은 다양한 비트 수준에서 최고의 성능을 달성했지만, 그래디언트 노름을 안정화하는 데 어려움을 겪어 신중한 학습률 조정이 필요했습니다. 이러한 한계를 해결하기 위해, 우리는 향상된 그래디언트 정규화 및 클리핑 기술을 통합한 Stable-SPAM을 제안합니다. 특히, Stable-SPAM은 (1) 스파이크 그래디언트에 대한 클리핑 임계값을 과거 최대값을 추적하여 적응적으로 업데이트하고, (2) 전체 그래디언트 행렬을 과거 l_2-노름 통계를 기반으로 정규화하며, (3) SPAM의 모멘텀 리셋을 상속받아 Adam의 첫 번째와 두 번째 모멘트를 주기적으로 리셋함으로써 스파이크 그래디언트의 누적을 완화합니다. 광범위한 실험을 통해 Stable-SPAM이 4비트 LLM 학습에서 그래디언트 노름을 효과적으로 안정화시키며, Adam과 SPAM에 비해 우수한 성능을 제공하는 것을 확인했습니다. 특히, Stable-SPAM으로 학습된 4비트 LLaMA-1B 모델은 Adam으로 학습된 BF16 LLaMA-1B보다 최대 2의 perplexity 차이로 우수한 성능을 보였습니다. 또한, 두 모델 모두 4비트로 학습할 때 Stable-SPAM은 Adam과 동일한 손실을 달성하면서도 약 절반의 학습 단계만을 필요로 했습니다. 코드는 https://github.com/TianjinYellow/StableSPAM.git에서 확인할 수 있습니다.

English

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.