Muon+: Towards Better Muon via One Additional Normalization Step

이 논문은 기울기 직교화 후 추가적인 정규화 단계를 도입한 Muon+ 를 제안하여 다양한 모델 규모와 아키텍처에서 기존 Muon 최적화기보다 훈련 및 검증 퍼플렉시티를 일관되게 향상시키는 것을 입증합니다.

핵심

이 논문은 인공지능 (AI) 을 더 똑똑하게 만드는 데 쓰이는 '학습 방법'을 조금 더 개선한 이야기를 담고 있습니다. 전문 용어 대신, 거대한 도서관을 정리하는 일에 비유해서 설명해 드릴게요.

📚 배경: 거대한 도서관 정리하기 (AI 학습)

지금 전 세계의 AI(예: 챗지피티, 라마 등) 는 방대한 양의 책 (데이터) 을 읽으며 학습합니다. 이 책을 읽을 때, AI 는 "어떤 부분을 더 집중해서 읽어야 할까?"를 결정하는 **학습 규칙 (옵티마이저)**이 필요합니다.

기존에 가장 많이 쓰인 규칙은 '무슨 일이든 다 조금씩 골고루 챙기는' 방식이었습니다. 하지만 AI 가 커질수록 이 방식은 너무 느리고 비효율적이었습니다.

그래서 최근 등장한 **'무온 (Muon)'**이라는 새로운 규칙이 주목받았습니다. 무온은 책장 (데이터) 을 정리할 때, 책들이 서로 겹치지 않고 바르게 서 있도록 (직교화) 도와주는 아주 똑똑한 정리법입니다. 덕분에 AI 가 훨씬 빠르게, 그리고 더 똑똑해졌습니다.

---

💡 이 연구의 핵심: "정리 후, 한 번 더 다듬기" (MUON+)

하지만 연구진들은 "무온도 훌륭하지만, **정리된 책장 위에 한 번 더 손질 (정규화)**을 가하면 더 완벽하지 않을까?"라고 생각했습니다.

이것이 바로 이 논문이 제안하는 **'MUON+(무온 플러스)'**입니다.

🧹 비유: 방 정리하기

1. 기존 무온 (Muon): 책들을 책장에 꽂을 때, 책들이 서로 뒤섞이지 않고 바르게 서 있도록 정리합니다. (책들이 서로 수직이 되도록 만듦)
2. 새로운 MUON+: 책들이 바르게 서 있는 것을 확인한 후, 책장 전체의 높이를 다시 맞춰주고 책들이 너무 튀어나오지 않도록 한 번 더 다듬습니다.

이 '한 번 더 다듬는 작업'이 바로 **정규화 (Normalization)**입니다. 단순히 책장을 정리하는 것만으로는 부족하고, 정리된 상태의 균형을 맞춰주는 과정이 필요하다는 것을 발견한 것입니다.

---

🚀 왜 이것이 중요한가요? (실제 효과)

연구진은 이 방법을 다양한 크기의 AI 모델 (작은 것부터 거대한 것까지) 에 적용해 보았습니다. 결과는 놀라웠습니다.

• 더 빠른 학습: 같은 양의 책을 읽어도, MUON+ 를 쓴 AI 가 더 빨리, 더 잘 이해했습니다.
• 더 똑똑한 AI: 학습이 끝난 후 AI 의 성능 (정답률) 이 기존 무온보다 consistently(일관되게) 더 높았습니다.
• 거대한 규모에서도 작동: 책이 100 권이든 100 만 권이든, 이 '한 번 더 다듬기' 방법은 항상 효과를 발휘했습니다.

🌟 핵심 메시지

이 논문은 **"기존에 아주 잘 작동하던 방법 (무온) 에 아주 작은 추가 단계 (정리 후 균형 맞추기) 만 더해도, AI 의 성능이 크게 좋아질 수 있다"**는 것을 증명했습니다.

마치 요리를 할 때, 재료를 다 잘라낸 후 (무온) 소금 간을 한 번 더 맞춰주는 (MUON+) 것과 같습니다. 재료 정리가 잘 되어 있다면, 마지막 간 맞추기 한 번이 요리의 맛을 결정적으로 바꿔주는 것처럼 말입니다.

이 방법은 앞으로 더 크고 강력한 AI 를 만들 때, 비용을 아끼면서도 더 좋은 성능을 내는 데 큰 도움이 될 것으로 기대됩니다.

기술 요약

논문 개요: MUON+

이 논문은 대규모 언어 모델 (LLM) 의 전학습 (Pre-training) 효율성을 높이기 위해 기존 Muon 옵티마이저에 단순하지만 효과적인 추가 정규화 (Normalization) 단계를 도입한 **MUON+**를 제안합니다. Muon 은 뉴턴 - 슈르츠 (Newton-Schulz) 반복을 통해 기울기 (또는 모멘텀) 행렬을 직교화하여 기울기 랭크 붕괴 (Rank Collapse) 를 방지하는 최적화 알고리즘으로, 최근 Kimi 나 GLM 과 같은 주요 모델의 전학습에 널리 사용되고 있습니다. 저자들은 직교화 단계 직후에 적용되는 정규화 단계가 성능 향상의 핵심 동인임을 발견하고, 이를 체계적으로 분석하여 MUON+ 를 개발했습니다.

---

1. 문제 정의 (Problem)

• 계산 비용의 증가: 모델 파라미터와 학습 데이터셋의 규모가 극대화됨에 따라 전학습의 계산 비용이 prohibitive(부담스러운) 수준에 도달했습니다.
• 옵티마이저의 한계: Adam/AdamW 가 여전히 주류이지만, 대규모 전학습의 계산 및 메모리 비용을 줄이기 위한 효율적인 옵티마이저에 대한 연구가 활발합니다.
• Muon 의 개선 필요성: Muon 은 기울기 랭크 붕괴를 해결하여 좋은 성능을 보이지만, 더 큰 규모의 모델과 더 긴 학습 기간 (Long-horizon) 에서 최적화 안정성과 최종 모델 품질을 더욱 향상시킬 수 있는 여지가 존재합니다.

2. 방법론 (Methodology)

**MUON+**는 Muon 의 업데이트 규칙을 기반으로 하되, 직교화 (Orthogonalization) 단계 직후에 정규화 (Normalization) 단계를 추가합니다.

• 기존 Muon 업데이트:

  1. 모멘텀 업데이트: $M_t = \mu M_{t-1} + (1-\mu)G_t$
  2. 직교화: $O_t = \text{Ortho}(M_t)$ (뉴턴 - 슈르츠 반복 등을 통해 근사)
  3. 가중치 업데이트: $W_t = W_{t-1} - \eta \sqrt{m/n} \cdot O_t$

• MUON+ 업데이트:

  1. 모멘텀 업데이트: 동일
  2. 직교화: 동일
  3. 정규화 (핵심 변경점): $O_t = \text{Norm}^{(d)}(\text{Ortho}(M_t))$
     • 여기서 $\text{Norm}^{(d)}$는 특정 방향 $d$ (행 또는 열) 로 적용되는 정규화 연산자입니다.
     • 정규화 방향: 열 단위 (Column-wise), 행 단위 (Row-wise), 또는 두 가지를 결합한 형태 (Col-Row, Row-Col) 를 고려합니다.
  4. 가중치 업데이트: 동일

• 핵심 통찰: 최근 NorMuon, Maon 등 Muon 의 변형 연구들이 정규화 단계를 포함하고 있음을 관찰했습니다. 저자들은 통제된 실험 (Ablation Study) 을 통해, 추가적인 적응적 스케일링 (Second-moment adaptation) 이나 매니폴드 형식보다 '직교화된 업데이트에 대한 구조적 정규화' 자체가 성능 향상의 주된 원인임을 규명했습니다.

3. 주요 기여 (Key Contributions)

1. MUON+ 제안: Muon 에 단순한 정규화 단계 하나를 추가함으로써 성능을 획기적으로 개선하는 새로운 옵티마이저를 제안했습니다.
2. 광범위한 실험 검증:
   • 모델 규모: 130M 에서 774M 파라미터 (GPT 스타일) 및 60M 에서 1B 파라미터 (LLaMA 스타일) 까지 다양한 규모의 모델을 대상으로 실험했습니다.
   • 학습 regime: 계산 최적화 (Compute-optimal, T2P 비율 $\approx 20$) 및 산업 수준의 긴 학습 기간 (Overtraining, T2P 비율 $\approx 200$) 모두에서 유효성을 입증했습니다.
3. 체계적인 분석 (Ablation Study):
   • 학습률 (Learning Rate) 에 대한 민감도 감소 효과 확인.
   • 다양한 정규화 방향 (Col, Row, Col-Row 등) 의 성능 비교.
   • 다양한 직교화 근사 방법 (You, Jordan, PolarExpress) 에 대한 무관성 (Orthogonalization-agnostic) 입증.
4. 오픈소스 제공: 구현 코드를 공개하여 커뮤니티의 재현과 활용을 장려했습니다.

4. 실험 결과 (Results)

• GPT 모델 (130M ~ 774M):
  • MUON+ 는 모든 모델 규모에서 Muon 보다 일관되게 더 낮은 검증 퍼플렉시티 (Validation Perplexity) 를 기록했습니다.
  • GPT-Small 은 -2.02, GPT-Base 는 -1.72, GPT-Large 는 -0.91 만큼의 퍼플렉시티 감소를 보였습니다.
• LLaMA 모델 (60M ~ 1B):
  • AdamW 와 Muon 을 베이스라인으로 비교했을 때, MUON+ 가 모든 규모에서 최상의 성능을 달성했습니다.
  • 특히 LLaMA-1B 에서 Muon 대비 -0.37, AdamW 대비 훨씬 큰 개선을 보였습니다.
• 과학습 (Overtraining) 실험 (T2P $\approx 200$):
  • GPT-Base 와 LLaMA-350M 을 720 억 토큰으로 학습시켰을 때, MUON+ 는 Muon 보다 여전히 우월한 성능을 유지하며 최적화 안정성이 뛰어났습니다. 이는 정규화 단계의 이점이 계산 최적화 regime 에 국한되지 않음을 의미합니다.
• 학습률 민감도:
  • MUON+ 는 Muon 에 비해 더 넓은 범위의 학습률에서 안정적인 성능을 보이며, 특히 비최적의 (너무 큰) 학습률 사용 시 성능 저하가 적었습니다.
• 정규화 방향:
  • 행 (Row) 단위 정규화가 열 (Col) 단위보다 일반적으로 더 좋은 성능을 보였으며, 행과 열을 결합한 방식 (Col-Row, Row-Col) 이 가장 우수한 결과를 기록했습니다.

5. 의의 및 결론 (Significance)

• 단순함의 힘: 복잡한 2 차 모멘트 적응이나 매니폴드 최적화 같은 복잡한 메커니즘 없이, 단순한 정규화 단계 하나만 추가함으로써 대규모 LLM 전학습의 성능을 일관되게 향상시킬 수 있음을 증명했습니다.
• 확장성: MUON+ 는 모델 규모가 커지고 학습 데이터가 늘어날수록 (Long-horizon) 그 효과가 유지되거나 강화되어, 차세대 초대규모 모델 전학습에 필수적인 최적화 전략이 될 수 있습니다.
• 실용성: 학습률 스케줄링을 재조정할 필요가 없으며, 기존 Muon 구현체에 쉽게 통합할 수 있어 실제 산업 적용에 매우 유리합니다.

결론적으로, 이 논문은 Muon 옵티마이저의 핵심 메커니즘인 '직교화' 이후의 '정규화'가 최적화 안정성과 최종 모델 품질을 결정하는 중요한 요소임을 규명하고, 이를 통해 **MUON+**라는 더 강력한 옵티마이저를 제시함으로써 대규모 언어 모델 학습 효율성 향상에 기여했습니다.