Improved Convergence Rates of Muon Optimizer for Nonconvex Optimization

이 논문은 기존 연구의 제한적인 가정을 피하고 직접적인 분석을 통해 뮤온 (Muon) 옵티마이저의 수렴 속도를 개선하고 더 넓은 문제 설정을 포괄하는 정교한 수렴 보장을 제시합니다.

핵심

1. 배경: 왜 새로운 도구가 필요할까?

AI 모델을 훈련시킨다는 것은, **수만 가지 재료가 섞인 거대한 냄비 (모델)**에서 가장 맛있는 맛 (최적의 성능) 을 찾아내는 과정입니다.

• 기존 도구 (Adam 등): 재료를 섞을 때 "이건 소금, 저건 설탕"처럼 각 재료를 따로따로 조절하는 방식입니다. 잘 쓰이지만, 재료가 너무 많고 섞이는 방식이 복잡하면 맛이 일정하지 않거나 시간이 오래 걸릴 수 있습니다.
• 무온 (Muon): 이 도구는 재료를 섞을 때 특이한 나침반을 사용합니다. 이 나침반은 재료를 섞는 방향을 '직교 (서로 수직)'하게 만들어줍니다. 마치 재료가 서로 부딪히지 않고 깔끔하게 섞이도록 도와주는 것이죠. 실제로는 이 도구가 기존 도구보다 훨씬 빠르고 안정적으로 맛을 낸다는 것이 알려져 있었습니다.

하지만 문제는 **"왜 이렇게 잘 되는지, 수학적으로 정확히 얼마나 빨리 맛을 낼 수 있는지"**에 대한 이론이 부족하다는 점입니다. 기존 연구들은 너무 가정을 많이 하거나, 실제 상황과 맞지 않는 조건에서만 증명되어 있었습니다.

2. 이 연구의 핵심: "더 정확한 지도를 만들다"

이 논문은 무온이 왜 잘 작동하는지, 그리고 어떤 조건에서 가장 빠르게 수렴 (최적의 맛에 도달) 하는지를 기존보다 훨씬 더 정교하고 일반적인 조건에서 증명했습니다.

🌟 핵심 비유: "배 (Batch Size) 를 키우는 전략"

무온의 성능을 결정하는 두 가지 중요한 변수가 있습니다.

1. 학습률 (Learning Rate): 한 번에 얼마나 큰 걸음으로 나아가는가?
2. 배치 크기 (Batch Size): 한 번에 몇 개의 재료를 함께 섞는가?

이 논문은 **"배치 크기를 점점 키워가면 무온은 훨씬 더 빠르게 목표를 달성한다"**는 놀라운 사실을 수학적으로 증명했습니다.

• 기존의 생각: "한 번에 재료를 조금씩 섞어가며 천천히 맛을 봐야 해." (작은 배치)
• 이 논문의 발견: "처음엔 조금씩 섞다가, 시간이 갈수록 한 번에 섞는 재료 양을 기하급수적으로 늘려보자!" (지수적으로 증가하는 배치)
  • 마치 요리를 할 때, 처음엔 작은 숟가락으로 맛을 보다가, 요리가 거의 완성될 때는 큰 주전자를 써서 한 번에 섞어주는 것과 같습니다. 이렇게 하면 수렴 속도가 기존 연구에서 알려진 것보다 훨씬 빨라집니다.

3. 주요 성과 (간단한 요약)

이 연구는 무온이 다음과 같은 조건에서 최고의 성능을 낸다고 증명했습니다.

1. 더 빠른 속도: 기존 연구들이 "느리게 갈 수도 있다"고 했던 것보다, 훨씬 빠른 속도로 최적의 상태에 도달할 수 있음을 보였습니다.
2. 더 넓은 적용: 무온이 작동하려면 "너무 까다로운 조건"이 필요하다는 오해를 불식시켰습니다. 일반적인 상황에서도 잘 작동함을 증명했습니다.
3. 실용적인 팁:
   • 학습률: 처음엔 크게 걸어가다가 점점 작게 걸어야 합니다.
   • 배치 크기: 시간이 지날수록 한 번에 섞는 재료 양을 두 배, 네 배로 늘려가면 (지수적 증가) 가장 효율적입니다.

4. 결론: 이 연구가 우리에게 주는 메시지

이 논문은 무온이라는 도구가 단순히 "실험적으로 잘 작동한다"는 것을 넘어, **"수학적으로도 매우 강력하고 효율적인 도구"**임을 증명했습니다.

한 줄 요약:

"AI 를 가르칠 때 무온을 쓴다면, 한 번에 섞는 재료의 양을 시간이 갈수록 점점 더 많이 늘려가면서 학습하면, 기존에 알던 어떤 방법보다 훨씬 빠르고 정확하게 최고의 모델을 만들 수 있다!"

이 연구는 AI 개발자들이 무온을 더 효과적으로 설정하고, 더 빠르고 안정적인 AI 모델을 만드는 데 이론적인 나침반이 되어줄 것입니다.

기술 요약

논문 요약: 비볼록 최적화를 위한 Muon 옵티마이저의 개선된 수렴 속도

1. 문제 제기 (Problem)

• 배경: 대규모 딥러닝 모델 (DNN) 학습에서 Muon (Momentum orthogonalized by Newton-Schulz) 옵티마이저는 일차 정보 (First-order information) 에 기반하여 업데이트 방향을 직교화 (Orthogonalization) 하는 독특한 메커니즘으로 인해 Adam 등 기존 옵티마이저 대비 뛰어난 성능과 확장성을 보여 주목받고 있습니다.
• 현황: Muon 의 실용적 유효성은 입증되었으나, 이에 대한 이론적 분석은 아직 미흡합니다. 기존 연구들 (Tang et al., 2025; Chang et al., 2025 등) 은 수렴성을 보장하지만 다음과 같은 한계가 있습니다:
  • 과도한 가정: PL(Polyak-Łojasiewicz) 조건과 같은 강력한 가정을 필요로 하거나, 업데이트 규칙에 대한 제한적인 분석 설정을 사용합니다.
  • 부족한 수렴 속도: 기존 분석에서 도출된 수렴 속도 (예: $O(T^{-1/4})$) 는 실제 성능보다 느리거나, 학습률과 배치 크기 설정에 따라 최적의 속도를 보장하지 못합니다.
  • 비완전한 분석: 신경망 파라미터 차원 ($n$) 이 분석 결과에 남아서 엄밀한 의미의 완전한 수렴 결과가 도출되지 않은 경우가 있습니다.
• 목표: 제한적인 가정 없이 Muon 의 업데이트 구조를 직접적이고 간소화된 방식으로 분석하여, 더 넓은 문제 설정에서 적용 가능하고 기존보다 빠른 수렴 속도를 보장하는 이론적 근거를 마련하는 것입니다.

2. 방법론 (Methodology)

• 가정 (Assumptions):
  • 비볼록 최적화: 손실 함수 $f$는 비볼록 (Nonconvex) 이며 미분 가능합니다.
  • 스무스함 (Smoothness): 각 데이터 포인트별 손실 함수는 $L_i$-스무스합니다 (그라디언트의 Lipschitz 연속성).
  • 확률적 그라디언트: 미니배치 그라디언트는 편향되지 않으며 (Unbiased), 분산이 유계 (Bounded) 입니다.
  • 제약 조건 부재: 기존 연구와 달리 업데이트 규칙에 대한 과도한 제약을 두지 않고, 일반적인 학습률 ($\eta_t$) 과 배치 크기 ($b_t$) 설정을 고려합니다.
• 분석 기법:
  • 감소 보조정리 (Descent Lemma): 스무스 함수의 성질을 활용하여 $f(W_t) - f(W_{t+1})$의 하한을 유도합니다.
  • 직교화 구조 분석: Muon 의 핵심인 업데이트 방향의 직교화 단계 (Step 9 in Algorithm 1) 를 수학적으로 정밀하게 분석하여 오차 항 ($\|\nabla f(W_t) - C_t\|_F$) 을 제어합니다.
  • 상한 유도: 학습률, 배치 크기, 모멘텀 파라미터 ($\beta$) 에 따른 전체 기대값 $E[\|\nabla f(W_t)\|_F]$의 상한을 유도하고, 이를 $O$ 표기법으로 변환하여 수렴 속도를 도출합니다.

3. 주요 기여 (Key Contributions)

1. 정밀한 수렴 상한 (Tight Convergence Upper Bound):
   • Muon 옵티마이저 (Nesterov 유무 포함) 가 생성하는 전체 그라디언트의 기대값 상한을 유도했습니다.
   • 이 상한은 학습률 ($\eta_t$), 배치 크기 ($b_t$), 모멘텀 ($\beta$) 에 의존하는 5~6 개의 항으로 구성되며, 기존 연구보다 더 정교한 구조를 가집니다.
2. 개선된 수렴 속도 (Improved Convergence Rates):
   • 다양한 학습률 전략 (상수, 코사인 어닐링, 다항식 감쇠, 감쇠 학습률) 과 배치 크기 전략 (상수, 지수 증가) 의 조합에 대해 구체적인 수렴 속도를 제시했습니다.
   • 핵심 발견:
     • 배치 크기 증가의 효과: 배치 크기를 고정하지 않고 증가시킬 때 (특히 지수적으로 증가, $b_t = b\delta^t$), 수렴 속도가 획기적으로 개선됨을 증명했습니다.
     • 최적 수렴 속도: 감쇠 학습률 ($\eta_t = \eta/\sqrt{t+1}$) 과 지수적으로 증가하는 배치 크기를 결합할 때, $O(\frac{\log T}{\sqrt{T}})$의 수렴 속도를 달성할 수 있음을 보였습니다. 이는 기존 $O(T^{-1/4})$나 $O(T^{-1/2})$보다 우월합니다.
3. 일반화된 조건 하의 보장:
   • PL 조건과 같은 강력한 가정을 배제하고, 일반적인 비볼록 최적화 환경에서도 Muon 의 수렴성을 보장합니다.

4. 주요 결과 (Key Results)

논문은 Table 1 과 Corollary 3.1 을 통해 다양한 시나리오에서의 수렴 속도를 정리했습니다. 주요 결과는 다음과 같습니다:

• 기존 결과와의 비교:
  • 기존 연구 (1, 4, 6, 7) 는 $O(T^{-1/4})$의 속도를 보였습니다.
  • Li & Hong (2025) 의 연구 (5) 는 배치 크기를 $T$로 설정 시 $O(T^{-1/2})$를 보였으나, 본 논문은 이를 개선했습니다.
• 본 논문의 개선된 결과 (R1 - R5):
  • 상수 학습률 + 상수 배치 크기: $O(\frac{1}{T} + \eta + \frac{1}{\sqrt{b}})$. $\eta = O(1/T), b=O(T^2)$일 때 $O(1/T)$ 달성.
  • 상수 학습률 + 지수 증가 배치 크기: $O(\frac{1}{T} + \eta)$. $\eta = O(1/T)$일 때 $O(1/T)$ 달성.
  • 감쇠 학습률 + 지수 증가 배치 크기: $O(\frac{\log T}{\sqrt{T}})$ 달성.
  • 이는 Muon 이 적절한 하이퍼파라미터 (학습률 스케줄링 및 배치 크기 스케일링) 설정 시, 기존 SGD 및 Adam 계열 옵티마이저 이론적 한계를 넘어서는 빠른 수렴을 가능하게 함을 의미합니다.

5. 의의 및 시사점 (Significance)

• 이론적 기반 확립: Muon 옵티마이저의 실용적 성공에 대한 이론적 근거를 강화하여, 대규모 모델 학습에서의 신뢰성을 높였습니다.
• 실무적 가이드라인 제공:
  • Muon 을 사용할 때 **배치 크기를 고정하지 않고 학습 과정 중 증가시키는 것 (Exponentially growing batch size)**이 수렴 속도와 안정성 측면에서 매우 유리함을 이론적으로 증명했습니다.
  • 학습률과 배치 크기를 총 반복 횟수 $T$의 함수로 적절히 결합하면 최적의 수렴 속도를 얻을 수 있음을 제시하여, 하이퍼파라미터 튜닝에 대한 명확한 지침을 제공합니다.
• 확장성: 본 논문의 분석 기법은 Muon 에 국한되지 않고, 직교화된 일차 최적화 방법 (Orthogonalized first-order methods) 전반에 적용 가능한 통찰을 제공합니다.

결론적으로, 이 논문은 Muon 옵티마이저가 제한된 가정 없이도 기존 방법론보다 우수한 수렴 속도를 가질 수 있음을 수학적으로 증명하였으며, 특히 증가하는 배치 크기 전략이 Muon 의 성능을 극대화하는 핵심 요소임을 밝혔습니다.