Newton–Schulz: 행렬 곱만으로 직교화하기

Muon 의 걸음 UVᵀ 는 매 스텝, 층마다 새로 구해야 한다. 큰 모델은 층마다 2차원 가중치 행렬이 여럿이라, 한 스텝에 직교화할 행렬이 수십에서 수백 개다.

“UVᵀ 를 얻으려면 매 스텝, 층마다 특이값 분해를 해야 하는가?”

실제 구현 — Newton–Schulz

켈러 조던(2024)이 공개한 Muon은 매 스텝 특이값 분해를 하지 않는다. 조던은 특이값 분해가 너무 느리고, Shampoo 구현들이 −1/4 제곱을 구하는 데 쓰던 결합 뉴턴 반복은 32비트 정밀도 이상에서만 안정해서 요즘 GPU 에서 느리다고 적었다. 대신 행렬 곱만으로 UVᵀ 에 다가가고 16비트(bfloat16)로도 안정하게 도는 Newton–Schulz 반복을 쓴다.

이름은 두 갈래에서 왔다. 하나는 뉴턴법이다. 수 a 의 역수 1/a 를 나눗셈 없이 구하는 뉴턴법 x ← x(2 − ax) 를 행렬에 옮기면 X ← X(2I − AX)(I 는 단위행렬)가 되고, 슐츠(G. Schulz)가 1933년에 행렬의 역을 이렇게 곱셈만 되풀이해 구하는 반복을 냈다. 같은 생각으로 행렬을 가장 가까운 직교 행렬(특이값이 모두 1인 행렬)로 다듬는 반복은 1970년 무렵 코바릭(Kovarik), 비에르크(Björck)와 보위(Bowie)의 논문에 나온다. Muon 이 쓰는 것은 그 반복의 5차 판이다.

\begin{gathered} \textcolor{#786800}{X}_{k+1} = \textcolor{#7888e8}{a}\,\textcolor{#786800}{X}_k + \textcolor{#7888e8}{b}\,(\textcolor{#786800}{X}_k \textcolor{#786800}{X}_k^{\top})\,\textcolor{#786800}{X}_k \\ \qquad\quad + \textcolor{#7888e8}{c}\,(\textcolor{#786800}{X}_k \textcolor{#786800}{X}_k^{\top})^2 \textcolor{#786800}{X}_k \\ \textcolor{#786800}{X}_0 = \frac{\textcolor{#c05080}{M}}{\|\textcolor{#c05080}{M}\|_F} \end{gathered}

\begin{array}{ll} \textcolor{#786800}{X}_k & k\text{번째 반복의 행렬} \\ \textcolor{#c05080}{M} & \text{모멘텀을 누적한 기울기 행렬} \\ \|\cdot\|_F & \text{프로베니우스 노름} \ (\ge \|\cdot\|_2) \\ \textcolor{#7888e8}{a},\ \textcolor{#7888e8}{b},\ \textcolor{#7888e8}{c} & \text{Muon: } (3.4445,\ -4.7750,\ 2.0315) \\ & \text{고전 3차: } (1.5,\ -0.5,\ 0) \end{array}

각 반복은 특이벡터는 그대로 두고 특이값 σ 만 다항식 aσ + bσ³ + cσ⁵ 로 옮긴다. 고전적인 3차 반복(1.5σ − 0.5σ³)은 σ 를 정확히 1로 보내지만 작은 σ 에서 느리다. 작은 σ 를 한 번에 몇 배로 키우는지는 첫 계수 a 가 정한다. 조던은 반복 끝의 값이 1 둘레 0.7~1.3 어딘가에 머물러도 학습 손실 곡선에 해가 없다는 것을 실험으로 보고, 그 범위를 지키는 한에서 a 를 가장 크게 하는 계수를 찾았다. 다섯 번 반복하면 충분했다. 그리고 Muon이 직교화하는 것은 날것의 기울기가 아니라 (네스테로프) 모멘텀을 누적한 업데이트다.

적용 범위도 정해져 있다. Muon은 은닉층의 2차원 가중치 행렬에만 쓴다. 특이값은 행렬에만 있으니 편향 같은 벡터·스칼라 파라미터는 애초에 대상이 아니다. 임베딩과 출력 헤드는 2차원이지만 AdamW로 학습한다. 임베딩은 입력이 한 칸만 1 인 벡터(원-핫 벡터)라서 은닉층과 다른 노름이 맞는다는 이론(라지(Large) 외 2024 의 모듈러 노름)이 있고, 출력 헤드를 AdamW 로 두는 것은 이론에서 나온 것이 아니라 실험으로 정한 것이라고 조던은 적었다.

규모를 키우면

문샷 AI(Moonshot AI)의 “Muon is Scalable for LLM Training”(류(Jingyuan Liu) 외 2025)은 가중치 감쇠를 더하고 업데이트의 RMS(제곱 평균의 제곱근, 성분의 전형적인 크기)를 AdamW와 맞추는 조정을 넣어, 계산 최적 학습에서 AdamW 대비 약 2배의 계산 효율을 보고했다. 같은 회사의 Kimi K2(2025)는 총 1조 파라미터 가운데 토큰 하나에 실제로 쓰이는 것은 320억인 전문가 혼합(mixture of experts, 입력마다 일부 하위 신경망만 쓰는 구조) 모델을 Muon에 QK-Clip(어텐션 로짓 폭주를 막는 장치)을 더한 MuonClip으로 15.5조 토큰 동안 손실 스파이크 없이 학습했다고 밝혔다.

파이썬

Muon 절에서 본 특이값 10, 1, 0.05 의 기울기를 4×3 행렬로 만들어 두 방식으로 직교화한다.

import numpy as np
rng = np.random.default_rng(0)
Uq, _ = np.linalg.qr(rng.standard_normal((4, 3)))
Vq, _ = np.linalg.qr(rng.standard_normal((3, 3)))
G = Uq @ np.diag([10.0, 1.0, 0.05]) @ Vq.T
U, s, Vt = np.linalg.svd(G, full_matrices=False)

def newton_schulz(M, steps, a, b, c):
    X = M / np.linalg.norm(M)                   # 프로베니우스 노름으로 나눠 특이값을 1 아래로
    for _ in range(steps):
        A = X @ X.T
        X = a * X + (b * A + c * A @ A) @ X
    return X

for name, coef in [("3차 (1.5, −0.5)", (1.5, -0.5, 0.0)), ("Muon 5차", (3.4445, -4.7750, 2.0315))]:
    for k in [5, 10]:
        X = newton_schulz(G, k, *coef)
        print(f"{name:14s} {k:2d}회  특이값 {np.linalg.svd(X, compute_uv=False).round(3)}")
print("UVᵀ 의 특이값        ", np.linalg.svd(U @ Vt, compute_uv=False).round(3))
# 3차 (1.5, −0.5)  5회  특이값 [1.    0.656 0.038]
# 3차 (1.5, −0.5) 10회  특이값 [1.    1.    0.281]
# Muon 5차         5회  특이값 [1.157 0.709 0.702]
# Muon 5차        10회  특이값 [1.117 1.105 0.694]
# UVᵀ 의 특이값         [1. 1. 1.]

프로베니우스 노름으로 나누면 특이값이 (0.995, 0.0995, 0.005) 에서 출발한다. 3차 반복은 열 번을 돌려도 가장 작은 특이값을 0.28 까지밖에 못 올린다. Muon의 5차 반복은 다섯 번 만에 셋 모두를 0.7~1.2 사이로 모은다. 정확한 UVᵀ 는 아니지만, 모든 방향이 비슷한 크기로 움직이기에는 충분하다.

수확

“Newton–Schulz 반복은 특이벡터는 그대로 두고 특이값만 다항식으로 옮겨, 행렬 곱만으로 UVᵀ 에 다가간다. 고전 3차 반복은 정확하지만 작은 특이값에서 느리고, Muon의 5차 계수는 정확히 1에 앉는 대신 빨리 1 둘레로 모은다.”

문제 14. Newton–Schulz를 세 번

G = U·diag(3, 1)·Vᵀ 인 2×2 행렬에 고전적 3차 Newton–Schulz X ← 1.5X − 0.5XXᵀX 를 세 번 적용해 UVᵀ 를 얻으려 한다. (가) 노름으로 나누지 않고 X₀ = G 로 시작하면 어떻게 되는가? (나) 가장 큰 특이값으로 나누면 처음부터 σ₁ = 1 인데, 왜 굳이 프로베니우스 노름으로 나누는가? (다) 프로베니우스 노름으로 나누고 시작하면 세 번이면 충분한가? 위젯의 「문제 14 불러오기」로 이 값을 넣어 볼 수 있다.

함께 풀기

X₀ = G 로 넣고 세 번 돌렸는데 행렬이 엄청나게 커져요. 가장 큰 특이값이 2천만이 넘어요.

특이값 하나만 따라가 봐요. σ = 3 이면 한 번 뒤에 얼마예요?

1.5·3 − 0.5·27 = −9. 그다음 351, 그다음 −2천만… 폭발이네요.

반복이 UVᵀ 로 가는 거라면서요. 왜 멀어져요?

f(σ) = 1.5σ − 0.5σ³ 이 1로 끌어당기는 범위가 정해져 있어. 0 < σ < √3 이면 1로 가고, σ 가 √3 을 넘으면 부호가 바뀌면서 튕겨 나가.

그래서 X₀ 를 어떻게 잡아요?

가장 큰 특이값으로 나누면 되는데… 그러려면 SVD를 해야 하잖아요. 그걸 피하려고 반복을 쓰는 건데.

프로베니우스 노름으로 나누면 돼. ‖G‖_F = √(Σσᵢ²) ≥ σ₁ 이니까 모든 특이값이 1 아래로 들어가. 성분 제곱합만 구하면 되니까 SVD 도 필요 없고. 여기서는 √10 으로 나눠서 (0.949, 0.316).

0.316 이면 본문 파이썬에서 느렸던 0.005 보다 훨씬 크니까, 세 번이면 둘 다 1이에요.

계산해 봐요.

0.949 는 0.996, 1.000 으로 금방 가는데, 0.316 은 0.459, 0.640, 0.829 예요. 세 번으로는 0.83 에서 멈춰요. 다섯 번은 돌려야 0.997 이에요.

수렴 범위 안에 있다는 것과 몇 번 안에 도착하는지는 다른 질문이에요. 작은 σ 근처에서 f(σ) ≈ 1.5σ 라 한 번에 1.5배밖에 안 커져요.

그래서 Muon은 5차 다항식으로 작은 σ 에서의 기울기를 3.4 배쯤으로 키운 거군요. 대신 정확히 1에 앉지 않고 1 둘레에서 출렁이게 두고요.

과제 채점을 "정답에 가까울수록 1.5배 가산"으로 하면, 이미 잘한 사람은 금방 만점인데 거의 안 한 사람은 몇 번을 돌려도 안 따라온다… 그런 느낌이요.

문제 15. Muon 은 계산을 얼마나 더 하는가

m × n 가중치 행렬(m ≤ n)에 Newton–Schulz 를 한 번 돌리는 데 드는 곱셈·덧셈 수는 많아야 6nm² 이고, 이 층이 학습 한 스텝(순전파와 역전파)에 쓰는 계산은 6nmB 다. B 는 그 스텝에 들어가는 토큰 수다. (가) m = n = 4096 인 층, 반복 T = 5 번, B = 2²⁰(약 105만) 토큰이면 Muon 이 더하는 계산은 원래 계산의 몇 % 인가? (나) 미세조정처럼 B = 8192 토큰이면? (다) Muon 의 추가 계산은 무엇에 따라 커지는가?

함께 풀기

한 스텝에 층마다 행렬 곱을 다섯 번씩 더 하잖아요. 학습이 적어도 두세 배는 느려지겠죠.

원래 한 스텝의 계산과 나눠 봐요.

추가는 T × 6nm², 원래는 6nmB 니까 비는 T·m/B. 5 × 4096 / 1048576 = 0.0195, 2% 도 안 돼요. 다섯 번이라도 행렬 하나에 한 번씩이고, 원래 계산은 토큰 백만 개가 그 행렬을 하나하나 거치니까요.

그럼 Muon 은 어디서나 2% 짜리 덤이네.

배치를 8192 토큰으로 줄이면요?

5 × 4096 / 8192 = 2.5, 250% 예요. 덤이 원래 계산보다 커요. 비율에 B 가 분모로 들어가 있으니 배치가 작을수록, 폭 m 이 클수록 덤이 커져요. 「2%」는 배치가 아주 큰 사전학습에서나 맞는 말이에요.

조던 블로그에도 이 계산이 있어요. 언어모델 속도 경주 설정(m = 768, B = 524288)에서 0.7%, 라마 405B 학습(m = 16384, B = 1600만)에서 0.5% 라서 1% 아래라고요.

조교가 채점 기준표를 만드는 데 한 시간이 걸려도, 답안이 300장이면 장당 12초고 3장이면 장당 20분인 거랑 같네요.

문제 16. 누적 없는 Shampoo는 UVᵀ 다 (킬러)

(가) G = UΣVᵀ (Σ 가역)일 때 (GGᵀ)^−1/4 G (GᵀG)^−1/4 = UVᵀ 임을 보여라. (나) G 가 4×3 이라 GGᵀ 가 특이할 때 이 식은 무슨 뜻인가? (다) 수치로 계산할 때 무엇을 조심해야 하는가?

함께 풀기

(가)는 GGᵀ = UΣ²Uᵀ 라서 (GGᵀ)^−1/4 = UΣ^−1/2Uᵀ, 같은 식으로 (GᵀG)^−1/4 = VΣ^−1/2Vᵀ. 곱하면 UΣ^−1/2·Σ·Σ^−1/2Vᵀ = UVᵀ.

여기서 U 는 몇 × 몇이에요?

U 가 정사각 직교행렬이라고 가정했어요. G 가 4×3 이면… GGᵀ 는 4×4 인데 랭크가 3 이하라 고유값 하나가 0이에요. 0의 −1/4 제곱은 없으니까, (나)의 답은 "정의되지 않는다, 세로로 긴 행렬에는 못 쓴다"예요.

그 0 고유값의 고유벡터는 G 와 어떤 관계예요?

GGᵀ 의 영공간이니까 G 의 열공간에 수직이에요. 아… G 의 열은 그 방향 성분이 아예 없어요. 그럼 그 방향에 무엇을 곱하든 G 에는 안 닿아요.

0 고유값을 그냥 버리는 유사역(pseudo-inverse) −1/4 제곱을 쓰면 되네요. 결과는 컴팩트 SVD 의 U_rV_rᵀ 이고요. 못 쓰는 게 아니라, 뜻을 "G 가 사는 부분공간 위에서"로 좁히면 돼요.

코드로 해 봤는데요. 고유값 분해(np.linalg.eig)해서 w**−0.25 를 하니까 결과가 전부 nan 이에요.

고유값을 찍어 봐요.

[100.0, 1.0, 0.0025, −6.9e−15]. 0이어야 할 게 음수로 나왔어요. 반올림 오차로요. 음수의 −1/4 제곱이라 nan 이고요.

절댓값을 씌우면 nan 은 사라지는데, 6.9e−15 의 −1/4 제곱이면 3천 배가 넘게 증폭돼요. 이번엔 그 방향 성분이 워낙 작아서 UVᵀ 와 차이가 8e−10 정도로 티가 안 났지만, 성분이 조금만 커도 잡음이 수천 배로 튀어요.

그럼 어떻게 해요?

문턱 이하 고유값은 버리면 돼요. 1e−10 아래를 0으로 두니까 UVᵀ 와 차이가 1.1e−12 예요.

실제 Shampoo 는 L + εI 처럼 작은 값을 더하지. 그런데 그러면 정확히 UVᵀ 가 아니야.

확인했어요. ε = 1e−6 이면 특이값이 (1, 1, 0.9998) 인데, ε = 0.01 이면 (1, 0.995, 0.447) 이에요. 특이값 0.05 인 방향은 σ² = 0.0025 가 ε 보다 작아서 거의 안 펴져요.

ε 이 "이보다 작은 특이값은 믿지 않겠다"는 문턱이네요. 그 아래 방향은 직교화하지 않고 원래 크기 쪽에 남겨 두는 거예요.

둘 다 같은 식의 다른 구멍을 찾았어요. 서연 학생은 성립 조건을, 민준 학생은 수치를. 정리하면 이래요. 누적 없는 Shampoo는 G 가 사는 부분공간 위에서 정확히 UVᵀ 다. 0 고유값은 버려야 하고, 반올림이 그것을 음수나 아주 작은 양수로 만들기 때문에 문턱이나 작은 값 더하기가 필요하다. 그리고 더한 값 ε 의 크기가 어떤 특이값까지 펴 줄지를 정한다. Newton–Schulz는 이 모든 것을 행렬 곱만으로, 고유값 분해 없이 해낸다는 데 값어치가 있어요.

해석학에서 "역함수는 치역 위에서만 정의된다"를 매번 강조하던 이유를 알겠어요. 정의역을 좁히는 게 반칙이 아니라 정의의 일부예요.

조교가 "제출 안 한 과제는 0점 처리가 아니라 평균에서 제외"라고 하던 거랑 같네요. 없는 걸 0으로 나누지 말고 빼 두는 거.