[논문리뷰] MARS-M: When Variance Reduction Meets Matrices

링크: 논문 PDF로 바로 열기

저자: Yifeng Liu, Angela Yuan, Quanquan Gu

핵심 연구 목표

본 논문은 대규모 언어 모델(LLM) 및 딥러닝 모델 훈련의 효율성과 안정성을 향상시키기 위해, 행렬 기반 전처리 옵티마이저의 장점과 분산 감소(variance reduction) 기법의 장점을 결합하는 것을 목표로 합니다. 특히, 기존 분산 감소 기법이 행렬 기반 옵티마이저와 통합될 때 발생하는 성능 저하 문제를 해결하고자 합니다.

핵심 방법론

저자들은 MARS-M이라는 새로운 옵티마이저를 제안합니다. 이는 MARS의 분산 감소 기술을 Muon의 Moonlight 버전에 통합한 것입니다. 이 방법론은 Newton-Schulz 반복법을 활용하여 그래디언트 모멘텀을 근사하고, 스케일링된 그래디언트 보정과 그래디언트 클리핑을 적용하여 안정성을 높입니다. 또한, 계산 효율성을 위한 근사 MARS-M 버전도 제시합니다.

주요 결과

MARS-M은 이론적으로 **O(T^-1/3)**의 수렴율을 달성하여 기존 Muon 옵티마이저의 Õ(T^-1/4) 수렴율을 개선했음을 증명합니다. GPT-2 모델을 이용한 LLM 훈련 실험에서 OpenWebText 및 FineWeb-Edu 100B 데이터셋에서 Moonlight 및 AdamW 대비 일관되게 더 낮은 훈련 및 검증 손실을 기록했습니다. 예를 들어, FineWeb-Edu 100B 데이터셋에서 소형 모델(0-shot) 훈련 시, **MARS-M (γ=0.025)**은 평균 45.55점을 달성하여 **Muon (44.95)**보다 우수했습니다. 또한, CIFAR-10 컴퓨터 비전 태스크에서도 Moonlight 대비 더 나은 테스트 정확도와 낮은 손실을 보였습니다.

AI 실무자를 위한 시사점

MARS-M은 LLM과 같은 대규모 딥러닝 모델 훈련에서 더 빠르고 안정적인 최적화를 가능하게 하는 실용적인 솔루션을 제공합니다. 이론적 수렴율 개선과 실증적 성능 향상은 AI 엔지니어들이 복잡한 모델을 효율적으로 훈련하고 더 나은 품질의 모델을 개발하는 데 기여할 수 있습니다. 근사 버전의 제공은 계산 리소스 제약이 있는 환경에서도 활용 가능성을 높입니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.