Muon 옵티마이저 원리 정리: 모멘텀 직교화, SVD 없는 다항식 반복, QK 클립까지

핵심 메시지

• Muon은 모멘텀 행렬을 직교화하는 옵티마이저로, 소형 언어 모델에서 AdamW 대비 약 두 배의 계산 효율을 보이며 주목받고 있다.
• Adam은 파라미터마다 두 개의 이동평균을 유지해 옵티마이저 상태가 모델 자체의 약 두 배 메모리를 차지하고, 모든 파라미터를 하나의 긴 벡터로만 취급한다.
• 선형 계층의 모멘텀 행렬은 실제로 거의 저계수가 되어 소수의 지배적 방향만 학습을 끌고 가는데, 직교화는 드물게 갱신되는 방향의 영향을 키운다.
• 매 스텝 SVD를 계산하는 것은 너무 비싸서, 홀수 차수 다항식을 다섯 번 반복해 특이값을 1 근처로 밀어 넣는 방식으로 대체한다.
• 모델을 키우면 어텐션 로짓이 계속 커져 학습이 불안정해지는데, 임계값을 넘을 때 쿼리·키 가중치를 되스케일하는 QK 클립으로 이를 억제한다.

쉽게 이해하기

기계학습에서 좋은 파라미터를 찾는 일을 최적화라고 하고, 그 도구가 옵티마이저다. 오랫동안 기본 선택지는 Adam이었다. Adam은 경사하강법 위에 두 가지 지수이동평균을 얹는다. 하나는 과거 기울기 자체인 모멘텀이고, 다른 하나는 기울기의 제곱이다. 원래 기울기로 바로 갱신하는 대신 모멘텀에 적응적 스케일을 곱해 쓰기 때문에 수렴이 빠르고 결과도 대체로 낫다. 대신 파라미터 하나마다 변수를 두 개씩 더 들고 있어야 해서 옵티마이저 상태가 모델 크기의 약 두 배를 차지한다.

더 근본적인 한계는 Adam이 모든 파라미터를 하나의 긴 벡터로 보고 값마다 따로 갱신한다는 점이다. 신경망에서 가장 흔한 선형 계층의 가중치는 사실 2차원 행렬인데, 벡터 기반 옵티마이저를 쓰면 이 행렬 모양의 모멘텀이 실제로는 거의 저계수 상태가 되는 경향이 있다. 소수의 지배적인 방향만 갱신을 주도하고 나머지 방향은 거의 기여하지 못한다는 뜻이다. Muon은 이 모멘텀 행렬을 직교화해서, 작게 또는 드물게 갱신되던 방향의 영향을 키운다. 사소해 보이는 이 방향들이 데이터의 미묘한 패턴을 잡는 데 결정적일 수 있기 때문이다.

직교화란 주어진 모멘텀 행렬 M에 가장 가까우면서 행과 열이 서로 직교하는 행렬 O를 찾는 일이다. 이 문제를 푸는 정석은 특이값 분해다. 어떤 선형 변환이든 회전, 축별 확대·축소, 다시 회전의 세 단계로 쪼갤 수 있다는 것이 특이값 분해이고, 가운데 대각 행렬의 특이값을 전부 1로 만들면 원하는 직교 행렬을 얻는다. 문제는 계산 비용이다. 학습의 매 갱신마다 특이값 분해를 돌릴 수는 없다.

그래서 쓰는 우회로가 홀수 차수 다항식이다. 행렬을 넣으면 자기 자신과 자기 전치를 곱한 항들의 가중합을 돌려주는 이 함수를 특이값 분해 형태로 풀어 보면, 회전 행렬들은 소거되고 결국 각 특이값에 같은 다항식을 따로 적용한 것과 같아진다. 즉 특이값 분해를 실제로 계산하지 않고도 특이값만 조작할 수 있다. 계수를 잘 고르면 0과 1 사이의 입력값이 반복 적용마다 1 쪽으로 끌려가고, 다섯 번쯤 반복하면 거의 1에 붙는다. 게다가 특이값이 정확히 1일 필요도 없고 대략 0.7에서 1.3 사이에만 들어오면 충분하다는 점이 실험적으로 확인되면서, 수렴을 더 빠르게 하는 계수 선택이 가능해졌다. 각 반복은 행렬 곱뿐이라 GPU에서 효율적으로 돌아간다.

다만 모델을 키우면 AdamW 대비 이득이 줄어드는 문제가 있어, 가중치 감쇠를 추가하고 2차원 행렬의 크기를 반영해 학습률을 조정하는 보완이 들어간다. 여기에 더해 학습이 진행될수록 어텐션 로짓이 점점 커져 불안정해지는 현상이 관찰된다. 임베딩 벡터가 보통 단위 노름으로 정규화되어 있으므로 로짓의 크기는 쿼리·키 가중치 행렬의 스케일에 달려 있다. 그래서 학습 중 로짓 최댓값을 감시하다가 임계값을 넘으면 비율을 계산해 두 행렬을 각각 그 비율의 제곱근으로 줄인다. 이것이 QK 클립이다. 멀티헤드에서는 헤드마다 별도의 스케일을 적용하고, 저계수 압축을 쓰는 멀티헤드 잠재 어텐션에서는 헤드별 회전 쿼리만 조정하고 모든 헤드가 공유하는 회전 키 행렬은 건드리지 않는다. 같은 행렬이 여러 번 축소되는 것을 막기 위해서다.

주요 인사이트

• Adam이 빠른 대가로 치르는 비용은 메모리다. 파라미터당 두 개의 추가 변수는 모델 자체보다 큰 옵티마이저 상태를 만든다.
• 행렬을 벡터로 펴 버리는 순간 층의 내부 구조 정보가 사라진다. Muon의 출발점은 그 구조를 되살리는 것이다.
• 직교화의 목적은 균형이다. 지배적인 방향에 가려져 있던 드문 방향을 끌어올려 학습이 소수의 축에만 쏠리지 않게 만든다.
• 특이값을 정확히 1로 맞출 필요가 없다는 실험적 관찰이 결정적이다. 목표를 구간으로 완화한 덕분에 더 공격적인 계수로 빠르게 수렴시킬 수 있다.
• 옵티마이저 교체는 그 자체로 끝나지 않는다. 대형 모델에서는 가중치 감쇠, 행렬 크기를 반영한 학습률, 로짓 억제 같은 안정화 장치가 함께 따라붙는다.

자주 묻는 질문

관련 AI 소식