MARS‑M: 행렬 기반 최적화와 분산 감소의 시너지

초록

MARS‑M은 행렬 기반 프리컨디셔닝 옵티마이저 Muon(및 경량 변형 Moonlight)에 MARS 스타일의 분산 감소 기법을 결합한 새로운 옵티마이저이다. 이론적으로는 표준 Muon이 달성하는 Õ(T⁻¹/⁴) 수렴률을 Õ(T⁻¹/³) 로 개선했으며, 실험에서는 GPT‑2 계열 모델의 언어 모델링 및 이미지 분류 작업에서 손실 감소와 다운스트림 성능 향상을 입증한다.

상세 분석

본 논문은 대규모 신경망, 특히 대형 언어 모델(LLM) 훈련에 있어 두 가지 주요 흐름을 통합한다. 첫 번째는 파라미터를 행렬 형태 그대로 다루어 2차 구조와 고유값 정보를 활용하는 Muon·Moonlight 계열의 행렬 기반 프리컨디셔닝 기법이다. Muon은 모멘텀 행렬 Mₜ 에 대해 SVD를 수행하고, Newton‑Schulz 반복을 통해 UₜVₜᵀ 근사를 얻어 업데이트에 적용한다. Moonlight은 0.2·p_{max(m,n)} 스케일링과 가중치 감쇠 λ 을 도입해 실제 LLM 훈련에서의 스텝 크기 불균형을 보정한다. 두 번째 흐름은 STORM 기반의 분산 감소 기법인 MARS이다. MARS는 기존 STORM의 그라디언트 차분 항에 스케일링 파라미터 γₜ 를 곱해, 고분산 스토캐스틱 그라디언트의 변동을 효과적으로 억제한다.

MARS‑M은 이러한 두 메커니즘을 자연스럽게 결합한다. 구체적으로, Moonlight의 모멘텀 업데이트 Uₜ 와 Mₜ 에 MARS의 교정 그라디언트 cₜ = ∇f(Xₜ,ξₜ)+γₜβ/(1−β)(∇f(Xₜ,ξₜ)−∇f(Xₜ₋₁,ξₜ₋₁)) 를 삽입한다. 교정된 그라디언트는 클리핑을 통해 ‖cₜ‖₂≤1 을 보장하고, 이후 β 와 (1−β) 가중치를 이용해 모멘텀 mₜ 을 업데이트한다. 최종 파라미터 업데이트는
Xₜ₊₁ = Xₜ − ηₜ·(0.2·p_{max(m,n)}·Oₜ + λXₜ)
와 같이 행렬 프리컨디셔너 Oₜ≈UₜVₜᵀ 에 스케일링과 가중치 감쇠를 동시에 적용한다.

이론적 분석에서는 표준 가정(목표 함수 F 의 L‑Lipschitz 연속성, 제한된 노이즈 분산, bounded second‑moment 등) 하에, MARS‑M이 기대 그라디언트의 제곱 노름 평균에 대해 Õ(T⁻¹/³) 수렴률을 달성함을 증명한다. 핵심 아이디어는 MARS의 스케일링이 분산을 O(1/T) 수준으로 감소시키면서, 행렬 프리컨디셔닝이 효과적인 방향 탐색을 제공해 전체 복합 오차를 O(T⁻¹/³) 으로 압축한다는 점이다. 이는 기존 Muon이 보인 Õ(T⁻¹/⁴) 수렴률보다 한 차수 높은 속도이며, 특히 대규모 배치와 높은 차원의 파라미터 공간에서 이득이 크게 나타난다.

실험에서는 GPT‑2‑small, ‑medium, ‑large 모델을 OpenWebText와 FineWeb‑Edu 100B 데이터셋에 각각 1‑2 배치 크기로 학습시켰다. MARS‑M은 동일 하이퍼파라미터 설정 하에 Muon·Moonlight 대비 훈련 손실이 평균 3‑5 % 낮았으며, 검증 퍼플렉시티도 유의미하게 개선되었다. 다운스트림 평가에서는 Hellaswag, SciQ와 같은 추론 벤치마크에서 정확도가 각각 1.8 점, 2.3 점 상승했으며, 이미지 분류(CIFAR‑10/100)에서는 Top‑1 정확도가 0.9 %~1.2 % 향상되었다. 또한, 근사 버전(이전 단계 그라디언트를 재사용)도 메모리와 연산량을 약 30 % 절감하면서도 성능 격차가 미미함을 확인했다.

전체적으로 MARS‑M은 행렬 기반 프리컨디셔닝과 최신 분산 감소 기법을 성공적으로 융합함으로써, 대규모 모델 훈련에서 이론적 수렴 속도와 실험적 효율성을 동시에 끌어올렸다. 향후 연구에서는 더 복잡한 비선형 프리컨디셔너(예: PolarGrad)와의 결합, 그리고 초대규모 LLM(수십 억 파라미터)에서의 스케일링 특성을 탐색할 여지가 있다.