The Spectral Dynamics and Noise Geometry of Muon

기술 요약: Muon의 스펙트럼 역학 및 노이즈 기하학

문제 정의
본 논문은 가중치를 그래디언트 자체가 아닌 그래디언트 행렬의 극값(polar factor)을 사용하여 업데이트하는 Muon 옵티마이저에 의해 유도되는 암묵적 정규화 편향(implicit regularization bias)을 조사한다. $G = U\Sigma V^\top$인 그래디언트에 대해, Muon은 $UV^\top$를 통해 업데이트하며, 이는 특잇값(singular values)의 상대적 크기는 버리되 활성 특이 방향(active singular directions)은 보존하는 방식이다. Muon은 트랜스포머 사전 학습에서 경험적인 성공을 보여주었으나, 그 이론적 선택 메커니즘은 여전히 불분명하다. 두 가지 경쟁하는 가설이 다뤄진다:

1. Muon은 단순히 행렬-그래디언트 정규화(matrix-gradient normalization)의 한 형태인가?
2. 스펙트럼-노름 기하학(spectral-norm geometry)이 저계수 해(low-rank solutions, 핵-노름 최소화)로의 편향을 의미하는가?

저자들은 두 가설 모두 불완전하다고 주장한다. 대신, 이들은 Muon이 모든 활성 특이 방향에 걸쳐 균등한 진폭을 할당하는 **평탄-업데이트-스펙트럼 편향(flat-update-spectrum bias)**을 유도한다고 제안한다. 본 논문은 특히 옵티마이저가 제로-로스 해들의 아핀 다양체(affine manifold)로부터 유한한 보간 함수(interpolant)를 선택해야 하는 과결정 상태의 회귀(underdetermined regression) 맥락에서, 이 편향을 이론적 및 경험적으로 규명하고자 한다.

방법론
본 연구는 연속 시간 역학계 분석, 변분법(variational calculus), 그리고 통제된 경험적 실험의 조합을 사용한다.

1. 이론적 프레임워크:

   • 투영된 폴라 플로우(Projected Polar Flow): 저자들은 $\dot{W} = -\eta P_\perp P(W)$인 연속 시간 흐름을 분석한다. 여기서 $P(W)$는 $W$의 극값이며, $P_\perp$는 보간 다양체 $\mathcal{M} = \{W : XW = Y\}$의 접공간(tangent space)으로의 직교 투영이다.
   • 공유-프레임(Shared-Frame, SF) 가정: 그래디언트와 가중치 행렬이 동일한 특이 프레임($U_G = U_W, V_G = V_W$)을 공유한다는 구조적 조건을 도입하여 분석을 단순화한다. 자가-폴라 플로우(self-polar flow)에 대해서는 이 가정 없이 정리들을 증명하였으며, SF 가정은 실제 그래디언트 기반의 Muon 플로우와의 간극을 메운다.
   • 범함수(Functionals): 스펙트럼 형상을 추적하기 위해 쌍별 스펙트럼 범함수 $R_{pw}(W) = -\sum_{i \neq j} \log(\sigma_i + \sigma_j)$와 그 척도 불변 버전인 $\tilde{R}_{pw}$를 활용한다.
   • 노이즈 기하학: 프레셰 미분(Fréchet derivative)을 분석하고, 일단계 선형화된 신호 대 잡음 교차(one-step linearized signal-to-noise crossover)를 사용하여 임계 배치 크기를 도출한다.

2. 경험적 검증:

   • 행렬 센싱(Matrix Sensing): 과결정 상태의 선형 회귀 작업에서 Muon을 핵-노름 최소화 도구(CVXPY로 해결) 및 경사 하강법(gradient descent)과 비교하는 통제된 실험을 수행한다.
   • 트랜스포머 사전 학습: 스펙트럼 역학(stable rank, spectral entropy)을 관찰하기 위해 NanoGPT (124M)의 소규모 사전 학습과 이에 대응하는 CIFAR-10용 ViT 대조군을 사용한다.
   • 어블레이션(Ablations): 가중치 감쇠 결합(coupled vs. decoupled) 및 공유-프레임 불일치에 대한 테스트를 통해 견고성을 검증한다.

주요 기여

1. 일단계 스펙트럼-편향 정리:
   명시적인 정렬 가정 하에, 저자들은 폴라 업데이트가 그래디언트 특이 방향을 따르되 현재 가중치의 스펙트럼은 무시하는 유계된 업데이트들 중 유일한 일단계 엔트로피 최대화 선택임을 증명한다. 이들은 **스펙트럼 평탄화 기준(spectral flattening criterion)**을 확립한다: 정규화된 스펙트럼은 측정 기하학이 가장 큰 특잇값에 대응하는 방향을 우선적으로 보존할 때( $\alpha_i q_i$를 포함하는 부호 기준으로 정량화됨)만 평탄해진다.

2. 정확한 특잇값 역학:
   투영된 폴라 플로우에 대한 정확한 역학을 도출하여, 특잇값이 $\dot{\sigma}_i = -\eta \alpha_i$의 속도로 감소함을 보여준다. 여기서 $\alpha_i$는 가중치의 좌측 특이 벡터와 측정 행렬의 영공간(null space) 사이의 정렬에 따라 결정된다. 이는 고정된 프로베니우스 노름 하에서, 스펙트럼 범함수의 유일한 최소화 도구는 평탄한 스펙트럼(모든 비제로 특잇값이 동일함)임을 보여주는 변분적 특성화를 이끈다.

3. 기존 설명과의 분리:

   • 정규화: 노름이 매칭된 경사 하강법은 Muon의 동작을 재현하지 못함을 실험으로 보여준다.
   • 핵-노름: 이론적 및 경험적 결과는 Muon의 극한점(limit points)이 핵-노름 최소화 도구와 일치하지 않음을 보여준다. 핵-노름 최소화는 스펙트럼 집중(저계수)을 선호하는 반면, Muon은 스펙트럼 평탄화(높은 stable rank)를 선호한다.

4. 레짐 의존적 효용성:
   본 논문은 Muon의 평탄-스펙트럼 편향이 보편적으로 우월한 것은 아님을 밝힌다. 소규모 NanoGPT 사전 학습에서 Muon은 stable rank를 보존하고 AdamW에 비해 검증 손실을 개선한다. 그러나 매칭된 소규모 ViT 대조군에서는 순위가 역전되며, 이는 이 편향이 많은 스펙트럼 방향이 활성화되어야 할 때만 유익함을 시사한다.

결과

• 이론적: 정리 1은 투영된 자가-폴라 플로우가 측정 기하학이 특정 부호 기준을 만족할 경우 시스템을 평탄한 스펙트럼으로 몰아넣음을 증명한다. 정리 2는 공유-프레임 가정이 근사적으로만 충족될 때의 편차를 제한한다. 정리 4는 폴라-맵의 노이즈 민감도 $S(\mu) = \sum_{i \neq j} (\sigma_i + \sigma_j)^{-2}$와 연결된 폴라-맵 SGD의 임계 배치 크기 $B_{crit}$ 공식을 제공한다.
• 경험적 (행렬 센싱): Muon은 최소 핵-노름 보간법에 비해 현저히 높은 핵-노름(1.29x–2.02x)과 높은 스펙트럼 엔트로피를 가진 해로 수렴한다. 이는 Muon이 핵-노름을 최소화한다는 가설을 반증한다.
• 경험적 (트랜스포머): NanoGPT에서 Muon은 언어 작업에서 계수를 붕괴시키는 경향이 있는 AdamW에 비해 더 높은 레이어별 스펙트럼 엔트로피와 stable rank를 유지한다. 반대로, ViT 대조군에서는 AdamW가 Muon보다 우수한 성능을 보여, "레짐 의존적" 결론을 뒷받침한다.
• 노이즈 민감도: 도출된 임계 배치 크기 공식은 폴라-맵이 작은 또는 퇴화된 특잇값에 의해 증폭된 노이즈에 얼마나 민감한지에 의해 결정됨을 보여준다. 고정된 프로베니우스 반경에서의 평탄한 스펙트럼은 이러한 민감도를 최소화한다.

의의 및 주장
본 논문은 Muon의 암묵적 편향을 저계수 또는 단순 정규화 효과가 아닌 평탄-업데이트-스펙트럼 메커니즘으로서 최초로 구조적으로 규명했다고 주장한다.

• 제한적 범위: 저자들은 결과가 **레짐 의전적(regime-dependent)**임을 명시적으로 밝힌다. Muon은 보편적인 옵티마이저가 아니다. 그 효용은 많은 방향이 활성화된 상태로 유지되는 것이 유익한 스펙트럼 레짐에서 발생한다.
• 이론적 한계: 정리들은 국소적이고 구조적이며, 단순한 특잇값 간격 및 특정 측정 기하학과 같은 가정에 의존한다. 본 논문은 트랜스포머 레이어를 연구된 과결정 상태의 회귀 모델로 축소하는 것이 검증되지 않은 가정(활성화 랭크에 관한 가정)을 필요로 함을 언급하며, 대규모 LLM 훈련의 모든 시나리오에서 Muon의 성공을 설명한다고 주장하지 않는다.
• 실무적 통찰: 이 연구는 실무자에게 진단 도구를 제공한다: 만약 애플리케이션이 stable rank를 보존하고 넓은 스펙트럼 활동을 필요로 한다면 Muon의 평탄-스펙트럼 편향이 유리하며, 문제가 본질적으로 저계수이거나 조건이 잘 갖춰져 있다면(well-conditioned) 이 편향은 불필요하거나 해로울 수 있다.

논문은 평탄화 기준을 전역적으로 보장하는 특정 측정 기하학을 규명하는 것과, 노이즈 민감도 분석을 직사각형 및 랭크 결핍 사례로 확장하는 것을 향후 과제로 제시하며 마무리된다.