🤖 machine learning

Muon $^p$ : Muon with Fractional Spectral Powers

이 논문은 그래디언트의 특이값(singular values)에 분수 스펙트럼 거듭제곱(fractional spectral powers)을 적용하여 Muon과 경사 하강법(gradient descent) 사이를 보간하고, 효율적인 이변량 재귀(bivariate recurrences)를 활용하여 이러한 업데이트를 근사하며, 가치 있는 특이 스펙트럼 정보를 보존하면서 수십억 규모 모델의 미세 조정(finetuning)에서 향상된 성능을 입증하는 새로운 옵티마이저인 Muon $^p$ 를 소개한다.

원저자: Yihe Dong, Will Sawin

게시일 2026-06-15

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yihe Dong, Will Sawin

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대하고 매우 똑똑한 로봇(신경망)에게 말하기, 코드 작성, 또는 수학 문제 풀기를 가르치고 있다고 상상해 보세요. 이 로봇을 가르치기 위해, 당신은 예시를 보여주고 실수를 바로잡아 주어야 합니다. 여기서 "옵티마이저(optimizer)"는 각 실수 후에 로봇의 뇌를 어떻게 조정할지 결정하는 교사의 방법입니다.

문제점: "전부 아니면 전무" 식의 선생님 (Muon)

최근에 Muon이라는 유명한 교수법이 화제가 되었습니다. 이것은 다음과 같이 작동합니다:

로봇의 실수들은 서로 다른 "방향"이나 "주파수"를 가지고 있다고 상상해 보세요. 어떤 방향은 매우 크고 명확한 반면(지배적 모드), 어떤 방향은 아주 작은 속삭임과 같습니다(작은 특이값).

기존의 선생님들(표준 경사 하강법 같은 방식)은 큰 소리가 나는 방향에는 크게 소리치며 교정하지만, 작은 속삭임은 무시하곤 했습니다.
Muon은 "평평한(flat)" 선생님이 되기로 결정했습니다. 그것은 모든 방향을 살펴보고 이렇게 말했습니다. "자, 이 모든 방향을 똑같이 취급하자!" 그리고 모든 교정의 볼륨을 평평하게 만들어, 작은 속삭임도 큰 외침과 똑같은 주의를 기울이게 했습니다.

함정: 이것은 로봇을 처음부터 가르칠 때(사전 학습, pretraining)는 훌륭하지만, 한 가지 중요한 정보를 버리게 됩니다. 바로 각 방향이 실제로 얼마나 강한가 하는 정보입니다. 때로는 그 큰 소리가 나는 방향들이 정말 중요한 것일 수도 있는데, 단순히 속삭임을 듣기 위해 사이렌 소리의 볼륨을 줄여버리는 것은 문제가 될 수 있습니다.

해결책: "딱 적당한" 선생님 (Muonp)

이 논문의 저자들은 Muonp를 소개했습니다. Muonp를 "골디락스(Goldilocks)" 존을 찾아내는 선생님이라고 생각하세요.

모든 볼륨을 완전히 평평하게 만드는 대신(Muon처럼), 혹은 원래의 크기를 그대로 유지하는 대신(기존 선생님들처럼), Muonp는 볼륨을 아주 조금만 낮춥니다. 그것은 수학적인 "디머 스위치(dimmer switch, 밝기 조절 스위치)"(p라고 불리는 숫자)를 사용하여, 큰 소리가 나는 방향의 원래 강도를 어느 정도 유지하면서도 작은 소리들을 끌어올려 줍니다.

만약 p = 0 이라면: 원래의 Muon입니다 (완전한 평탄화).
만 if p = 1 이라면: 표준적인 선생님입니다 (변화 없음).
만약 p가 그 사이의 값(예: 1/3)이라면: 이것이 바로 Muonp입니다. 원래의 형태를 유지하면서도 교정의 내용을 부드럽게 다듬어 줍니다.

어려운 점: 마법 같은 수학적 트릭

당신은 이렇게 생각할 수도 있습니다. "그래, 그냥 볼륨을 조금만 줄이면 되겠네." 하지만 거대한 행렬(로봇의 뇌)의 세계에서 이 계산을 수행하는 것은 매우 어렵습니다. 보통 이 방향들의 볼륨을 조절하려면, 뇌를 해체하고, 모든 부분을 측정하고, 다시 조립해야 합니다. 이는 느리고 비용이 많이 드는 작업입니다.

저자들은 이 "볼륨 조절"을 단순한 한 단계 레시피로는 할 수 없다는 것을 증ato했습니다. 원래의 뇌 구조를 기억하면서 볼륨을 조절하기 위해서는 더 복잡한 두 단계 레시피가 필요합니다.

그들은 컴퓨터가 이미 갖추고 있는 빠른 수학 칩(GPU)을 사용하여 이 "조절된" 교정값을 계산할 수 있는 새롭고 빠른 레시피("이변량 다항식 재귀(bivariate polynomial recurrences)"라고 불리는 것)를 발명했습니다. 이 마법 같은 지름길 덕분에 Muonp는 Muon만큼 빠르면서도 더 똑똑합니다.

발견한 사실: 직업에 따라 다르다

논문 저자들은 이 새로운 선생님을 십억 단위 규모의 모델에 테스트하여 흥미로운 성능 차이를 발견했습니다:

새로운 뇌를 만들 때 (사전 학습, Pretraining):
만약 로봇을 처음부터 가르치고 있다면, "평평한" 선생님(Muon)이 실제로 더 낫습니다. 기초를 탄탄히 쌓기 위해서는 가능한 모든 방향을 똑같이 탐색해야 하기 때문입니다. 기존의 "큰" 방향들을 붙잡고 있으려 하는 Muonp는 이 단계에서는 약간 성능이 떨어졌습니다.
전문가를 미세 조정할 때 (미세 조정, Finetuning):
이미 똑똑한 로봇을 데려와서 특정 새로운 기술(예: 코딩이나 수학)을 가르친다면, Muonp가 승자입니다.
- 왜 그럴까요? 이 단계에서 로봇은 이미 기초를 알고 있습니다. "큰 소리"가 나는 방향들이 바로 새로운 과제를 수행하는 데 중요한 데요. Muonp는 이 중요한 방향들에 집중하면서도 작은 소리들을 도와줍니다.
- 결과: 수학적 추론, 코딩, 언어 이해와 같은 과제에서 Muonp는 기존의 Muon이나 표준적인 선생님들보다 로봇을 더 똑똑하고, 빠르고, 정확하게 만들었습니다.

"커리큘럼" 아이디어

저자들은 또한 영리한 트릭을 시도했습니다. 학습의 대부분은 "평평한" 선생님(Muon)을 사용하고, 마지막 몇 단계에서 "디머(dimmer)" 선생님(Muonp)으로 전환하는 것입니다. 이것은 놀라운 효과를 냈습니다. 마치 학생에게 넓은 붓으로 기초를 가르친 다음, 마지막 세부 사항을 다듬을 때는 아주 가는 펜으로 바꾸는 것과 같습니다.

요약

Muonp는 인기 있는 Muon 옵티마이저의 더 똑똑하고 유연한 버전입니다. 이것은 모든 학습 방향을 똑같이 만들지 않고, 대신 그 중요성을 부드럽게 조절합니다.

수학: 컴퓨터 속도를 늦추지 않고 이러한 조정을 계산할 수 있는 영리하고 빠른 지름길을 사용합니다.
결과: Muonp는 대규모 AI 모델의 **미세 조정(fine-tuning)**을 위한 최고의 도구이며, 코딩이나 수학 같은 특정 기술을 이전보다 더 잘 마스터하도록 돕습니다. 하지만 **처음부터 학습(pretraining)**할 때는 기존의 "평평한" Muon이 여전히 챔피언입니다.

기술 요약: Muonp – 분수 스펙트럼 거듭제곱을 이용한 Muon

문제 정의

Muon 옵티마이저는 특히 대규모 언어 모델(LLM)의 심층 신경망 학습에서 강력한 도구로 부상했습니다. 이 방식은 그래디언트 행렬 $G = USV^\top$ 를 그 극값(polar factor)인 $UV^\top$ 로 대체함으로써, 특잇값 스펙트럼(singular value spectrum)을 효과적으로 평탄화(flattening)합니다. 이러한 "스펙트럼 평탄화"는 균형 잡힌 학습을 촉진하고, 지배적인 모드(dominant modes)가 학습 업데이트를 독점하는 것을 방지합니다.

그러나 저자들은 Muon 방식의 결정적인 한계를 지적합니다: 완전한 평탄화는 모든 특잇값 정보를 버린다는 점입니다. 이러한 균일성은 사전 학습(pretraining) 단계에서 광범위한 특징 발견에 유익하지만, 어떤 잔차 방향(residual directions)을 변화시키는 것이 가장 가치 있는지를 나타내는 특잇값의 크기 정보가 중요한 적응(adaptation) 작업(예: 파인튜닝)에서는 최선이 아닐 수 있습니다. 본 연구가 다루는 핵심 질문은 다음과 같습니다: Muon의 행렬 수준 스펙트럼 이점을 유지하면서도, 효과적인 적응에 필요한 크기 구조(magnitude structure)를 어떻게 보존할 것인가?

방법론

이론적 정식화: Muonp

저자들은 분수 스펙트럼 거듭제곱을 사용하여 가중치를 업데이트하는 Muon의 일반화 버전인 Muonp를 소개합니다:
$\text{Update} \propto US^pV^\top$
여기서 $p \in (0, 1)$ 은 유리수입니다.

기하학적 해석: 저자들은 이 업데이트 규칙이 $q = 1 + 1/p$ 인 Schatten $q$ -노름(Schatten $q$ -norm) 하에서의 최급 강하(steepest-descent) 방향에 해당함을 증명합니다(정리 2.1).
보간(Interpolation):
- $p = 0$ : 표준 Muon( $UV^\top$ )에 해당하며, 균일성을 극대화합니다.
- $p = 1$ : 표준 경사 하강법( $USV^\top = G$ )에 해당합니다.
- $0 < p < 1$ : 두 방식 사이를 보간하며, 부분적인 크기 정보를 유지합니다.

알고리즘적 과제 및 해결책

$US^pV^\top$ 를 직접 계산하려면 특잇값 분해(SVD)가 필요하며, 이는 대규모 모델에서 계산 비용이 매우 높습니다. 저자들은 명시적인 SVD 없이 이 업데이트를 근사하는 방법을 제시합니다:

단변수 반복의 불가능성: 저자들은 임의의 유리수 $p$ 에 대해 고정된 단변수 다항식 반복 $S_{n+1} = f(S_n)$ 이 $S^p$ 로 수렴할 수 없음을 증명합니다(보조정리 2.2). $p=0$ 인 극값(polar factor)과 달리, 분수 거듭제곱은 알고리즘이 원래의 행렬을 "기억"해야 합니다.
이변수 다항식 재귀(Bivariate Polynomial Recurrence): 이를 해결하기 위해 저자들은 저차 홀수 이변수 재귀식을 도출합니다. 알고리즘은 특잇값이 $S_{n+1} = f(S_n, S)$ $S_{n + 1} = f (S_{n}, S)$ 를 통해 진화하는 행렬 $G_n = US_nV^\top$ $G_{n} = U S_{n} V^{⊤}$ 를 반복적으로 계산합니다.
- 정리 2.3은 임의의 홀수 이변수 다항식 $f(x, y)$ 에 대해, 항 $Uf(S_n, S)V^\top$ 를 행렬 곱셈과 전치만을 사용하여 계산할 수 있으며, 이를 통해 Muon의 계산 구조를 보존할 수 있음을 보여줍니다.
- 정리 2.4는 임-의의 $p \in (0, 1)$ 에 대해 $S^p$ 로 수렴하는 그러한 다항식이 존재함을 입증합니다.
명시적 구성: 논문은 이러한 다항식에 대한 명시적인 공식(예: $p=1/3$ 인 경우 $f(x, y) = x + c(y - x^3)$ )을 제공합니다. 계수는 $S^p$ 가 끌개(attracting) 고정점이 되도록 선택됩니다.

구현 세부 사항

복잡도: 계산 복잡도는 여전히 $O(\omega k N)$ 입니다 (여기서 $\omega$ 는 행렬 곱셈 복잡도, $k$ 는 단계당 곱셈 횟수, $N$ 은 반복 횟수). Muonp는 상수 배만큼의 행렬 곱셈 차이만 있을 뿐 Muon과 다를 바 없습니다.
안정성: bfloat16으로 구현했을 때 수치적으로 안정적이며, 현대 하드웨어에 적합합니다.
커리큘럼 호환성: Muon( $p=0$ )이 Muonp의 특수 사례이므로, 학습 중 두 방식을 전환(예: Muon으로 사전 학습 후 Muonp로 파인튜닝)할 때 상태 초기화 없이 다항식 근사치만 변경하면 됩니다.

주요 기여

Muonp 옵티마이저: 명시적인 SVD 없이 임의의 $p \in (0, 1)$ 에 대해 $US^pV^\top$ 를 계산하는 일반화된 옵티마이저를 도입했습니다.
이론적 한계 규명: 분수 스펙트럼 거듭제곱은 고정된 단변수 다항식 반복으로는 계산될 수 없으며, 새로운 이변수 접근 방식이 필요함을 증명했습니다.
구성적 알고리즘: Muon의 행렬 곱셈 전용 구조를 유지하면서 분수 거듭제곱을 근사하는 저차 홀수 이변수 재귀식을 도출했습니다.
실증적 검증: 수학적 추론, 코딩 및 다운스트림 작업의 파인튜닝 시나리오에서 Muonp가 Muon 및 AdamW보다 크게 우수한 성능을 보임을 입증했습니다.
스펙트럼 기하학 분석: 최적의 스펙트럼 기하학은 레지임(regime)에 따라 달라진다는 통찰을 제공합니다: 균일한 업데이트( $p=0$ )는 사전 학습에 유리하고, 부분적 보존( $0 < p < 1$ )은 파인튜닝에 유리합니다.

실험 결과

파인튜닝 성능

Muonp는 Llama3-1B 및 Llama3-3B 모델을 파인튜닝할 때 Muon 및 AdamW에 비해 일관된 개선을 보여줍니다:

수학 및 코드 추론: GSM8K, Numina, OpenCodeInstruct에서 Muonp는 Muon에 비해 pass@8 정확도를 향상시키고 손실(loss)을 줄입니다.
Perplexity: Muonp는 Fineweb 및 Pile of Law 데이터셋에서 더 낮은 검증 perplexity를 달성합니다.
다운스트림 작업: Muonp는 일련의 추론 벤치마크(PIQA, Winogrande, LAMBADA 등)에서 Muon과 AdamW 모두보다 높은 정확도를 기록합니다.
옵티마이저 불일치: 사전 학습 시 동일한 옵티마이저(Muon)를 사용했음에도 불구하고, 파인튜닝 시 Muonp가 Muon보다 우수한 성능을 보입니다. 이는 이 이점이 단순히 옵티마이저의 일관성 때문이 아니라 스펙트럼 기하학의 근본적인 특성임을 시사합니다.

사전 학습 vs 파인튜닝의 이분법

본 논문은 학습 단계에 따른 성능의 뚜렷한 이분법을 밝혀냈습니다:

사전 학습: Muon ( $p=0$ )이 Muonp보다 우수합니다. 저자들은 사전 학습이 광범위한 특징을 발견하고 지배적인 방향이 업데이트를 독점하는 것을 방지하기 위해 "등방화(isotropization)"가 필요하다고 주장합니다.
파인튜닝: Muonp ( $0 < p < 1$ )가 Muon보다 우수합니다. 이 단계에서 모델은 낮은 내재적 차원(intrinsic dimension)을 가지며, 특잇값 정보를 유지하는 것이 가장 관련 있는 잔차 방향에 적응을 집중하는 데 도움을 줍니다.
커리큘럼 실험: 사전 학습의 마지막 단계에서 Muon에서 Muonp로 전환하는 커리큘럼은 즉각적인 손실 감소를 가져오며, 이는 $p$ 를 증가시키는 것이 작은 특잇값 방향에 대한 유효 학습률을 적응시킨다는 가설을 입증합니다.

실행 시간

Muonp는 Muon과 대등한 실행 시간을 유지합니다. 예를 들어, $p=1/3$ (3차 재귀)의 경우, Muonp는 Muon(일부 구현에서 $p=0$ 을 위해 5차 재귀 사용)보다 약간 빠르며, H100 GPU에서 단계당 약 100~120ms의 시간을 소요합니다.

의의 및 주장

본 논문은 Muonp를 Muon 옵티마이저의 원칙적인 일반화로 규정합니다. 주요 의의는 다음과 같습니다:

간극 해소: Muon의 균일성과 경사 하강법의 크기 민감도 사이를 체계적으로 보간하여, "스펙트럼 평탄화"와 "크기 보존" 사이의 트레이드오프를 해결합니다.
레지임 의존적 최적화: 최적의 스펙트럼 기하학은 보편적이지 않다는 증거를 제시합니다: 사전 학습에는 완전한 평탄화가 필요하고, 파인튜닝에는 부분적인 평탄화가 필요합니다.
실용적 효율성: SVD의 비용 없이 이러한 이론적 이점을 달 순 있으며, 수십억 개의 파라미터 모델로 확장하는 데 필수적인 하드웨어 친화적인 행렬 곱셈 전용 구조를 유지합니다.
향후 방향: 특정 작업에 최적의 $p$ 를 자동으로 선택하고, Muonp를 사전 학습 커리큘럼에 통합하는 것이 유망한 연구 분야임을 제안합니다.

결론적으로, 특잇값 스펙트럼을 보존하는 것이 적응 시나리오에서 상당한 이득을 가져올 수 있으며, Muonp는 이를 달성하기 위한 계산 효율적이고 원칙적인 메커로니즘을 제공합니다.