MARS-M: 분산 감소 기법과 행렬의 만남

초록

행렬 기반 사전조건 최적화 도구인 Muon과 같은 방법들은 최근 대규모 언어 모델(LLM)을 포함한 대규모 신경망 훈련 시 스칼라 기반 최적화 도구보다 더 효율적인 것으로 입증되었습니다. 한편, LLM 사전훈련을 위한 최적화 도구에 대한 최근 벤치마크 결과, MARS와 같은 분산 감소 기법을 적용한 최적화 도구가 분산 감소를 사용하지 않는 표준 최적화 도구 대비 상당한 속도 향상을 달성할 수 있음을 보여주었습니다. 본 논문에서는 두 방법의 장점을 모두 취하기 위해 MARS의 분산 감소 기법과 Muon을 통합한 새로운 최적화 도구인 MARS-M을 소개합니다. 표준 규칙성 조건 하에서 Muon-M이 Muon이 달성한 𝒪(T^{-1/4}) 속도보다 향상된 𝒪(T^{-1/3})의 속도로 1차 정류점에 수렴함을 증명합니다. 언어 모델링 및 컴퓨터 비전 작업에 대한 실험 결과는 MARS-M이 다양한 다운스트림 벤치마크에서 지속적으로 더 낮은 손실과 향상된 성능을 제공함을 입증합니다. MARS-M의 구현은 https://github.com/AGI-Arena/MARS/MARS_M에서 확인할 수 있습니다.

English

Matrix-based preconditioned optimizers, such as Muon, have recently been shown to be more efficient than scalar-based optimizers for training large-scale neural networks, including large language models (LLMs). On the other hand, recent benchmarks on optimizers for LLM pre-training have demonstrated that variance-reduction techniques such as MARS can achieve substantial speedups over standard optimizers that do not employ variance reduction. In this paper, to achieve the best of both worlds, we introduce MARS-M, a new optimizer that integrates the variance reduction technique in MARS with Muon. Under standard regularity conditions, we prove that Muon-M converges to a first-order stationary point at a rate of mathcal{O}(T^{-1/3}), which improves upon mathcal{O}(T^{-1/4}) rate attained by Muon. Our empirical results on language modeling and computer vision tasks demonstrate that MARS-M consistently yields lower losses and improved performance across various downstream benchmarks. The implementation of MARS-M is available at https://github.com/AGI-Arena/MARS/MARS_M.

MARS-M: 분산 감소 기법과 행렬의 만남

MARS-M: When Variance Reduction Meets Matrices

초록

Support