TrasMuon: Trust-Region Adaptive Scaling for Orthogonalized Momentum Optimizers

이 논문은 뉴턴-슈어츠 반복을 통한 직교화를 수행하는 Muon 옵티마이저의 크기 불안정성 문제를 해결하기 위해, 전역 RMS 보정과 에너지 기반 신뢰 영역 클리핑을 결합한 TrasMuon 을 제안하며, 이는 웜업 단계 없이도 기존 베이스라인보다 빠른 수렴과 향상된 안정성을 입증합니다.

핵심

1. 문제 상황: "빠르지만 위험한 스포츠카"

AI 모델을 학습시키는 과정은 미지의 땅을 빠르게 달리는 스포츠카를 운전하는 것과 비슷합니다.

• 기존의 방법 (Muon): 최근 등장한 'Muon'이라는 도구는 차의 방향을 아주 정교하게 조절합니다. 마치 네비게이션이 도로의 곡선을 완벽하게 예측해서 차가 미끄러지지 않고 빠르게 코너를 도는 것처럼요. 그래서 학습 속도가 매우 빠릅니다.
• 하지만 문제점이 있습니다: 이 도구는 방향만 조절할 뿐, **차의 속도 (크기)**를 잘 조절하지 못합니다. 갑자기 큰 돌 (데이터의 이상치) 이 튀어나오면, 차가 너무 빠르게 달려서 심하게 흔들리거나 심지어 전복될 수 있습니다. 이를 기술 용어로 '학습이 불안정해진다'고 말합니다.

2. TrasMuon 의 등장: "스마트한 속도 제한 장치"

저자들은 이 문제를 해결하기 위해 TrasMuon을 만들었습니다. 이 도구는 기존 'Muon'의 빠른 방향 전환 능력을 유지하면서, 두 가지 새로운 안전 장치를 추가했습니다.

① 전역 속도 조절기 (Global RMS Calibration)

• 비유: 차가 달리는 전체 도로의 평균 속도를 측정해서, 너무 빠르거나 느리지 않게 전체적인 속도를 맞춰주는 장치입니다.
• 효과: AI 가 학습하는 동안, 각 층 (Layer) 마다 학습 속도가 들쑥날쑥하지 않고 일정하게 유지되도록 도와줍니다.

② 신뢰 구역 (Trust Region) 과 에너지 차단기

• 비유: 도로에 갑자기 **한쪽 차선에만 엄청난 폭포수 (에너지 폭발)**가 쏟아진다고 상상해 보세요. 기존 방식은 그 폭포수를 피할 수 없어 차가 뒤집힙니다.
• TrasMuon 의 해결책: 이 도구는 "어떤 차선에서 물이 너무 많이 쏟아지고 있나?"를 실시간으로 감지합니다. 그리고 그 특정 차선으로만 가는 물의 양을 줄여주는 (Damping) 장치를 작동시킵니다.
• 핵심: 다른 차선은 그대로 빠르게 달리게 하면서, 위험한 차선만 살짝 브레이크를 밟는 것입니다. 이를 **'신뢰 구역 (Trust Region)'**이라고 부릅니다. 위험한 영역에 들어가지 못하게 경계를 설정한 셈입니다.

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

이 논문의 실험 결과에 따르면 TrasMuon 은 다음과 같은 장점이 있습니다:

1. 더 빠른 학습: 처음 학습을 시작할 때 (Warm-up 단계 없이도) 훨씬 빠르게 안정화됩니다. 마치 스포츠카가 시동을 걸자마자 바로 최고 속도로 달리는 것처럼요.
2. 더 튼튼한 안정성: 데이터에 갑자기 이상한 값 (Outlier) 이 섞여도 AI 가 넘어지지 않고 견딜 수 있습니다.
3. 편의성: 기존에는 학습 속도를 조절하는 '학습률 (Learning Rate)'을 아주 정교하게 맞춰주지 않으면 AI 가 망가졌는데, TrasMuon 은 그런 세밀한 조절 없이도 잘 작동합니다.

4. 요약: "방향은 Muon, 안전은 TrasMuon"

• 기존 (Muon): 방향은 잘 잡지만, 돌발 상황에 약함.
• TrasMuon: 방향은 Muon 만큼 잘 잡으면서, 위험한 상황 (에너지 폭발) 이 감지되면 그 부분만 자동으로 브레이크를 밟아주는 지능형 시스템.

결론적으로 TrasMuon 은 AI 를 더 빠르고, 더 안전하게, 그리고 더 편하게 훈련시킬 수 있게 해주는 **'스마트한 안전장치'**가 달린 새로운 운전 기술이라고 할 수 있습니다.

기술 요약

1. 문제 제기 (Problem)

현대 대규모 기초 모델 (Foundation Models) 의 학습에서 옵티마이저 선택은 수렴 속도, 안정성, 계산 비용에 결정적인 영향을 미칩니다. 최근 Muon 스타일의 옵티마이저는 뉴턴 - 슈르츠 (Newton-Schulz, NS) 반복을 통해 업데이트 방향을 직교화 (Orthogonalization) 하여, Adam 계열 방법보다 우수한 업데이트 기하학을 제공합니다.

그러나 Muon 스타일 옵티마이저에는 다음과 같은 근본적인 한계가 존재합니다:

• 크기 정보의 손실: 직교화 과정이 업데이트의 방향 (기하학) 만을 제어할 뿐, 크기 (Magnitude) 정보를 버립니다. 이로 인해 학습률 (Step-size) 하이퍼파라미터에 매우 민감해집니다.
• 고에너지 버스트 (High-energy Bursts) 에 대한 취약성: 실제 학습 신호는 두꺼운 꼬리 (Heavy-tailed) 를 가지며, 특정 특징 축 (Feature axes) 에 에너지가 집중되는 순간적인 버스트가 발생합니다. 크기 보정이 부재할 경우, 이러한 버스트가 손실 급증 (Loss Spikes) 을 유발하고 안정적인 학습률 영역을 좁힙니다.
• 워업 (Warmup) 의존성: 크기 보정이 없으면 학습 초기의 불안정성을 막기 위해 긴 워업 기간과 세밀한 스케줄 조정이 필수적입니다.

2. 방법론 (Methodology)

저자들은 TrasMuon (Trust-Region Adaptive Scaling for Muon) 을 제안합니다. 이는 Muon 스타일의 직교화된 방향을 유지하면서, 크기 안정성을 확보하기 위해 전역 RMS 보정 (Global RMS Calibration) 과 신뢰 영역 기반 적응형 스케일링 (Trust-Region Adaptive Scaling) 을 결합한 알고리즘입니다.

핵심 구성 요소

1. 구조화된 혼합 인자 (Structured Mixing Factor):

   • Muon 과 동일하게 뉴턴 - 슈르츠 (NS) 반복을 사용하여 모멘텀의 극성 인자 (Polar Factor) 를 근사합니다.
   • 이를 통해 업데이트 방향이 회전 불변성 (Rotation-invariant) 을 가지며, 전역 특징 혼합 (Global feature mixing) 을 촉진합니다.
   • 행렬 업데이트는 $\Delta W_t = -\hat{\eta}_t O^{base}_t \text{diag}(c_t)$ 형태로 분해됩니다. 여기서 $O^{base}_t$ 는 직교화된 방향입니다.

2. 전역 RMS 보정 (Global RMS Calibration):

   • 학습률 $\hat{\eta}_t$ 를 층의 형태와 무관하게 일정한 RMS 범위를 갖도록 보정합니다.
   • $\hat{\eta}_t = \eta \cdot \frac{\sqrt{d_{out}d_{in}}}{\|O^{base}_t\|_F + \epsilon}$ 공식을 사용하여, 계층별 및 시간별 크기 변동을 줄이고 학습률 민감도를 낮춥니다.

3. 신뢰 영역 기반 특징별 감쇠 (Trust-Region Feature-wise Damping):

   • 문제: 전역 보정은 전체 크기는 제어하지만, 특정 열 (Column/Feature) 에 에너지가 집중되는 버스트를 막지 못합니다.
   • 해결: 각 열의 에너지 $E_j$ 와 robust 기준치 $E_{ref}$ (중앙값 기반) 의 비율 $r_j = E_j / E_{ref}$ 를 계산합니다.
   • 클리핑 (Clipping): 비율이 임계값을 초과하는 열에 대해 감쇠 계수 $c_j \in [c_{min}, 1]$ 을 적용합니다.
     • $c_j = \text{clip}\left(\frac{1}{1 + \alpha \log(1 + r_j)}, c_{min}, 1\right)$
   • 이는 신뢰 영역 (Trust Region) 개념으로, 에너지가 집중된 "폭발적인" 특징 축만 선택적으로 억제하여 손실 급증을 방지합니다.

4. 시간적 안정화 (Temporal Smoothing & Schedule-Free Averaging):

   • 감쇠 신호가 노이즈에 민감하지 않도록 지수 이동 평균 (EMA) 을 적용합니다.
   • 또한, Schedule-Free 방식 (Effective-time weighting) 을 도입하여 워업 기간이나 학습 스케줄 변화에 따른 편향을 줄이고, 장기적인 평균을 통해 안정성을 확보합니다.

3. 주요 기여 (Key Contributions)

1. 알고리즘 제안: Muon 스타일의 직교화 혼합과 전역 RMS 보정, 그리고 특징별 에너지 신뢰 영역 클리핑을 통합한 TrasMuon 알고리즘을 개발했습니다.
2. 기초 모델 학습에서의 입증:
   • 워업 유무 관계없이 빠른 초기 수렴과 높은 안정성을 보입니다.
   • 특히 워업이 없는 설정에서도 TrasMuon 은 AdamW 나 기존 Muon 보다 훨씬 빠르게 손실을 감소시키고, 손실 급증을 효과적으로 억제합니다.
3. 강건성 (Robustness):
   • 두꺼운 꼬리 분포 (Heavy-tailed) 와 구조화된 비정상성 (Structured non-stationarity) 하에서도 손실 급증을 줄이고 더 나은 최종 성능을 달성합니다.
   • PINN(물리 정보 신경망) 과 같은 비정상적인 샘플링 환경에서도 뛰어난 안정성을 입증했습니다.

4. 실험 결과 (Results)

• 언어 모델 사전 학습 (Qwen3-0.6B, GPT-2):
  • 워업 포함: TrasMuon 은 AdamW 보다 2.35 배, Muon 보다 1.75 배 빠르게 목표 손실 (7.0) 에 도달했습니다.
  • 워업 제거: 워업 없이도 TrasMuon 은 AdamW 보다 6.21 배, Muon 보다 1.73 배 빠르게 수렴하며, 초기 단계의 손실 진동을 크게 줄였습니다.
• 비전 변환기 (ViT-Base on ImageNet-100):
  • AdamW, Muon, NorMuon 과 비교하여 TrasMuon 이 가장 빠른 손실 감소와 가장 높은 검증 정확도 (77.47%) 를 기록했습니다.
  • 시드 간 변동성 (Variance) 이 가장 낮았습니다.
• PINN 벤치마크 (Helmholtz 방정식):
  • ROI(관심 영역) 샘플링으로 인한 비정상적인 분포 변화를 가하는 스트레스 테스트에서, TrasMuon 은 기존 Muon 과 유사한 수렴 속도를 유지하면서도 극단적인 변동성을 줄이고 최종 해의 정확도를 높였습니다.
• 메커니즘 연구 (Toy Problem):
  • 특정 열에 인위적인 버스트를 주입했을 때, TrasMuon 은 해당 열의 에너지 비율 증가를 감지하여 즉시 감쇠를 적용하고 손실 급증을 억제하는 것을 확인했습니다.
  • 특징 축의 의미가 무너지는 경우 (열의 의미가 섞일 때) 는 성능 향상이 사라지므로, 제안된 메커니즘이 특징 기반 (Feature-wise) 으로 작동함을 증명했습니다.

5. 의의 및 결론 (Significance)

TrasMuon 은 행렬 구조화된 업데이트 (Muon 계열) 와 예측 가능한 크기 제어 (적응형 스케일링) 를 성공적으로 결합했습니다.

• 실용성: 대규모 모델 사전 학습에서 워업 기간과 학습 스케줄 조정에 대한 의존도를 줄여, "Drop-in" 가능한 실용적인 옵티마이저가 됩니다.
• 안정성: 고에너지 버스트와 두꺼운 꼬리 노이즈가 존재하는 현실적인 학습 환경에서, 손실 급증을 방지하고 학습의 안정성을 크게 향상시킵니다.
• 효율성: 복잡한 곡률 추정 (Curvature estimation) 없이도 뉴턴 - 슈르츠 반복을 통해 효율적인 직교화를 수행하며, 크기 보정을 통해 최적화 효율을 극대화합니다.

결론적으로 TrasMuon 은 현대 대규모 모델 학습에서 직교화 옵티마이저의 한계를 극복하고, 더 빠르고 안정적인 학습을 가능하게 하는 차세대 옵티마이저로 평가됩니다.