Tabular Foundation Model은 왜 Muon을 쓰는가 : Norm을 고르는 문제

Norm을 고르면 Dualizer가 정해지고, Dualizer가 모델 학습의 업데이트 방향을 정한다. Adam과 Muon의 차이는 학습률이 아니라 weight를 독립적인 숫자 모음으로 볼지 선형변환으로 볼지의 차이다.

Photo by Afif Ramdhasuma / Unsplash

Tabular Foundation Model의 대표주자들 중 하나인 TabICL은 optimizer로 AdamW가 아니라 Muon을 쓴다. Tabular foundation model도 결국 표를 in-context로 처리하는 transformer이니 LLM 쪽에서 검증된 optimizer를 그대로 가져오는 게 이상할 일은 아니다. 그런데 왜 하필 Muon인지, 그리고 왜 그 이득이 tabular에서 특별히 클 만한지는 업데이트 공식만 들여다봐서는 잘 보이지 않고 이에 대해서 아직 논의하진 않은 상태이다. 이를 이해하기 위해서는 Optimizer를 어떤 Norm 위에서의 Steepest descent로 읽는 관점이 필요하다.

Muon이 하는 일

Muon(MomentUm Orthogonalized by Newton-Schulz)은 모멘텀을 쌓은 gradient를 orthogonalize해서 업데이트로 쓰는 optimizer다. 보통의 SGD-momentum처럼 gradient $G_t$에 모멘텀을 쌓고($M_t = \mu M_{t-1} + G_t$), 그 행렬의 SVD $M_t = U\Sigma V^\top$에서 특이값을 전부 1로 만든 $UV^\top$로 방향을 바꿔 업데이트한다.

$$W_t = W_{t-1} - \eta\cdot \text{ortho}(M_t)$$

$$\text{ortho}(M_t) = UV^\top \quad (M_t = U\Sigma V^\top)$$

매 스텝 SVD를 직접 계산하는 건 너무 오래 걸린다. 그래서 실제로는 Newton-Schulz 근사를 쓴다. bf16 행렬곱으로 quintic iteration을 다섯번 돌리면 $UV^\top$에 충분히 가까워진다.

왜 Newton-Schulz 근사가 $UV^\top$을 근사하는지에 대해선 이 글의 논지에 벗어나 논의하지 않겠으나, 이 글에 구체적으로 설명되어 있다.

여기서 특이값을 평평하게 펴는 이 연산이 왜 좋은 업데이트인지는 정의만 봐서는 드러나지 않는다. 이걸 보려면 optimizer를 업데이트 공식이 아니라 norm으로 봐야 한다.

업데이트를 고르는 문제

파라미터를 업데이트한다는 건 결국 $\arg\min_{\Delta W}\mathcal{L}(W + \Delta W)$를 구하는 일이다. $\mathcal{L}(W + \Delta W)$를 2차까지 테일러 전개해보면

$$\mathcal{L}(W + \Delta W) = \mathcal{L}(W) + \langle\nabla\mathcal{L}(W), \Delta W\rangle_F + \frac{1}{2}\langle\Delta W, H(W)\Delta W\rangle_F + \cdots$$

가 된다. 여기서 2차 텀을 Hessian 대신 $\Delta W$의 크기에 대한 penalty로 취급해버리면

$$\Delta W_{\text{target}} = \arg\min_{\Delta W}\left\lbrace \langle\nabla\mathcal{L}(W), \Delta W\rangle_F + \frac{\lambda}{2}\lVert\Delta W\rVert^2 \right\rbrace$$

라는 문제가 남는다. gradient와의 내적을 줄이면서 업데이트 크기에 penalty를 주는 이 방식을 Steepest Descent라 한다.

Dual norm이 방향을 결정한다

Bernstein et al.의 정리를 빌리면, optimizer에서 exponential moving average를 걷어냈을 때 남는 것은 특정 norm 아래의 steepest descent다. 위 문제의 해는

$$\Delta W_{\text{target}} = -\frac{\lVert\nabla\mathcal{L}(W)\rVert^{\dagger}}{\lambda}$$

$$\text{dualizer} _ {\lVert\cdot\rVert}(\nabla\mathcal{L}(W)), \qquad \text{dualizer} _ {\lVert\cdot\rVert}(X) = \arg\max_{\lVert T\rVert = 1}\langle X, T\rangle_F$$

이고 $\lVert\cdot\rVert^{\dagger}$는 dual norm이다. dualizer를 읽어보면 단위구 위에서 gradient와 가장 잘 정렬된 방향을 찾는 연산이다. 그러니까 어떤 norm을 쓰느냐가 단위구의 모양을 정하고, 단위구의 모양이 업데이트 방향을 정한다.

Norm을 바꾸면 Optimizer가 바뀐다

익숙한 norm들을 위 식에 넣어보면 우리가 쓰는 optimizer들이 그대로 나온다.

norm	Euclidean	$\ell_\infty$ (좌표별 max)	spectral	Schatten-$p$
단위구 모양	공(sphere)	상자(box)	특이값 $\le 1$	앞의 둘 사이
dualizer	$G/\lVert G\rVert$	$\text{sign}(G)$	$UV^\top$	$\sigma_i^{q-1}$
optimizer	SGD	Adam	Muon	SGD ↔ Muon
weight를 보는 관점	값의 벡터	숫자 모음	선형변환	선형변환

Euclidean norm(행렬에서는 Frobenius norm)을 쓰면 dualizer가 gradient의 크기만 1로 맞춰 돌려주니 방향은 $G$ 그대로고, 결국 SGD로 귀결된다. Adam이 사실상 sign gradient descent라는 이야기도 여기서 나온다. $\ell_\infty$의 단위구는 상자라서 dualizer가 각 좌표를 독립적으로 모서리까지 밀어붙인다. 좌표를 따로따로 보는 만큼, 그 좌표들이 모여 만드는 행렬 구조는 보지 않는다. 반대로 spectral norm의 단위구는 특이값이 1 이하인 행렬들의 집합이고, 그 안에서 gradient와 가장 정렬된 원소는 특이값을 모두 1로 채운 $UV^\top$이다. Muon의 orthogonalize는 결국 이 norm에서 유도되는 답이었던 것이다.

$\ell_\infty$와 spectral의 단위구와 각 norm의 dualizer 후보.

이 차이가 업데이트 방향에서 어떻게 갈리는지 위의 간단한 예시를 통해 살펴보자. $2\times2$ 행렬에서 자유도를 둘만 남긴 단면($\begin{pmatrix} a & -b \cr b & a \end{pmatrix}$ 꼴)을 쓰면 두 norm의 단위구를 같은 평면에 놓고 볼 수 있다. $\ell_\infty$ 쪽에서 dualizer가 고를 수 있는 답은 상자의 네 모서리 (꼭짓점) 뿐이다. gradient가 어디를 가리켜도 업데이트는 그중 하나로 향하고, 그림에서 $G$와 $\text{sign}(G)$의 방향이 벌어지는 결과가 생긴다. 이에 반해 spectral은 특이값이 전부 1인 행렬, 즉 orthogonal 행렬 전체가 후보다. 이 후보들은 끊어져 있지 않고 연속으로 이어져서 $UV^\top$는 원 위 어디든 갈 수 있고, 그래서 gradient가 회전하면 업데이트도 같이 회전한다.

이 단면에서 spectral 단위구가 동그랗게 나오는 건 단면을 그렇게 잘랐기 때문이다. 일반적인 행렬 공간에서 spectral 단위구는 공 모양이 아니고, 위 표가 Euclidean과 spectral을 다른 칸에 둔 것도 spectral이 애초에 non-Euclidean norm이라서다. 위 예시 단면에서 봐야 할 부분은 dualizer의 후보가 유한개냐 연속이냐뿐이다.

여기까지 놓고 보면 Muon은 Adam과 다른 종류의 optimizer가 아니다. $\ell_\infty$는 행렬의 성분을 상자에 담아 성분마다 부호를 뒤집고, spectral은 특이값을 상자에 담아 특이값을 전부 1로 편다. 담는 대상을 성분에서 특이값으로 옮긴 Adam이 Muon이라고 읽어도 크게 틀리지 않는다.

Schatten-p norm

표의 마지막 줄을 위해 Schatten-p norm을 잠시 살펴보자면,

$$\lVert A\rVert_{S_p} := \left( \sum_{i=1}^{\min(m, n)} |\sigma_i(A)|^p \right)^{1/p}$$

여기서 $\sigma_i(A)$는 $A$의 특이값이다. $p$ 값에 따라 알고 있는 norm들이 나온다. $p=1$이면 nuclear norm, $p=2$면 Frobenius norm, $p=\infty$면 특이값의 최댓값인 spectral norm이다. 그리고 Schatten-$p$의 dual norm은 Schatten-$q$이고, 이때 $p$와 $q$는 $1/p + 1/q = 1$을 만족하는 conjugate exponent다.

$p$를 2에서 $\infty$로 올리는 건 SGD에서 Muon으로 연속적으로 옮겨가는 것과 같다. 그 사이 구간의 dualizer는 특이값을 그대로 두지도, 완전히 펴지도 않고 $\sigma_i^{q-1}$로 리스케일한다.

대규모로 가면 보정이 필요하다

Muon을 대규모 학습에 그대로 얹으면 오히려 AdamW보다 성능이 떨어진다. Moonshot(Kimi 모델을 만드는 팀)이 이걸 발견하고 두 가지 보정을 붙여 대규모에서도 작동하는 Muon을 만들었다. 이 글에서만 Moonlight Muon이라 부르자.

하나는 decoupled weight decay다. Adam에서 AdamW로 넘어간 것과 같은 이유로, 대규모에서 weight와 output의 norm이 계속 자라는 것을 억제한다.

다른 하나는 update RMS 스케일 매칭이다. orthogonalize된 업데이트의 RMS는 행렬 shape마다 제각각이라 여러 파라미터가 하나의 학습률을 공유하기 어렵다. 그래서 파라미터별로 $0.2\sqrt{\max(n,m)}$을 곱해서 AdamW의 전형적인 update RMS(대략 0.2) 근처로 맞춘다. 이 보정 덕분에 AdamW용으로 이미 튜닝해둔 학습률을 거의 그대로 재사용할 수 있다.

$$W_t = W_{t-1} - \eta\left(0.2\sqrt{\max(n,m)}\cdot\text{ortho}(M_t) + \lambda W_{t-1}\right)$$

TabICL이 쓰는 Muon도 이 Moonlight Muon이고, 하이퍼파라미터까지 거의 그대로 가져다 쓴다.

MuonClip과 그 이후

From Kimi Team, Kimi K2: Open Agentic Intelligence

Moonlight Muon으로 학습해도 step을 길게 가면 attention logit이 터지는 현상이 남았다. Kimi K2를 발표하면서 Moonshot 팀이 여기에 QK-Clip을 더한 것이 MuonClip이다.

MuonClip은 head별 최대 attention logit $S^h_{max}$가 threshold $\tau = 100$을 넘으면 그 head의 query와 key weight를 각각 $\sqrt{\tau/S^h_{max}}$로 rescale한다. Q와 K에 각각 $\sqrt{\tau/S}$가 곱해지니 둘이 만드는 logit은 정확히 $\tau/S$배가 되어 upper bound가 $\tau$ 가 된다. 목표 logit 스케일을 $\gamma_h = \min(1, \tau/S^h_{max})$로 두고 각 weight에는 $\sqrt{\gamma_h}$를 곱하는 셈이다.

Kimi K3는 이 안정화를 사후 보정으로 남겨두지 않고 optimizer 구조 안으로 흡수하는 Per-Head Muon으로 갔다. attention head 단위로 독립 최적화하면서 head별 learning rate scheduling까지 주는 방식이고, projection 행렬 전체에 Newton-Schulz를 돌리는 것보다 오히려 저렴하다고 한다.

Tabular에서는 어떻게 읽을 수 있을까

앞의 dualizer 이야기를 tabular foundation model에 대입해보자. 그러면 AdamW냐 Muon이냐는 이 모델의 weight 행렬에 어떤 geometry가 맞느냐의 문제가 된다.

From Franz Louis Cesista(leloykun), Deep Learning Optimizers as Steepest Descent in Normed Spaces

TabPFN이나 TabICL 계열은 표를 in-context로 처리하는 transformer이고, 그 weight 행렬들은 우리가 명시적으로 normalize해둔 activation 공간 사이를 잇는다. embedding마다 LayerNorm이나 RMSNorm이 걸려 있다는 뜻이다. leloykun의 글에서는, 입출력이 Euclidean으로 정규화되어 있으면 그 사이를 잇는 weight에 유도되는 norm은 spectral norm이다. 그렇다면 그 weight를 업데이트하는 steepest descent의 올바른 dualizer는 $UV^\top$, 즉 Muon이고, 원소별 sign을 쓰는 AdamW는 행렬의 행과 열 방향 구조를 무시하는 선택이 된다. 정규화가 잘 된 transformer에서 Muon은 여러 후보 중 하나가 아니라 norm 기하학에서 유도되는 선택인 것이다.

이득의 크기까지 생각하면 tabular 쪽이 특히 유리해 보이는 이유가 있다. tabular in-context learning은 스케일도 의미도 제각각인 컬럼과 드물게만 등장하는 방향을 한 모델이 다 표현해야 한다. 그런데 transformer gradient는 condition number가 커서 소수의 방향이 업데이트를 지배한다. Schatten-p를 높여 dualizer가 특이값을 펴면 업데이트의 stable rank가 커지고, 과소대표된 드문 방향까지 학습에 실린다. weight-erasure가 완화되는 것이다. 이걸 tabular in-context learning으로 옮겨 읽으면, 지배적인 패턴에 붕괴하지 않고 드문 feature나 sample 상호작용까지 attend하도록 배우는 쪽에 가깝다.

norm을 고른다는 것

Adam은 좌표별 max norm 아래의 steepest descent이고, Muon은 spectral norm 아래의 steepest descent다. 둘의 차이는 학습률이나 모멘텀 같은 하이퍼파라미터가 아니라 weight를 독립적인 숫자 모음으로 볼지 선형변환으로 볼지의 차이다. optimizer를 고른다는 건 이 모델의 weight가 어떤 공간에 사는지를 고르는 일에 가깝다.

Tabular foundation model은 activation을 정규화해둔 transformer이고, 이질적인 컬럼과 드문 방향을 함께 다뤄야 한다. 두 조건 모두 spectral 쪽이 max보다 해석하기 유의하다. TabICL이 Moonlight Muon을 쓰는 것도 그 방향의 신호로 보인다.