Muon 직교화 업데이트, 행렬 이차식 분석으로 효과와 한계 규명

Muon의 핵심인 직교화 업데이트가 언제 효과를 내고 언제 느려지는지 수학적으로 따진 논문이 arXiv에 공개됐다. 논문은 행렬 이차식 문제에서 이 업데이트(PolarGD)를 분석했다. 학습률 조정의 영향을 없애려고 직선 탐색 아래서 갱신 방향을 비교한 결과, 기울기가 저랭크에서 멀고 가장 큰 특이값 방향의 곡률이 클 때 직교화가 이득이라고 밝혔다.

Muon은 행렬 형태의 파라미터를 대략 직교화한 방향으로 갱신하며, 대형 언어 모델 학습에서 AdamW보다 빠른 경우가 많다. 다만 직교화가 왜, 언제 도움이 되는지는 분명하지 않았다. 연구진은 최적화 과정을 닫힌 형태로 분석할 수 있는 행렬 이차식을 골랐다.

빠른 경우와 느린 경우

오차 행렬의 행들이 서로 직교하고 노름이 같으면 PolarGD는 헤시안과 상관없이 한 단계에 최솟값에 닿는다. 같은 조건에서 경사 하강법이 한 단계에 닿으려면 헤시안의 고윳값이 모두 같아야 한다. 부호 하강법과 달리 PolarGD는 문제를 회전시켜도 영향을 받지 않는다.

여러 단계에 걸쳐서는 PolarGD가 현재 반복값으로 전처리한 경사 하강법처럼 움직인다. 속도는 헤시안이 아니라 반복값의 조건수가 좌우한다. 반복값에 아주 작은 특이값이 있으면 경사 하강법보다 훨씬 느려지는 구간이 생긴다.

Muon이 쓰는 몇 번의 Newton-Schulz 반복은 이런 작은 특이값을 경사 하강법과 같은 방식으로 다룬다. 그래서 반복 횟수가 적으면 정확한 직교화보다 더 빨리 수렴할 수 있다. 트랜스포머를 Muon으로 학습시켰을 때 어텐션 행렬 특이값의 약 4분의 1이 작은 특이값 구간, 곧 PolarGD가 경사 하강법보다 훨씬 느려지는 구간에 해당했다.