How Much Orthogonalization Does Muon Need?

기술 요약: Muon에는 어느 정도의 직교화(Orthogonalization)가 필요한가?

문제 정의
Muon 옵티마이저는 불량한 조건의 모멘텀 업데이트를 근사적인 반직교(semi-orthogonal) 업데이트로 대체함으로써 신경망 학습을 개선합니다. 표준 수치 선형 대수학은 모든 특이값(singular values)을 1로 매핑하는 고정밀 극분해(polar decomposition)를 목표로 하지만, Muon의 맥락에서는 엄격한 정확도가 불필요할 수 있음을 시사합니다. 본 연구가 다루는 핵심 질문은 다음과 같습니다: Muon이 학습 품질을 유지하기 위해 실제로 요구하는 직교화 정확도는 어느 정도인가? 구체적으로, 본 논문은 기존 방식(Muon-Jordan 및 Polar Express 등)에서 사용되는 고차 다항식 반복의 계산 비용이 정당한지, 아니면 더 저렴하고 낮은 차수의 근사치가 유사한 학습 결과를 달-성할 수 있는지를 조사합니다.

방법론
저자들은 Muon의 저정밀도(bfloat16) 특이값 밴드에 특화되어 설계된 **완화된 3차 뉴턴-슐츠 스케줄(relaxed cubic Newton–Schulz schedule, 이하 cubic5)**을 제안합니다.

1. 완화된 타겟(Relaxed Targets): 클래식 극분해와 달리, 이 방법은 기계 정밀도(machine precision)를 목표로 하는 대신 $[0.7, 1.3]$의 완화된 밴드를 목표로 합니다. 이는 bfloat16 환경에서 유효한 하한선($l_0 = 7 \times 10^{-3}$) 미만의 특이값은 실질적으로 의미가 없다는 점을 인정하는 것입니다.
2. 적응형 계수 도출(Adaptive Coefficient Derivation): 저자들은 각 반복마다 적응형 3차 다항식 $f_t(x) = a_t x + b_t x^3$를 도출합니다. 계수는 현재의 최악의 경우 하한 특이값 경계($l_t$)와 고정된 상한 피크($u=1.3$)를 기반으로 계산됩니다. 이를 통해 다항식이 보호된 구간 $[l_t, r_t]$를 타겟 밴드로 매핑하도록 보장하며, 새로운 하한을 최대화하기 위해 양 끝점을 균형 있게 조절합니다.
3. 계산 효율성:
   • Cubic5: 단일 3차 단계는 두 번의 지배적인 행렬 곱셈(그람 행렬 형성 및 이를 다시 곱함)을 필요로 합니다. 5단계 수행 시 총 10회의 지배적인 행렬 곱셈이 필요합니다.
   • Baseline (Muon-Jordan/Polar Express): 5차 단계는 세 번의 지배적인 행렬 곱셈을 필요로 합니다. 5단계 수행 시 총 15회의 지배적인 행렬 곱셈이 필요합니다.
   • 이 방법은 작은 그람 행렬에서의 연산을 최소화하기 위해 키가 큰(tall) 행렬에 대한 "전치 트릭(transpose trick)"을 활용합니다.

주요 기여

• 완화된 3차 스케줄의 도출: bfloat16의 유효 하한에 적응된 계수를 가진 5단계 3차 뉴턴-슐츠 스케줄을 도출하여, 5번의 반복 내에 타겟 스펙트럼 밴드에 도달함을 입증했습니다.
• 비용-품질 프런티어(Cost-Quality Frontier) 특성화: 본 연구는 학습 품질이 극분해의 정확도에 의해 단조적으로 결정되지 않음을 확립했습니다. 저렴한 비용의 3차 근사법이 고비용의 5차 방식 및 특정 영역에서의 정밀 SVD와 대등한 성능을 달mi할 수 있음을 보여줍니다.
• 다양한 규모에서의 경험적 검증: 저자들은 다음을 통해 방법을 검증했습니다:
  • 합성 진단(Synthetic Diagnostics): 저정밀도 행렬 산술에서 스칼라 다항식 전이가 유효함을 확인했습니다.
  • NanoGPT 절제 실험(Ablations): GPT-2 Small에 대해 cubic5, 절단된 Polar Express, 절단된 Muon-Jordan, 명시적 FP32 SVD를 비교했습니다.
  • 대규모 학습: Megatron-LM을 사용하여 1B에서 4B 파라미터 범위의 하이브리드 MoE/Mamba 모델(Nemotron-3-Nano)을 학습시켰습니다.

결과

• 합성 벤치마크: cubic5 방식은 보호된 범위 $[7 \times 10^{-3}, 1]$ 내의 특이값들을 수치적 폭발 없이 타겟 밴드로 성공적으로 매핑했습니다. NVIDIA H200 GPU에서의 마이크로 벤치마크 결과, cubic5는 행렬 형태(직사각형 vs 정방형)에 따라 1.16x에서 1.52x 사이의 속도 향상을 보이며 Polar Express 대비 직교화 서브루틴 비용을 줄였습니다.
• GPT-2 Small 절제 실험:
  • 5단계 cubic5의 최종 검증 손실(validation loss)은 3.2811입니다.
  • 5단계 Polar Express는 3.2798을 기록했습니다.
  • 5단계 Muon-Jordan은 3.2822를 기록했습니다.
  • 명시적 FP32 SVD는 3.2809를 기록했습니다.
  • 관찰 사항: 가장 강력한 변체들 간의 검증 손실 차이는 0.003 미만입니다. 주목할 점은, 정밀 SVD가 근사 다항식 방법보다 우수하지 않았으며, Muon-Jordan에 5번째 단계를 추가했을 때 4단계보다 성능이 오히려 약간 저하되었다는 것입니다.
• 하이브리드 MoE/Mamba 스케일링 (1B–4B):
  • 1B 및 2B 모델에서 cubic5는 Muon-Jordan보다 약간 더 우수한 성능을 보였습니다(검증 손실 약 $10^{-3}$ 낮음).
  • 4B 모델에서는 cubic5가 약간 높았으나($\approx 10^{-3}$), 초기 단계 이후의 학습 궤적은 거의 구별할 수 없을 정도로 유사했습니다.
  • 상대적인 학습 손실 차이는 빠르게 감소하여 나머지 학습 예산 동안 0 근처에서 변동했습니다.

의의 및 주장
본 논문은 cubic5가 보편적으로 우월한 업데이트 규칙이라기보다, 실용적이고 저비용인 Muon 직교화 변체라고 겸손하게 주장합니다.

• 옵티마이저 vs. 분해 정확도: 결과는 Muon이 엄격한 극분해 정확도보다는 "스펙트럼 재형성(spectral reshaping)" 효과로부터 이득을 얻는다는 가설을 뒷받침합니다. 최적의 학습을 유도하는 정확한 특이값 맵은 수학적 1로의 수렴뿐만 아니라, 학습 단계와 학습률에 따라 달라집니다.
• 실용적 유용성: cubic5는 비용-품질 프런티어 상의 원칙적인 지점을 제공합니다. 이는 테스트된 영역에서 표준 Muon-Jordan 5차 업데이트와 대등한 학습 품질을 유지하면서도, 지배적인 행렬 곱셈 횟수를 약 1/3 줄임으로써(15회 대비 10회) 효율성을 확보합니다.
• 한계점: 저자들은 이러한 결과가 **학습 품질의 동등성(parity)**을 확립하는 것이지, 엔드 투 엔드 처리량(throughput)의 속도 향상을 보장하는 것이 아님을 명시했습니다. 실제 벽시계 시간(wall-clock) 속도 향상은 하드웨어 커널, 배치 크기, 그리고 전체 옵티마이저 루프 내에서 직교화 단계가 차지하는 비중에 달려 있습니다. 본 논문은 cubic5가 수학적 의미에서 Polar Express보다 더 나은 극분해 알고리즘이라고 주장하는 것이 아니라, 특정 Muon 학습 맥락에서 더 효율적인 스펙트럼 변환임을 주장하는 것입니다.