MuonGS는 가중치 노름을 고정하지 않고도 Muon의 학습 지속력을 높인다
MuonGS: Pretraining LLMs with Relaxed Weight-Norm Control
무엇인가
이 논문은 LLM 사전학습용 Muon 계열 옵티마이저에서 가중치 노름 제어가 왜 도움이 되는지에 대한 통일된 설명이 아직 없다는 문제에서 출발한다. AdamW의 decoupled weight decay, MuonW, 고정 반경 제약을 쓰는 MuonH, nGPT의 단위 노름 벡터 같은 방법들이 각각 다른 방식으로 노름을 누르고 있지만, 어떤 제어가 본질이고 어디까지 완화해도 되는지는 정리되어 있지 않았다. 저자들은 단일 스텝의 노름이 아니라 학습 전 구간에 걸친 탐색 능력을 보자는 관점을 택한다.
어떻게 동작하나
핵심 지표는 누적 유효 학습률 C_T = Σ η_t / ‖W_t‖_F 다. Muon 업데이트의 상대적 크기는 η_t/‖W_t‖_F 에 업데이트 방향의 노름을 곱한 값으로 결정되므로, 이 비율을 학습 구간마다 더한 값이 그 구간에서 가중치가 상대적으로 얼마나 움직일 수 있는지를 규정한다. 실제로 명제 3.1은 ‖W_T − W_s‖_F / max{r_s, r_T} ≤ √p (C_T − C_s) 를 보인다. 스케일 불변 손실을 가정하면 손실 변화량도 2L√p (C_T − C_s) 로 상한이 잡힌다. 즉 C_T 가 충분히 쌓이지 않으면 상대적 이동도, 손실 감소도 제한된다. 중요한 점은 이 상한이 전체 구간 [0,T] 뿐 아니라 임의의 구간 [s,T] 에도 성립한다는 것으로, 후반 학습에서도 C_T 가 계속 늘어야 탐색이 유지된다는 뜻이다.
무엇과 다른가
이 지표로 세 가지 Muon 변형을 비교하면 차이가 뚜렷하다. 상수 학습률에서 vanilla Muon은 노름이 선형으로 커져 C_T 가 Θ(log T) 에 그치고, 후반 구간 증가분 C_{2T} − C_T 는 log2/√p 로 수렴해 사실상 멈춘다. 반면 MuonW는 C_T ≥ ηT/R_λ (R_λ = max{r_0, √p/λ}) 를, 고정 노름 제어는 C_T = ηT/R 을 정확히 보장해 최소 선형 증가를 얻는다. 학습률을 키우는 우회로도 통하지 않는다. GPT-160M 사전학습에서 학습률을 16배 올려도 vanilla Muon의 C_20,000 은 0.165에서 0.190으로밖에 오르지 않았고, MuonW의 최적 실행들은 0.3~0.4에 도달하며 검증 퍼플렉시티도 21.98~22.02 대 22.49로 더 낮았다.
어떻게 쓰나
여기서 저자들이 얻는 결론은 노름을 유계로 묶는 것이 선형 누적을 보장하는 충분조건이지 필요조건은 아니라는 것이다. 노름 성장이 r_t ≤ r_0 + c·t^a (0 ≤ a < 1) 처럼 충분히 느리면 C_{2T} − C_T = Ω(T^{1−a}) 가 되어 로그 최악의 경우보다는 빠르게 쌓인다. 노름 제어를 느슨하게 풀 수 있는 여지가 생기는 지점이다.
전제와 한계
MuonGS는 이 여지를 쓰는 구체적인 방법으로, Muon 업데이트에서 방사 방향 성분만 제거한다. W_{t+1} = W_t − η_t·Õ_t, Õ_t = O_t − (⟨W_t, O_t⟩_F / r_t^2)·W_t 라는 Gram-Schmidt 투영이 전부다. 현재 가중치와 직교하는 방향 중 Muon 제안 방향에 가장 가까운 것을 고르는 닫힌 형태 해이므로 반복 최적화가 필요 없고, 모멘텀이나 근사 직교화가 섞인 O_t 에도 그대로 적용된다. 결과적으로 ‖Õ_t‖_F ≤ √p 이고 노름 제곱의 증가분이 스텝 크기의 2차항만 남아 r_{t+1}^2 = r_t^2 + η_t^2‖Õ_t‖_F^2 가 된다. AdamP가 스케일 불변으로 판정된 가중치에 같은 접선 공간 투영을 적용하는 것과 유사하지만, MuonGS는 스케일 불변 판정 없이 Muon의 행렬 업데이트에 직접 적용한다.
이론적 보장은 정리 4.1과 명제 4.2로 정리된다. 상수 학습률에서 r_T^2 ≤ r_0^2 + pΣη_t^2 이므로 노름은 O(√t) 로만 자라고, C_T ≥ (2/(pη))(√(r_0^2 + pη^2T) − r_0) = Ω(√T), 후반 구간도 C_{2T} − C_T ≥ ηT/√(r_0^2 + 2pη^2T) = Ω(√T) 이다. 노름을 고정하지 않고도 vanilla Muon의 로그 최악보다 나은 누적을, 그것도 후반까지 유지한다. 매끄러운 스케일 불변 목적함수에 대해서는 후반 절반 구간에서 구면 그래디언트 노름이 K/√T 이하인 반복점이 존재한다는 정상성 보장 O(T^{-1/2}) 도 제시한다.
실험은 FineWeb의 100B 토큰 부분집합에서 Modded-NanoGPT 계열 GPT로 진행했다. 160M 모델(12층, hidden 768, 6헤드)은 3~25B 토큰, 400M 모델(24층, hidden 1024, 8헤드)은 8~63B 토큰으로 약 1~8 Chinchilla 예산에 해당하며, 시퀀스 길이 1024, 전역 배치 512, BF16 혼합 정밀도를 썼다. 은닉 행렬 가중치는 Nesterov 모멘텀 0.95를 쓰는 MuonGS·MuonH·Muon/MuonW로, 임베딩·출력 헤드·1D 파라미터는 AdamW(β=(0.8,0.95), eps 1e-10, weight decay 1e-3, 최대 학습률 0.7/0.004/0.015)로 학습했다. 평가는 FineWeb 검증셋 약 10M 토큰의 교차 엔트로피와 퍼플렉시티다. 학습이 길어질수록 MuonGS가 MuonH보다 낮은 검증 퍼플렉시티를 보였고, MuonH 손실을 αD^{−B} + β 형태로 피팅해 동일 손실 도달에 필요한 토큰 예산을 역산한 결과 토큰 환산 속도 향상이 약 7~15%(초록 기준 최대 약 15%)로 추정됐다. MoE 실험은 32층, hidden 512, 64개 라우티드 전문가(top-7)에 층당 공유 전문가 1개를 더한 약 135M 활성·1B 전체 파라미터 모델을 Ultra-FineWeb에서 10.8B·21.6B 토큰(활성 파라미터의 80배·160배)으로 학습했고, 두 예산 모두에서 MuonGS가 MuonH보다 낮은 최적 검증 퍼플렉시티를 기록했다.
저자들이 명시한 한계는 분명하다. 누적 보장은 최악의 경우 상한이며, C_T 의 충분한 증가는 지속적 탐색을 위한 필요조건일 뿐 충분조건이 아니다. 손실 감소 상한을 얻으려면 스케일 불변성과 립시츠 가정이 필요하고, 실험은 최대 400M 밀집 모델과 1B MoE, 최대 63B 토큰 규모에 머문다. 또한 주어진 학습 예산에 어떤 노름 성장 속도가 가장 좋은지는 여전히 열린 질문으로 남겨두었다.