사전학습 데이터 규모가 커지면 Muon의 파인튜닝 우위가 뒤집힌다

무엇인가

이 논문이 푸는 문제는 최근 LLM 사전학습에서 AdamW를 위협하는 Muon의 이점이 과연 파인튜닝 이후에도 남는가이다. Muon은 같은 계산량에서 사전학습 손실을 더 빨리 낮춘다고 보고됐지만, 선행 연구들은 사전학습 손실을 최소화하는 선택이 파인튜닝 후 최선의 모델을 주지 않는다고 지적한다. 저자들은 성능을 '학습–망각 트레이드오프', 즉 파인튜닝으로 목표 과제를 얼마나 잘 배우면서 사전학습 지식을 얼마나 보존하는가로 정의하고, 두 옵티마이저로 사전학습한 베이스 모델을 이 기준으로 비교한다.

어떻게 동작하나

비교 지표는 세 가지다. 먼저 각 베이스 모델마다 파인튜닝 러닝레이트를 넓게 스윕해 얻을 수 있는 (파인튜닝 손실, 사전학습 손실) 점들의 집합을 만들고, 그 왼쪽 아래 파레토 프런티어를 '학습–망각 프런티어'라 부른다. 두 옵티마이저의 프런티어 사이 부호 있는 평균 최소 유클리드 거리를 '파인튜닝 갭'으로 정의해, 양수면 AdamW 프런티어가, 음수면 Muon 프런티어가 우세하다고 읽는다. 데이터 예산은 토큰–파라미터 비율(TPR)로 측정하고, AdamW가 Muon을 앞지르기 시작하는 최소 TPR을 '크로스오버 TPR'이라 한다. 실험은 사전학습 손실을 최소화하도록 하이퍼파라미터를 고른 설정과, 손실을 무시하고 파인튜닝 후 프런티어가 가장 좋아지도록 고른 설정 두 가지로 나눠 수행한다.

무엇과 다른가

실험 규모는 OLMo 스타일 30M~600M 모델을 1억 5천만 개에서 1540억 개 DCLM 토큰으로 사전학습하고, GSM8K·StackMathQA·Alpaca·MusicPile 네 데이터셋에서 파인튜닝하는 방식이다. 파인튜닝 러닝레이트는 AdamW 10⁻⁴~10⁻², Muon 6×10⁻⁴~2×10⁻²로 스윕했고, 총 사전학습 실행 약 600회, 파인튜닝 실행 약 29,000회를 돌렸다. 핵심 결과는 StackMathQA 기준으로 사전학습 데이터가 24억 토큰일 때는 Muon이 파레토 우세(갭 음수)지만, 그보다 데이터가 많아지면 AdamW가 더 낮은 파인튜닝 손실과 더 적은 망각을 보이며 갭 부호가 뒤집힌다는 것이다. 이 역전은 최저 사전학습 손실 설정과 최적 프런티어 설정 모두에서, 그리고 나머지 데이터셋에서도 나타난다. 또한 크로스오버 TPR은 모델이 커질수록 작아져, StackMathQA에서 30M 모델은 320 TPR, 600M 모델은 10 TPR이었다. 저자들은 그 원인으로 모델 크기가 커질수록 Muon의 사전학습 손실 우위가 줄고, 같은 파인튜닝 손실에서 Muon의 망각(degradation)이 AdamW보다 커진다는 점을 든다.

어떻게 쓰나

이론 설명의 중심 개념은 margin-to-spread ratio(MSR)다. 각 예측에서 정답 로짓과 오답 로짓 평균의 차이를, 오답 로짓들의 퍼짐(표준편차)으로 나눈 값으로, 클수록 출력 분포의 작은 교란에 덜 민감하다. 저자들은 키–값 쌍을 외우는 선형 연상 기억 모델에서 데이터 예산 γ=N/d를 키우며 Muon(정확한 직교화 버전)과 전체 배치 경사하강법(GD)을 비교한다. 0에서 출발한 단 한 번의 업데이트를 해석적으로 풀어, γ가 충분히 크면 GD의 MSR이 Muon보다 커지고 두 MSR 비율이 3π/8≈1.18에 1/γ 속도로 아래에서 수렴함을 증명한다(정리 4.2). 구체적으로 ρ_GD≈√(d/γ)(1−1/γ), ρ_Muon≈(8/3π)√(d/γ)(1−1/(10γ))이고 비율은 (3π/8)(1−9/(10γ))로 근사된다. 10⁴ 스텝까지 학습한 시뮬레이션에서도 같은 예산 추세가 유지되며, 섭동 허용도(σ50)도 같은 방향으로 움직인다.

전제와 한계

실제 언어모델에서도 이 추세를 확인한다. 5개 모델 크기의 AdamW·Muon 사전학습 모델에 대해 DCLM 홀드아웃 텍스트에서 정답 다음 토큰 로짓과 나머지 어휘 로짓 사이 마진의 중앙값으로 MSR을 계산했더니, 데이터 예산이 작을 때는 Muon의 MSR이, 충분히 커지면 AdamW의 MSR이 더 컸고 두 비율은 1을 넘어섰다. 다만 MSR 크로스오버와 파인튜닝 갭 크로스오버가 정확히 같은 예산에서 나타나지는 않는다. 학습–망각 트레이드오프가 사전학습 손실에도 함께 의존하기 때문이다.

개발자 관점의 의미는 분명하다. 사전학습 손실 하나만으로 옵티마이저를 고르면, 데이터 예산이 큰 구간에서 파인튜닝 적응력과 지식 보존을 놓칠 수 있다. 특히 사전학습 토큰 수가 모델 파라미터의 수십 배를 넘어가는 본격적인 학습을 계획한다면, Muon 대비 AdamW의 파인튜닝 후 학습–망각 프런티어를 자체 데이터로 확인할 가치가 있다. 또 하나의 실무적 관찰은 AdamW로 사전학습한 모델은 AdamW로 파인튜닝할 때 일관되게 더 좋았고, Muon 사전학습 모델은 두 파인튜닝 옵티마이저 사이에 일관된 우위가 없었다는 점이다.

저자들이 밝힌 한계와 전제도 분명하다. 이론 분석은 키와 값이 독립 표준 가우시안인 등방성 연상 기억 모델에 기반하므로, 다른 데이터 분포에서는 그래디언트 스펙트럼과 경쟁 점수 분산이 달라져 옵티마이저 순서가 바뀔 수 있다. 또한 파인튜닝 데이터가 사전학습 데이터와 매우 가까우면 이는 사실상 사전학습 연장이므로 역전이 나타난다고 단정하지 않는다. MSR과 파인튜닝 크로스오버의 예산이 정확히 일치하지 않는다는 점, 그리고 사전학습 손실이 적응성의 좋은 대리 지표가 아니라는 결론이 Muon을 넘어 하이퍼파라미터·데이터·아키텍처 선택 전반에 적용될 수 있다는 점을 향후 과제로 남긴다.