適応型オプティマイザにおけるマスキング更新の驚くべき有効性について

要旨

大規模言語モデル（LLM）の学習は、高度に洗練された前処理行列を備えた密な適応型オプティマイザにほぼ独占的に依存している。我々はこれに異を唱え、パラメータ更新をランダムにマスキングする手法が極めて有効であることを実証する。RMSPropのマスキング変種は、最新の最先端オプティマイザを一貫して凌駕する性能を示した。分析の結果、ランダムマスキングが曲率依存の幾何学的正則化を誘発し、最適化軌道を平滑化することが明らかになった。この知見に基づき、我々は運動量-勾配アラインメントを用いてマスキング更新を調整するMomentum-aligned gradient masking（Magma）を提案する。大規模なLLM事前学習実験により、Magmaが計算オーバーヘッドを無視できる程度に抑えつつ、適応型オプティマイザの単純な代替として一貫した性能向上をもたらすことを実証した。特に1Bパラメータモデルでは、MagmaはAdamと比較して19％以上、Muonと比較して9％以上のパープレキシティ低減を達成した。

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