HTMuon: Improving Muon via Heavy-Tailed Spectral Correction

이 논문은 무언 (Muon) 알고리즘의 중력 꼬리 분포 억제 문제를 해결하고 성능을 향상시키기 위해 헤비-테일 자기 정규화 이론에 기반한 'HTMuon'을 제안하고, 이를 통해 LLM 사전 학습 및 이미지 분류에서 기존 최첨단 방법보다 우수한 결과를 입증합니다.

핵심

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

기존의 상황 (Muon 이라는 요리사)
최근 'Muon'이라는 새로운 학습 방법이 화제가 되었습니다. Muon 은 기존 방법 (Adam 등) 보다 훨씬 잘 작동해서 많은 AI 모델이 이걸 쓰고 있습니다.

• 비유: Muon 은 등산할 때 모든 방향을 똑같은 힘으로 밀어주는 나침반 같은 역할을 합니다. 등산로 (데이터) 의 지형에 따라 방향을 잘 잡아주지만, 모든 발걸음에 똑같은 힘을 줍니다.

문제점 (너무 똑같은 힘의 부작용)
하지만 연구자들은 Muon 에 숨겨진 두 가지 문제를 발견했습니다.

1. 소음 (Noise) 을 너무 믿는다: 등산할 때 가끔은 바람이나 나뭇잎이 흔들리는 '소음'이 있습니다. Muon 은 이 소음이 섞인 작은 방향에도 똑같은 힘을 주어, 길을 잃게 만들 수 있습니다.
2. 다양성을 죽인다: 잘 훈련된 AI 는 마치 거대한 오케스트라처럼 다양한 주파수 (중요한 정보와 덜 중요한 정보) 를 섞어 가지고 있어야 합니다. 하지만 Muon 은 모든 주파수를 평평하게 만들어버려, AI 가 가진 '깊이'와 '다양성'을 잃게 만듭니다.

2. 해결책: HTMuon (무거운 꼬리를 가진 마법 도구)

이 문제를 해결하기 위해 제안된 것이 HTMuon입니다. 이름의 'HT'는 **'Heavy-Tailed (무거운 꼬리)'**를 의미합니다.

HTMuon 의 핵심 아이디어: "중요한 것은 더 강하게, 덜 중요한 것은 약하게"

• 비유 (등산의 재해석):
  HTMuon 은 등산할 때 힘을 다르게 분배합니다.
  • 중요한 방향 (신호): 등산로가 뚜렷하고 중요한 길이라면, Muon 보다 더 강하게 밀어줍니다.
  • 덜 중요한 방향 (소음): 바람에 흔들리는 나뭇잎 같은 사소한 방향은 힘을 아주 약하게 줍니다.
  • 결과: 이렇게 하면 AI 가 '중요한 정보'에 집중하고 '소음'은 무시하게 되어, 더 깊은 이해 (Heavy-tailed spectrum) 를 얻게 됩니다.

어떻게 작동할까요?
Muon 은 모든 방향의 힘을 '1'로 고정했지만, HTMuon 은 힘을 **0.125 승 (약 1/8 승)**으로 조절합니다.

• 수학적으로: 큰 힘은 조금 줄이고, 작은 힘은 상대적으로 더 줄여서, 전체적인 힘의 분포가 '꼬리가 길고 무거운' 모양이 됩니다.
• 일상적으로: "모두에게 똑같은 점수를 주는 게 아니라, 진짜 잘한 학생 (중요한 정보) 에게는 더 큰 점수를, 엉뚱한 학생 (소음) 에게는 점수를 아예 안 주는" 식으로 학습을 조절합니다.

3. 실험 결과: 실제로 효과가 있을까요?

연구진은 이 방법을 다양한 AI 모델 (LLaMA, GPT-2 등) 과 이미지 인식 모델 (ResNet, ViT) 에 적용해 보았습니다.

• 성능 향상: 기존 최고의 방법 (Muon 포함) 보다 **더 낮은 실수율 (Perplexity)**을 기록했습니다. 예를 들어, LLaMA 모델을 학습시킬 때 실수가 약 1.0 정도 줄었습니다. 이는 AI 가 더 자연스러운 문장을 생성한다는 뜻입니다.
• 이미지 인식: 고양이와 개를 구별하는 작업에서도 기존 방법보다 더 정확하게 분류했습니다.
• 플러그인 기능: HTMuon 은 기존 Muon 을 대체할 수도 있지만, 기존 Muon 위에 얹어서 쓰는 '부스터' 역할도 할 수 있습니다. Muon + HTMuon 조합은 더 강력한 성능을 냅니다.

4. 속도 문제 해결 (빠르게 달리는 HTMuon)

일반적으로 이런 정교한 계산은 시간이 많이 걸립니다. 하지만 연구진은 두 가지 방법을 개발하여 속도를 높였습니다.

1. 간격 조정: 매번 계산하지 않고, 몇 번의 학습을 거친 후 한 번만 계산합니다. (예: 5 번 중 1 번만 정밀 계산)
2. 수치적 근사: 정확한 계산 대신, 아주 빠르고 정확한 근사치를 사용합니다.

이렇게 하면 성능은 유지하면서 계산 비용은 Muon 과 비슷하게 줄였습니다.

5. 결론: 왜 이 연구가 중요한가요?

이 논문은 **"AI 가 배우는 방식에 '다양성'과 '중요도'를 더하는 것"**이 얼마나 중요한지 보여줍니다.

• 기존의 Muon: "모든 길을 똑같이 걸어라." (빠르지만 가끔 길을 잃을 수 있음)
• 새로운 HTMuon: "중요한 길은 확실히 가고, 사소한 길은 가볍게 지나가라." (더 깊고 정확한 학습)

이 방법은 앞으로 더 크고 똑똑한 AI 를 만들 때, 더 적은 비용으로 더 좋은 성능을 내는 데 큰 기여를 할 것으로 기대됩니다. 마치 요리를 할 때, 모든 재료를 같은 양으로 넣는 게 아니라, 핵심 향신료는 더 많이 넣고 잡내는 덜 넣어서 훨씬 맛있는 요리를 만드는 것과 같습니다.

기술 요약

1. 문제 제기 (Problem Statement)

최근 대규모 언어 모델 (LLM) 학습에서 Muon 옵티마이저는 파라미터 간의 기하학적 상호의존성 (interdependencies) 을 효과적으로 포착하여 AdamW 와 같은 기존 벡터 기반 옵티마이저보다 우수한 성능을 보여주었습니다. Muon 은 모멘텀 행렬에 대한 직교화 (orthogonalization) 업데이트 규칙을 사용하여 모든 특이값 (singular values) 을 1 로 설정합니다.

그러나 이 논문은 Muon 의 이러한 직교화 업데이트가 다음과 같은 두 가지 근본적인 한계를 가지고 있다고 주장합니다.

1. 노이즈 우세 방향의 과대평가: 모든 특이값을 1 로 균일하게 설정함으로써, 작은 특이값에 해당하는 방향 (일반적으로 노이즈가 지배적인 방향) 에도 동일한 가중치를 부여합니다. 이는 학습 후기 단계에서 모델이 노이즈에 민감하게 반응하게 만들어 성능을 저하시킬 수 있습니다.
2. Heavy-Tailed Spectrum 의 억제: 잘 훈련된 신경망은 가중치 행렬의 고유값 분포 (ESD, Empirical Spectral Density) 가 Heavy-Tailed (무거운 꼬리) 특성을 보이는 경향이 있으며, 이는 모델의 일반화 성능과 강한 상관관계가 있습니다 (Heavy-Tailed Self-Regularization, HT-SR 이론). Muon 의 균일한 업데이트는 이러한 Heavy-Tailed 특성을 억제하고 Light-Tailed (가벼운 꼬리) 스펙트럼을 생성하여, 궁극적으로 모델이 도달할 수 있는 최종 성능의 상한을 제한합니다.

2. 제안 방법: HTMuon (Methodology)

저자들은 Muon 의 장점을 유지하면서 Heavy-Tailed 특성을 강화하기 위해 HTMuon을 제안합니다.

• 핵심 아이디어: Muon 의 모멘텀 행렬 $M_t$의 특이값을 모두 1 로 만드는 대신, $p$ ($0 < p < 1$) 승을 취하여 변환합니다.
  • 업데이트 공식: $O_t = U_t \Sigma_t^p V_t^\top$
  • 여기서 $U_t \Sigma_t V_t^\top$는 $M_t$의 SVD(특이값 분해) 입니다.
• 매개변수 $p$의 역할:
  • $p=1$: SGDM(벡터 기반) 과 동일해지며 파라미터 간 상호의존성을 무시합니다.
  • $p=0$: Muon 과 동일해지며 모든 특이값이 1 이 되어 Heavy-Tailed 특성이 약화됩니다.
  • $0 < p < 1$(권장값$p=0.125$): Muon 의 행렬 기반 구조를 유지하여 파라미터 간 상호의존성을 포착하면서도, 작은 특이값을 상대적으로 더 작게 만들어 노이즈 방향을 억제하고 Heavy-Tailed 업데이트를 유도합니다.
• 이론적 배경: HTMuon 은 Schatten-$q$ 노름 제약 하의 Steepest Descent와 동치임을 이론적으로 증명했습니다. 이는 Muon 이 Schatten-$\infty$ 노름 제약 하의 Steepest Descent 라는 기존 결과를 일반화한 것입니다.

3. 주요 기여 (Key Contributions)

1. Muon 의 한계 규명 및 HT-SR 이론 적용: Muon 의 직교화 업데이트가 Heavy-Tailed 스펙트럼 형성을 억제하고 노이즈 우세 방향을 강조하여 일반화 성능을 제한한다는 것을 실험적으로 증명했습니다.
2. HTMuon 옵티마이저 제안: 파라미터 간 상호의존성 모델링 능력을 유지하면서 Heavy-Tailed 업데이트를 생성하는 새로운 옵티마이저를 설계했습니다.
3. 광범위한 실험적 검증:
   • LLM 프리트레이닝: LLaMA (60M~1B), GPT-2 모델 학습에서 Adam, AdamW, Muon 및 최신 변형 옵티마이저 (NorMuon, AdaMuon, COSMOS 등) 대비 일관된 성능 향상을 보였습니다.
     • 예: LLaMA-135M(C4 데이터셋) 에서 Muon 대비 Perplexity 를 0.98만큼 감소시켰습니다.
   • 이미지 분류: CIFAR-10/100 및 ImageNet-1K 에서 ResNet 및 ViT 아키텍처를 사용하여 기존 옵티마이저보다 높은 정확도를 달성했습니다.
4. 효율적인 구현 및 가속화: SVD 연산의 오버헤드를 줄이기 위해 두 가지 가속화 전략을 제안했습니다.
   • 간격 업데이트 (Interval-based): 모든 스텝이 아닌 특정 간격 (예: 5 스텝) 마다 HTMuon 을 적용하고 나머지는 Muon 을 사용하는 방식.
   • 수치적 반복법 (HTMuon NS): SVD 대신 Newton-Schulz 반복법을 사용하여 행렬의 분수 거듭제곱을 근사하는 방식. 이는 계산 비용을 크게 줄이면서도 Muon 대비 우수한 성능을 유지합니다.
5. 이론적 분석: 부드러운 비볼록 (smooth non-convex) 환경에서 HTMuon 이 Muon 및 SGDM 과 동일한 샘플 복잡도 상한 ($O(\epsilon^{-4})$) 을 가진다는 수렴 분석을 제공했습니다.

4. 실험 결과 (Results)

• 성능 향상:
  • LLaMA-60M (C4): Muon 대비 PPL 0.92 감소 (28.80 $\to$ 27.88).
  • LLaMA-135M (C4): Muon 대비 PPL 0.98 감소 (22.23 $\to$ 21.25).
  • LLaMA-1B (C4): Muon 대비 PPL 0.16 감소 (14.33 $\to$ 14.17).
  • 이미지 분류: CIFAR-100 에서 ResNet50 기준 Muon 대비 정확도 0.31%p 향상.
• Heavy-Tailed 특성 확인: HTMuon 으로 학습된 모델의 가중치 행렬은 Muon 으로 학습된 모델보다 더 낮은 Power Law 지수 ($\alpha$) 를 가지며, 이는 더 무거운 꼬리 (Heavy-Tailed) 분포를 의미합니다. 이는 HT-SR 이론과 일치하며 더 나은 일반화 성능과 연결됩니다.
• 하류 작업 (Downstream Tasks): LLaMA-1B 모델을 7 가지 상식 추론 벤치마크에서 평가한 결과, HTMuon 은 평균 점수에서 Muon 보다 1.05 포인트 높은 성능을 보였습니다.
• 효율성: HTMuon NS(간격 5) 는 Muon 대비 약 6~11% 의 추가 오버헤드만 발생시키면서도 성능은 우월했습니다.

5. 의의 및 결론 (Significance)

이 논문은 옵티마이저 설계에 있어 Heavy-Tailed Self-Regularization (HT-SR) 이론의 중요성을 부각시켰습니다. 단순히 모멘텀을 정규화하는 것을 넘어, 업데이트 스펙트럼의 꼬리 분포를 제어함으로써 모델이 노이즈에 덜 민감하고 더 강력한 일반화 능력을 갖도록 유도할 수 있음을 증명했습니다.

HTMuon 은 기존 Muon 의 강점 (파라미터 간 기하학적 관계 포착) 을 잃지 않으면서 Heavy-Tailed 특성을 도입하여, LLM 프리트레이닝 및 이미지 분류 등 다양한 작업에서 State-of-the-Art 성능을 달성했습니다. 또한, 제안된 가속화 기법과 플러그인 모듈로서의 호환성 (기존 Muon 변형과 결합 가능) 은 실제 대규모 모델 학습 환경에서의 실용성을 높여주며, 향후 옵티마이저 연구에 새로운 방향성을 제시합니다.