기술 요약: 잔차 연결(Residual Connection)로서의 Muon

문제 정의

Muon은 CNN부터 대규모 언어 모델(LLM)에 이르는 다양한 아키텍처 전반에서 최첨단 성능을 달este하며, 대규모 신경망 학습에 매우 효과적인 옵티마이저로 부상했습니다. 이러한 경험적 성공에도 불구하고, Muon의 효능에 대한 기존 설명들은 스펙트럼 노름(spectral-norm) 최급 강하, 비유클리드 신뢰 영역 최적화, 스펙트럼 노름 볼(spectral-norm ball) 상의 프랭크-울프(Frank–Wolfe) 방법, 롱테일 학습(long-tail learning), 곡률 분석 등 다양하고 정교한 수학적 프레임워크에 의존하고 있습니다.

본 논문은 광범위한 수학적 훈련을 요구하지 않는, 더 직관적이고 기계적인 설명을 제공할 필요성을 다룹니다. 구체적으로, 단순히 즉각적인 그래디언트 충실도(gradient fidelity)나 국소 손실 감소에만 집중하는 것이 아니라, 하위 레이어를 위한 표현의 사용성(representation usability) 관점에서 최적화 역학을 재구성함으로써 왜 Muon이 잘 작동하는지를 설명하고자 합니다.

방법론

저자들은 새로운 해석을 제안합니다: Muon은 암시적 잔차 연결(implicit residual connection)로서 기능한다.

이론적 프레임워크

1. 직교화로서의 잔차성(Orthogonalization as Residuality): 표준 잔차 네트워크에서 업데이트는 명시적입니다: $x_{l+1} = x_l + \text{Block}(x_l)$, 또는 가중치 공간에서는 $W \leftarrow W + I$ (또는 직교 행렬 $R$)입니다. Muon은 표준 업데이트 $W \leftarrow W + \Delta W$를 직교화된 업데이트 $W \leftarrow W + \text{Orth}(\Delta W)$로 대체합니다.
2. 트레이드오프(The Trade-off): 본 논문은 그래디언트 충실도(즉각적인 손실 감소를 위해 원시 그래디언트에 근접함)와 표현 보존(하위 레이어가 활용하기 쉬운 표현 구조를 유지함) 사이의 트레이드오프를 상정합니다. 업데이트를 직교화하는 것은 즉각적인 그래디언트 충실도를 희생하지만, 상위 표현의 사용성을 보존하는 잔차 유사 변환을 암시적으로 주입합니다.

실험적 검증

이 가설을 테스트하기 위해, 저자들은 복잡한 비선형 심층 네트워크의 복잡성을 피하고 메커니즘을 고립시키기 위해 선형 모델을 사용하는 통제된 실험을 수행했습니다.

1. 2단계 선형 실험:

   • 1단계 (국소 최적화): 두 개의 층으로 구성된 선형 네트워크($\hat{y} = W_2 W_1 x$)가 타겟 변환 $A$를 맞추도록 학습됩니다. $W_2$는 고정되며, $W_1$은 로컬 타겟 $W_1^*$를 향해 최적화됩니다.
     • 가설: SGD는 로컬 목적 함수를 위한 최급 강하 방향을 따르는 반면, Muon은 업데이트를 직교화하므로 여기서 SGD가 더 우수할 것입니다.
   • 2단계 (하위 적응): 학습된 $W_1$을 고정하고, $W_2$를 재초기화하여 $\|W_2 W_1 - A\|_F$를 최소화하도록 SGD를 사용하여 최적화합니다.
     • 가설: 만약 Muon의 "암시적 잔차" 효과가 성립한다면, 1단계에서 Muon에 의해 학습된 $W_1$은 로컬 타겟으로부터는 더 멀리 떨어져 있음에도 불구하고, $W_2$가 적응하기 더 쉬운 상태가 될 것입니다.

2. $\tau$-스케줄 실험:

   • 엔드 투 엔드(end-to-end) 학습을 시뮬레이션하기 위해, 저자들은 레이어를 교대하는 길이 $\tau$의 세그먼트(예: $W_1$에 대해 $\tau$ 스텝, 그 다음 $W_2$에 대해 $\tau$ 스텝)로 업데이트하는 스케줄을 도입했습니다.
   • 이 설정은 "로컬 비용"( $W_1$ 업데이트 중)이 "하위 이득"( $W_2$ 업데이트 중)과 어떻게 상호작용하는지 연속적인 학습 루프 내에서 관찰할 수 있게 합니다.

주요 결과

2단계 선형 실험

• 1단계: 예측대로, 600 스텝 후 SGD가 Muon($0.0447$)보다 더 낮은 로컬 재구성 오차($0.0312$)를 달성했습니다. 이는 Muon이 그래디언트 충실도를 희생함을 확인시켜 줍니다.
• 2단계: 고정된 $W_1$에 대해 $W_2$를 최적화할 때, Muon으로 학습된 모델은 SGD로 학습된 모델(13,436 스텝)보다 타겟 오차 임계값에 도달하는 데 더 적은 스텝 수(13,321 스텝)가 필요했습니다.
• 해석: Muon에 의해 학습된 표현은 로컬 타겟에는 최적이 아닐지라도, 더 "평탄한" 단일 스펙트럼(높은 스펙트럼 평탄도)을 가져 하위 레이어의 최적화 문제를 더 잘 조건화(well-conditioned)된 상태로 만들었습니다.

$\tau$-스케줄 실험

• Muon은 완전한 공동 업데이트($\tau=0$)와 교대 업데이트($\tau > 0$)를 포함한 모든 테스트된 스케줄에서 SGD보다 일관되게 더 빠르게 타겟 합성 손실(composition loss)에 도달했습니다.
• 역학 분석:
  • "블루 세그먼트"($W_1$만 업데이트하는 동안) 동안 SGD는 종종 초기 손실을 더 빠르게 줄였으며, 이는 Muon의 그래디언트 충실도 비용을 반영합니다.
  • "레드 세그먼트"($W_2$를 업데이트하는 동안) 동안, Muon이 학습한 궤적은 더 빠른 수렴을 가능하게 했습니다.
  • 결정적으로, Muon의 궤적에서 개선된 $W_2$는 후속 $W_1$ 서브 문제를 해결하기 더 쉽게 만들어, 초기 로컬 지연을 보상하는 피드백 루프를 생성했습니다.

주요 기여

1. 기계적 해석: 본 논문은 Muon의 성공에 대한 단순하고 직관적인 설명을 제공합니다: 그것은 암시적 잔차 연결로서 기능합니다. 이는 그래디언트 강하라는 즉각적인 필요성과 표현 보존이라는 장기적인 필요 사이의 균형을 맞춥니다.
2. 새로운 설계 관점: 이는 옵티마이저 설계를 레이어별 서브 문제에서의 즉각적인 손실 감소만으로 판단하는 것에서, 업데이트가 하위 레이어를 위한 표현을 어떻게 형성하는지를 고려하는 것으로 전환시킵니다.
3. 경험적 증거: 통제된 선형 실험을 통해, 저자들은 로컬 최적화 속도를 희생하는 것이 하위 최적화 문제의 조건수(condition number)를 개선함으로써 엔드 투 엔드 수렴 속도를 높일 수 있음을 입증했습니다.

의의 및 주장

본 논문은 이러한 해석이 기존의 설명(스펙트럼, 곡률, 신뢰 영역)을 대체하는 것이 아니라 보완한다고 주장합니다. Muon을 암시적 잔차 연결로 프레이밍함으로써, 저자들은 다음과 같이 제안합니다:

• 옵티마이저는 로컬 및 하위 목적을 모두 균형 있게 고려해야 합니다: 효과적인 옵티마이저는 레이어별 서브 문제에 대한 즉각적인 손실 감소만으로 평가될 필요가 없습니다.
• 맥락이 중요합니다: Muon과 같은 동작(로컬 충실도를 희생하여 하위 사용성을 높이는 것)의 이점은 하위 레이어가 더 잘 조건화된 표현을 활용할 수 있는 충분한 용량과 학습 시간을 가질 때, 특히 결합된 변환의 깊은 체인에서 가장 가치 있습니다.
• 미래 방향: 이 관점은 레이어의 스펙트럼, 조건수 또는 하위 진행 상황에 따라 직교화 또는 스펙트럼 평탄화의 강도를 조절하는 적응형 옵티마이저 설계를 시사합니다.

저자들은 자신의 실험이 의도적으로 통제되었으며 선형적이라는 점을 인정하며 겸손한 어조를 유지합니다. 이들이 모든 현대적 비선형 네트워크에서 이 메커니즘이 작동함을 증명했다고 주장하는 것이 아니라, Muon의 관찰된 경험적 이점에 대한 타당하고 직관적인 메커니즘을 제시하고 있다고 주장합니다.