Denoise First, Orthogonalize Later: Understanding Momentum in Muon via Spectral Filtering

기술 요약: "먼저 노이즈를 제거하고, 나중에 직교화하라: 스펙트럼 필터링을 통한 Muon의 모멘텀 이해"

문제 정의
Muon 옵티마이저는 뉴턴-슐츠(Newton–Schulz) 반복법을 사용하여 모멘텀 행렬을 대략적으로 직교화함으로써 대규모 언어 모델(LLM) 학습에서 강력한 경험적 성능을 입증해 왔습니다. 그러나 Muon 내에서 모멘텀의 이론적 역할은 여전히 불분명합니다. 기존의 분석들은 스펙트럼 업데이트만을 고립시켜 연구하기 위해 모멘텀을 제거하거나, 모멘텀을 유지하면서도 그것이 왜 실질적인 이점을 제공하는지에 대해서는 설명하지 못합니다. 또한, 모멘텀에 관한 기존 이론 문헌들은 무거운 모멘텀이 수렴 속도를 저하시키거나 아무런 이점이 없다고 제안하는 경우가 많은데, 이는 모멘텀이 필수적이라는 실제 관찰 결과와 상충됩니다. 아울러, 연산의 특정 순서—모멘텀을 적용한 후 직교화하는 방식(Pre-polar)과 직교화를 먼저 수행한 후 모멘텀을 적용하는 방식(Post-polar)—에 대한 엄밀한 설명이 부족하며, 실제로는 Pre-polar 방식이 Post-polar보다 훨씬 뛰어난 성능을 보입니다.

방법론
저자들은 Muon을 분석하기 위해 구조화된 **신호-및-섭동 그래디언트 모델(signal-plus-perturbation gradient model)**을 제안합니다. 저자들은 시점 $t$에서의 그래디언트 $G_t$가 일관된 저계수(low-rank) 신호 성분 $G^{\text{sig}}_t$와 유계 분산을 가진 평균 0의 직교(BVMZOS) 섭동 성분 $\Xi_t$로 분해된다고 가정합니다.

핵-이론적 기여는 모멘텀 버퍼 $M_t = \beta M_{t-1} + (1-\beta)G_t$를 **스펙트럼 필터(spectral filter)**로 분석하는 것입니다. 저자들은 가중치 업데이트 $\Delta W_t$에 대한 세 가지 파이프라인 변형을 조사합니다:

1. Pre-polar (표준 Muon): $\Delta W_t \propto O(M_t)$. 즉, 모멘텀을 먼저 적용한 후 극값 인자(polar factor) $O(\cdot)$를 적용합니다.
2. Post-polar (Orthogonal-SGDM): $\Delta W_t \propto \tilde{M}_t$. 즉, 원시 그래디언트에 먼저 극값 인자를 적용한 후 모멘텀을 누적합니다.
3. Polar-only: $\Delta W_t \propto O(G_t)$. 모멘텀이 없는 베이스라인입니다.

분석은 유효 윈도우 크기 $T = 1/(1-\beta)$에 의존하며, 그래디언트 행렬의 특이 부공간(singular subspace) 섭동을 정량화하기 위해 위든(Wedin)의 $\sin \Theta$ 정리와 같은 랜덤 행렬 이론 도구를 활용합니다.

주요 기여

1. 필터로서의 모멘텀 (정리 1 및 따름정리 1):
   저자들은 그래디언트 모델 하에서 모멘텀이 섭동(노이즈)을 억제하는 동시에 일관된 신호를 보존하는 저역 통과 필터(low-pass filter) 역할을 함을 증명합니다. 구체적으로, 유효 윈도우 크기 $T$가 증가함에 따라, 꼬리 부분의 특이값(노이즈에 의해 지배됨)은 $O(T^{-1/4})$의 속도로 감소하는 반면, 선행 특이값(신호)은 0에서 멀리 떨어진 상태를 유지합니다. 이는 신호와 노이즈 성분 사이에 확장된 **스펙트럼 갭(spectral gap)**을 생성합니다. 결과적으로, 모멘텀 버퍼 $M_t$의 상위-$r$ 특이 부공간은 원시 그래디언트의 부공간보다 실제 신호 부공간과 더 신뢰성 있게 정렬됩니다.

2. 비가환성 및 파이프라인 순서 (정리 2):
   본 논문은 모멘텀과 직교화가 교환 법칙이 성립하지 않음을 입증합니다.

   • Pre-polar의 우위: 모멘텀이 먼저 적용되면 그래디언트 스트림의 노이즈를 제거하여, 이후의 극값 인자가 안정된 특이 부공간 구조를 가진 행렬에 작용할 수 있게 합니다. 저자들은 Pre-polar가 Post-polar 및 Polar-only 베이스라인보다 더 높은 기대 신호 정렬도를 달리는 것을 증명합니다.
   • 대안의 실패: Post-polar의 경우, 노이즈가 섞인 그래디언트를 먼저 직교화하면 신호와 노이즈 방향 사이의 진폭 차이가 "지워집니다". 극값 인자는 모든 특이값을 1로 대체하므로, 이후의 모멘텀 단계에서 신호와 노이즘을 구분하는 것이 불가능해집니다. 이론적으로 Post-polar와 Polar-only 파이프라인은 특히 낮은 신호 대 잡음비(SNR) 환경에서 신호 방향을 효과적으로 복구하는 데 실패합니다.

3. 낮은 SNR에서의 정량적 분리 (정리 3):
   랭크-1 스파이크 가우시안 모델(rank-1 spiked Gaussian model)에서, 저자들은 낮은 SNR 환경에서 Polar-only 파이프라인의 신호 정렬도가 행렬 차원 $m \to \infty$에 따라 소멸함을 보여줍니다. 반면, Pre-polar 파이프라인은 모멘텀 윈도우 $T$가 충분히 큰 경우 타겟 신호를 계속해서 복구해냅니다.

결과
이론적 발견은 광범 ***실험을 통해 검증되었습니다***:

• 스펙트럼 갭 시각화: NanoGPT에 대한 정적 프로브(stationary probes)를 사용하여, 모멘텀 계수 $\beta$를 높일수록 모멘텀 버퍼의 선행 특이값과 꼬리 특이값 사이의 간격이 넓어진다는 것을 보여주며, 이는 정리 1과 일치합니다.
• 부공간 신뢰도: 부공간 정렬 오차의 경험적 측정 결과, $T$가 증가함에 따라 모멘텀 버퍼의 특이 부공간과 실제 신호 부공간 사이의 편차가 감소함을 확인하여 따름정리 1을 검증했습니다.
• 파이프라인 비교: 합성 모델, CIFAR-10, NanoGPT/LLaMA 350M 학습 전반에 걸쳐, Pre-polar 파이프라인은 신호 정렬도와 검증 손실 측면에서 Post-polar 및 Polar-only보다 일관되게 우수한 성능을 보였습니다. 특히, Post-polar의 성능은 $\beta$에 관계없이 평탄하게 유지되었는데, 이는 직교화가 노이즈로 인해 손실된 신호 방향을 복구할 수 없다는 이론적 예측을 확인시켜 줍니다.
• 엔드 투 엔드 학습: LLM 사전 학습 작업(NanoGPT 및 LLaMA 350M)에서 Pre-polar 파이프라인(표준 Muon)은 다른 변형들에 비해 낮은 검증 손실을 달성하였으며, 이는 모멘텀이 단순한 가속 기법이 아니라 신호 복구를 위한 필수 요소라는 주장을 뒷받침합니다.

의의
본 논문은 왜 Muon에 모멘텀이 필요한지, 그리고 왜 연산 순서(직교화 전 노이즈 제거)가 중요한지에 대한 최초의 이론적 근거를 제공합니다. 이는 모멘텀을 행렬 기반 옵티마이저에서 단순한 가속 메커니즘이 아닌, 노이즈에 맞서 특이 부공간을 안정화하는 스펙트럼 필터로 재정의합니다. "먼저 노이즈를 제거하고, 그 다음 직교화하라"는 지침은 비선형 스펙트럼 단계를 사용하는 미래의 행렬 인식 옵티마이저를 위한 설계 원칙을 제시하며, 신호와 노이즈를 분리하는 작업이 스펙트럼 정규화나 직교화에 앞서 수행되어야 함을 시사합니다. 저자들은 본 논문의 이론이 Muon의 이점을 설명하지만, 이러한 이점을 예측하지 못하는 현재의 이론적 경계의 한계 또한 지적하며, 표준 최적화 이론과 스펙트럼 옵티마이저의 실제 동작 사이의 간극을 짚어냈습니다.