적응형 옵티마이저에서 마스킹 업데이트의 놀라운 효과에 관한 연구

초록

대규모 언어 모델(LLM) 훈련은 점점 더 정교한 선조건자(Preconditioner)를 적용한 밀집 적응형 최적화 알고리즘에 거의 전적으로 의존해 왔습니다. 본 연구는 이러한 관행에 도전하여, 매개변수 업데이트를 무작위로 마스킹하는 방법이 매우 효과적일 수 있음을 보여줍니다. 특히 RMSProp의 마스킹 변형이 최신 최첨단 최적화 알고리즘들을 지속적으로 능가하는 성능을 발휘했습니다. 우리의 분석에 따르면, 무작위 마스킹은 곡률에 의존하는 기하학적 정규화(Geometric Regularization)를 유발하여 최적화 궤적을 매끄럽게 만듭니다. 이러한 발견에 동기를 부여받아, 우리는 모멘텀-기울기 정렬(Momentum-gradient Alignment)을 사용하여 마스킹된 업데이트를 조절하는 Momentum-aligned gradient masking(Magma)을 제안합니다. 광범위한 LLM 사전 훈련 실험을 통해 Magma가 일관된 성능 향상과 무시할 수 있는 계산 오버헤드로 기존 적응형 최적화 알고리즘의 간단한 대체재가 될 수 있음을 입증했습니다. 특히 10억 개 매개변수 규모의 모델에서 Magma는 Adam 및 Muon 대비 각각 19% 이상, 9% 이상의 perplexity 감소를 달성했습니다.

English

Training large language models (LLMs) relies almost exclusively on dense adaptive optimizers with increasingly sophisticated preconditioners. We challenge this by showing that randomly masking parameter updates can be highly effective, with a masked variant of RMSProp consistently outperforming recent state-of-the-art optimizers. Our analysis reveals that the random masking induces a curvature-dependent geometric regularization that smooths the optimization trajectory. Motivated by this finding, we introduce Momentum-aligned gradient masking (Magma), which modulates the masked updates using momentum-gradient alignment. Extensive LLM pre-training experiments show that Magma is a simple drop-in replacement for adaptive optimizers with consistent gains and negligible computational overhead. Notably, for the 1B model size, Magma reduces perplexity by over 19\% and 9\% compared to Adam and Muon, respectively.

적응형 옵티마이저에서 마스킹 업데이트의 놀라운 효과에 관한 연구

On Surprising Effectiveness of Masking Updates in Adaptive Optimizers

초록

Support