Polar Express: 반복마다 다른 다항식

Muon 은 뉴턴–슐츠 반복에 같은 5차 다항식 p(σ) = 3.4445σ − 4.7750σ³ + 2.0315σ⁵ 를 다섯 번 썼다. σ(시그마)는 반복에 들어가는 특잇값이다. 이 계수는 첫 계수를 가장 크게 해 작은 특잇값을 빨리 키우는 대신, 끝의 특잇값이 1 근처 0.7~1.2 어딘가에서 출렁이도록 허락한 것이었다. 그런데 다섯 번을 돌린다면 왜 다섯 번 모두 같은 다항식이어야 할까? 첫 반복에 들어오는 특잇값은 0 근처에서 1까지 넓게 퍼져 있지만, 몇 번 돌린 뒤에는 이미 1 근처에 모여 있다. 넓게 퍼진 특잇값을 끌어올리는 일과 1 근처에서 다듬는 일은 서로 다른 다항식이 더 잘할 것이다.

역사: 반복마다 가장 나쁜 경우를 줄이는 다항식

노아 암셀(Noah Amsel), 데이비드 퍼손(David Persson), 크리스토퍼 무스코(Christopher Musco), 로버트 가워(Robert M. Gower)는 2025년 5월 「The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm」을 arXiv 에 올렸다. 이들은 반복마다 「지금 특잇값이 있을 수 있는 구간 [ℓ, u] 전체에서 1과 가장 멀리 떨어진 값」을 가장 작게 하는 다항식을 새로 골랐다. 가장 나쁜 경우를 가장 작게 하는 미니맥스 문제를 반복마다 푼 것이다. 한 번 고르고 나면 다음 구간이 정해지므로, 다섯 다항식의 계수를 미리 계산해 표로 둘 수 있다. 이들은 16비트 실수(bfloat16)에서도 깨지지 않게 안전 여유를 넣었고, Muon 에 끼워 GPT-2 크기 모델을 학습시켰을 때 여러 학습률에서 검증 손실이 조금씩 좋아졌다고 보고했다. Dimuon 의 재현 결과는 이 계수에 1.01 의 안전 여유를 둔 것이었다.

입력 특잇값이 0.001 에서 1 사이 어디에나 있을 수 있을 때, 반복마다 특잇값이 있을 수 있는 구간. Muon(초록)은 위쪽 끝이 1.20 에 묶인 채 아래쪽 끝만 천천히 올라와 다섯 번 뒤 0.47~1.20 이다. Polar Express(파랑)는 중간에 구간 위쪽 끝이 2 근처까지 넓어지지만 다섯 번 뒤 0.84~1.12 로 좁아진다. 암셀 팀(2025)의 생각을 이 책에서 다시 그린 것(sources/dimuon/figs.py)

특잇값 0.995, 0.0995, 0.005 를 다섯 번

Muon 을 다룰 때 쓴 4×3 그래디언트(특잇값 10, 1, 0.05)를 다시 쓰자. 프로베니우스 노름으로 나누면 특잇값이 0.995, 0.0995, 0.005 에서 출발한다. Polar Express 는 입력을 1.01 배 더 나누고(0.985, 0.0985, 0.0049), 계수도 그에 맞춰 바꾼다(sources/dimuon/s1_polar_express.py).

반복	Muon: 방향 1, 2, 3	Polar Express: 방향 1, 2, 3
출발	0.995, 0.0995, 0.005	0.985, 0.0985, 0.0049
1번	0.705, 0.338, 0.017	1.461, 0.787, 0.040
2번	1.109, 0.989, 0.059	0.468, 1.962, 0.164
3번	0.715, 0.710, 0.202	1.553, 1.614, 0.629
4번	1.097, 1.103, 0.658	0.440, 0.463, 1.514
5번	0.702, 0.709, 1.157	0.871, 0.902, 0.997
‖X₅ − UVᵀ‖_F	0.445	0.162

Polar Express 의 중간 값은 Muon 보다 훨씬 크게 출렁인다. 2번 뒤 둘째 특잇값은 1.962 까지 올라가고, 4번 뒤 첫째 특잇값은 0.440 까지 내려간다. 그런데 다섯 번 뒤에는 세 특잇값이 0.87~1.00 에 모이고, 직교 행렬 UVᵀ 와의 거리는 Muon 의 3분의 1쯤이다. 앞의 다항식들은 작은 특잇값을 끌어올리는 데 힘을 쓰고, 그 사이 큰 특잇값이 크게 흔들려도 마지막 다항식들이 1 근처로 되돌려 놓는다.

왼쪽: 입력 특잇값마다 다섯 번 반복한 뒤 가는 곳. Muon(초록)은 0.68~1.13 사이를 오가고, Polar Express(파랑)는 0.001 위에서 0.84~1.12 에 머문다. 0.001 보다 작은 특잇값은 둘 다 아직 1에 닿지 못한다. 오른쪽: 한 번의 반복이 쓰는 다항식. Muon 의 다항식 하나와 Polar Express 의 다섯 다항식(옅은 파랑이 첫째)(sources/dimuon/figs.py)

직교화를 반복 몇 번으로 줄이는 계산은 「행렬의 부호 함수」나 「극분해」라는 이름으로 수치해석에서 오래 다뤄 온 문제이고, Polar Express 라는 이름도 극분해(polar decomposition)에서 왔다. Muon 은 그 문제에 「정확할 필요는 없고 1 근처면 된다」는 조건을 걸었고, Polar Express 는 같은 조건 아래에서 다섯 번의 행렬 곱으로 더 1에 가까이 가는 법을 찾았다.

문제 3. 0.83 을 0.2 로 떨어뜨리는 다항식

Polar Express 의 첫 다항식은 1.01 로 나눈 입력이 0.822 인 특잇값(원래 0.83)을 0.202 로 떨어뜨린다. Muon 의 다항식은 같은 입력을 0.929 로 보낸다. (가) 첫 반복만 보면 어느 쪽이 나은가? (나) 다섯 번 뒤에는 Polar Express 가 0.877, Muon 이 1.122 다. 0.005 에서 출발한 특잇값은 Polar Express 가 1.000, Muon 이 1.154 다. 첫 다항식은 무엇을 겨냥하고 있었는가? (다) 반복을 세 번에서 멈추면 어떤 일이 생기는가?

(가)만 보면 Muon 이 나아 보여. 0.83 이 0.93 이면 1에 가까워졌는데, 0.2 면 오히려 멀어졌잖아.

근데 표를 보면 0.005 는 Polar Express 첫 번에 0.041 로 여덟 배쯤 커져. Muon 은 0.017 이고. 첫 다항식은 0 에서의 기울기가 8.2 라서 작은 값을 확 키우는 대신, 큰 값 쪽에서 크게 휘는 것 같아.

그러면 첫 다항식을 고를 때 무엇을 가장 걱정했을까요? 다섯 번 뒤에 가장 1에서 먼 특잇값은 어떤 특잇값이죠?

가장 작은 특잇값이요. 큰 건 한 번만 크게 흔들려도 다음 다항식들이 1로 돌려놓을 수 있는데, 작은 건 다섯 번 안에 1까지 올라와야 해요. 0.005 가 Muon 에서도 마지막에 1.154 니까 겨우 닿은 거고요.

그러니까 첫 다항식은 0.83 같은 값이 0.2 로 떨어져도 괜찮다고 보고, 다음 반복에서 0.2 가 다시 0.8 근처로 올라온다는 걸 알고 고른 거네. 실제로 0.202 → 0.799 → 1.855 → 1.335 → 0.877 이야.

다섯 다항식은 한 묶음으로 고른 거예요. 각 다항식은 「지금 특잇값이 있을 수 있는 구간」을 받아 다음 구간을 가장 좁게 만들어요. 그러면 (다)는요?

세 번에서 멈추면… 0.83 은 1.855 고 둘째 방향은 1.614 예요. Muon 은 세 번 뒤에도 0.7~1.1 근처였는데, Polar Express 는 중간에 끊으면 오히려 더 나빠요.

등산로 지름길이 중간에 한 번 내려갔다 다시 올라가도 정상에는 먼저 닿는 것처럼, 끝까지 가야 빠른 거네요. 다섯 번이라는 횟수가 계수 안에 들어 있는 거고요.