DynMuon: Muon의 동적 스펙트럴 셰이핑 관점

Mu on 스타일의 업데이트가 특정 학습 단계에서 고유한 강점과 한계를 보인다. 본 논문은 p를 스펙트럴(exponent)으로 표현한 확장 가능한 가족을 도입하고, 초반에는 high-curvature 방향을 강조하는 양의 p가 초기에 가속을 촉진하며, 후반에는 flat 방향에 residual signal이 남아 있을 때 mildly negative p가 수렴을 개선한다는 모형을 제시한다. 이 모형에 기초해 DynMuon은 p를 학습 진행에 따라 동적으로 감소시키는 스케줄을 제시하고, Muon 대비 다양한 규모의 Transformer에서 더 낮은 검증 손실 및 더 빠른 목표 손실 도달을 보인다.

• 출발점과 한계: Muon은 M = UΣV^⊤를 UV^⊤로 바꾸어 특이 방향은 보존하되 특이값을 평탄화한다. 그러나 학습 진행에 따라 최적의 스펙트럼 변형이 달라질 수 있으며, 고정된 p가 항상 최적은 아니다.
• 이 논문의 해결 원리: D^(p) = UΣ^pV^⊤를 업데이트에 적용하고, 로컬 커브처(H) 방향에 따라 G_t의 모드별로 p의 영향을 다르게 받는다는 모형을 도출한다. G_t ≈ κ_t H E_t + Ξ_t로 근사하고, (G_tG_t^⊤)^{(p-1)/2}를 H^{(p-1)/2}로 근사하는 등 곡률-기반 재가중으로 업데이트를 재배치한다.
• 달라지는 점: p가 큰 경우 strong 모드의 수렴이 빠르고, p가 음수로 작아질수록 flat 모드의 수렴에 더 큰 기여를 한다. late-stage에서 residual signal은 flat 모드에 집중되며, 이는 mildly negative p가 이들 방향의 contraction을 증가시키되 노이즈 증가를 관리해야 함을 시사한다.
• DynMuon의 핵심 아이디어: 학습 초기에 양의 p를 유지하고, 이후 점진적으로 p를 감소시키는 로지ス틱 스케줄로 전과 후를 연결한다. NS 보정과 A^p/2 보정의 조합으로 계산 효율을 유지한다.
• 실험적 요약: 127M/601M/1.11B 파라미터에서 검증 손실이 더 낮게 나오고, Muon 대비 10.6%에서 26.5% 더 빠른 목표 도달이 확인된다.

1. 업데이트 연산: D^(p)t = U Σ^p V^⊤; W{t+1} = W_t - η D^(p)t, G_t ≈ ∇L(W_t) ≈ κ_t H E_t; D^(p)t와 H의 관계를 통해 모드별 업데이트를 재가중한다. 2) 모드 분해: δ{i,t+1} = (1 - η_t κ_t h_i^{(p+1)/2}) δ{i,t} - η_t h_i^{(p-1)/2} ξ_{i,t}; 3) 모드 신호-노이즈 트레이드오프: h_i는 모드의 곡률이며 ci,t는 노이즈 수준, δ_i,t은 잔여 신호의 모드별 에너지다. 4) 단계 의존성 예측: 잔여 신호가 큰 방향은 high-curvature, 잔여 신호가 작은 방향은 flat 방향에 집중되며, p가 이를 어떻게 바꾸는지 분석한다. 5) DynMuon 알고리즘: Algorithm 1에서 logistic schedule로 pt를 계산하고, pt ≥ 1/4일 때는 원래 업데이트를 사용, 0 ≤ pt < 1/4일 때는 Newton–Schulz 보정, pt < 0일 때는 Fast–Spectral 보정을 사용한다. 6) Fast–Spectral 구현: Algorithm 2에서 Xn := X/||X||_F, Y^µ는 Muon 업데이트로, A := XnXn^⊤, E := A - I, C := I + δE + 1/2 δ(δ-1) E^2, Xe := ||X||_F^p C Y^µ로 최종 업데이트를 얻는다. 7) 실험 설계: GPT-style 및 Qwen-style 디코더-형 Transformer에서 10B/20B 토큰, 127M/601M/1.11B 파라미터 규모의 모델을 대상으로 Muon, NorMuon, DynMuon 비교. 학습률 기본값은 Muon/DynMuon 0.01, AdamW 0.002; pmax = 1, pmin = -0.25을 기본값으로 사용하며 τ, w로 스케줄링을 제어한다.

주요 결과: 10B 토큰 규모에서 127M 파라미터의 Best Val. Loss가 Muon 3.190에서 DynMuon 3.171로 개선되며, Steps to Target는 16000에서 12500으로 감소해 21.9%의 Step Saving을 달성한다. 601M 파라미터 규모에서 Muon 2.872에서 DynMuon 2.858로 개선되며, Steps 16000에서 13950으로 감소해 12.8%의 Step Saving을 기록한다. 1.11B 파라미터 규모에서 Muon 2.788에서 DynMuon 2.776로 개선되고, Steps 16000에서 14300으로 감소해 10.6%의 Step Saving이 나타난다. 20B 토큰 예에서 127M 파라미터의 Best Val. Loss는 Muon 3.139에서 DynMuon 3.124로 감소하며, Steps 30400에서 22350으로 감소해 26.5%의 Step Saving이 확인된다. 601M/1.11B 파라미터에서 각각 30400에서 25000(절감 17.8%), 30400에서 26450(절감 13.0%)으로 감소한다. Qwen-Style Architecture에서도 DynMuon이 Muon보다 낮은 Validation Loss를 달성한다. 127M 모델에서 FineWeb-Edu 데이터셋으로의 실험에서도 DynMuon의 성능이 우수하고 pmin=-0.25가 최적으로 나타난다. Exact SVD 기반 비교에서도 DynMuon은 실측 정확도에서 비슷한 추적 능력을 보이며 Muon 대비 우수하며, 동일 스케줄의 경우 속도 차이는 약 3배 미만이다. 학습 시 Seed robustness 실험에서도 DynMuon이 Muon에 비해 변동성이 작고 안정적이다. NorMuon 대비 DynMuon 우위가 확인되며, ROI에서의 효과가 재현된다. CE–Brier 인터폴레이션에서도 mild negative p가 전반적으로 이득을 보이며, 지나치게 음수인 p는 성능 저하를 유발한다.

1. 아키텍처 구조: M = U Σ^p V^⊤를 사용해 업데이트를 재가중하며, DynMuon은 ν_t에 따라 pt를 계산하는 로직을 따른다. 2) 핵심 메커니즘의 수학적 기반: D^(p)t = U Σ^p V^⊤, ∇L(W_t) ≈ κ_t H E_t, H = Q Λ Q^⊤, δ{i,t+1} = (1 − η_t κ_t h_i^{(p+1)/2}) δ_{i,t} − η_t h_i^{(p−1)/2} ξ_{i,t}, E_t = W_t − W⋆. 3) Prior work 대비 차별점: fixed p를 고수하는 Muon 대비, p를 학습 진행에 따라 조정하는 동적 스케줄링(DynMuon)과 NS 근사 기반의 효율적 구현. 4) 구현 및 학습 세부사항: Algorithm 1 DynMuon, Algorithm 2 Fast–Spectral, pmax=1, pmin=-0.25, τ=0.02, w=0.01 등; NS 근사에서 Yµ는 Muon 업데이트에 해당하고, A^p/2의 보정은 I를 기준으로 한 2차 테일러 확장으로 근사한다. 5) 수치적 검증: 10B/20B 토큰 사이즈에서 3개의 모델 규모에서 DynMuon이 Muon 대비 더 낮은 검증 손실과 더 빠른 수렴을 나타냄.