🤖 machine learning

Approximate Muon with low-rank adapters

이 논문은 목적 함수를 선형화하고 이를 최소제곱법을 통해 해결함으로써, 기존에는 LoRA와 호환되지 않았던 SFT 및 ReLoRA 사전 학습과 같은 환경에서 Muon의 성능 이점을 가능하게 하는 저차원 매개변수 효율적 미세 조정(PEFT)을 위한 효율적인 방법인 sMuon을 소개한다.

원저자: Ben Anson, Conor Houghton, Edward Milsom

게시일 2026-08-17

📖 3 분 읽기☕ 가벼운 읽기

CC0 1.0

원저자: Ben Anson, Conor Houghton, Edward Milsom

원본 논문은 CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.0/)에 따라 공공 도메인에 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 거대하고 매우 똑똑한 로봇에게 시 쓰기나 수학 문제 풀기와 같은 새로운 기술을 가르치려 한다고 상상해 보세요. 로봇의 전체 뇌를 처음부터 다시 훈련시키는 것은 시간이 너무 오래 걸리고 비용도 엄청나게 들기 때문에, 그럴 수는 없습니다. 대신, 기존의 뇌에 작고 가벼운 "훈련 모듈"을 부착하고 싶습니다. 이것을 **매개변수 효율적 미세 조정(Parameter-Efficient Fine-Tuning, PEFT)**이라고 부릅니다. 이것은 범용 로봇에게 특화된 헤드셋을 씌워주는 것과 같습니다. 이 헤드셋은 로봇의 핵심 하드웨어를 변경하지 않고도 새로운 기술을 가르쳐 줍니다.

이러한 헤드셋을 만드는 인기 있는 방법 중 하나는 LoRA(Low-Rank Adaptation)입니다. 이것은 마치 두 장의 얇고 평평한 종이로 헤드셋을 만드는 것과 같아서, 이 종이들을 쌓았을 때 새로운 기술의 형태가 만들어집니다. 로봇을 가르치기 위해서는 헤드셋이 어떻게 모양을 바꾸어야 더 나아질 수 있는지 알려주는 수학적 코치인 **옵티마이저(optimizer)**가 필요합니다. 가장 흔한 코치는 AdamW이지만, 더 새롭고 강력한 코치인 Muon이 있습니다. Muon은 로봇에게 첫 번째 교훈을 주는 사전 학습(pretraining) 단계에서 믿을 수 없을 정도로 효율적인 것으로 유명합니다. 이는 로봇의 뇌 업데이트가 마치 무용수가 한 발로 완벽하게 균형을 잡으며 회전하는 것처럼, 완벽하게 균형 잡힌 "직교(orthogonal)" 방식으로 일어나도록 보장하기 때문입니다.

하지만 문제가 하나 있습니다. LoRA 헤드셋을 훈련할 때 Muon을 사용하려고 하면 수학적으로 오류가 발생합니다. Muon은 로봇 뇌의 전체 레이어를 균형 있게 맞추고 싶어 하지만, LoRA 헤드셋은 그 레이어의 아주 작은, 낮은 차원의 조각일 뿐입니다. 이것은 마치 오케스트라 전체를 지휘하기 위해 지휘봉을 휘두르는데, 정작 지휘는 단 한 명의 바이올리니스트에게만 적용되어 기하학적 구조가 맞지 않는 상황과 같습니다. 표준적인 해결책은 오케스트라는 무시하고 바이올리니스트에게 자기 혼자서만 돌라고 말하는 것이었지만, 이는 Muon의 진정한 묘미를 놓치는 일입니다. 이 논문은 이 질문을 던집니다: 수학적 구조를 깨뜨리거나 속도를 늦추지 않으면서도 Muon이 이 작은 헤드셋에서도 작동하게 만들 방법이 있을까요?

저자인 벤 안슨(Ben Anson), 코너 호튼(Conor Houghton), 에드워드 밀섬(Edward Milsom)은 그렇다고 답합니다. 그들은 sMuon(small Muon)이라는 새로운 방법을 소개합니다. 불가능한 것을 강요하는 대신, 그들은 문제를 "선형화(linearize)"하는 영리한 수학적 트릭을 사용합니다. 흔들거리는 거대한 우산을 들고 외줄 타기를 한다고 상상해 보세요. 우산 전체의 균형을 한꺼번에 잡으려 하는 대신, 아주 짧은 순간 동안 우산을 그냥 곧은 막대기라고 가정하고, 완벽한 발걸음을 계산한 다음 다시 조정하는 것입니다. 이렇게 함으로써, 그들은 완벽한 Muon의 춤을 근사(approximate)할 수 있는 최적의 LoRA 헤드셋 업데이트 방식을 찾아냅니다.

논문의 결과에 따르면, 이 새로운 sMuon 방식은 놀라울 정도로 잘 작동합니다. 다양한 언어 모델을 대상으로 테스트했을 때, sMuon은 표준 AdamW 코치 및 다른 Muon 시도들보다 종종 더 나은 성능을 보였습니다. 특히 이미 Muon으로 사전 학습된 로봇(Moonlight라고 불림)에 사용했을 때, sMuon은 11개의 서로 다른 작업 중 6개에서 최고 점수를 기록하며 매우 인상적인 모습을 보였습니다. 사전 학습 실험에서 sMuon은 오차율을 낮추는 데 있어 가장 가까운 경쟁 방식인 Riemannion과 대등한 최고의 결과를 냈지만, 그 과정은 약 30% 더 빨랐습니다.

결정적으로, 저자들은 자신들의 방법이 단순히 이론적인 아이디어에 그치지 않고 실용적이라는 점을 보여줍니다. 다른 고급 방법들은 복잡한 행렬을 분해하고 재조립하는 것과 같은 무겁고 느린 계산을 요구하지만, sMuan은 거의 전적으로 단순한 행렬 곱셈(현대 컴퓨터 칩들이 빛의 속도로 처리하도록 설계된 바로 그 수학)에 의존합니다. 이는 sMuon이 실제 훈련 과정을 늦추지 않고도 실전 훈련에 사용될 만큼 빠르다는 것을 의미합니다. 이 논문은 LoRA 부분을 별개의 관련 없는 조각으로 취급하는 대신 전체 레이어의 기하학적 구조를 존중함으로써, 아주 적은 추가 비용만으로 상당한 성능 향상을 얻을 수 있음을 시사합니다. 이는 때때로 거대한 로봇에게 새로운 기술을 가르치는 가장 좋은 방법은, 작고 똑똑하며 수학적으로 우아한 자극을 주는 것임을 상기시켜 줍니다.

=== 요약 ===

기술적 요약: 저계수 어댑터를 이용한 Muon 근사 (sMuon)

문제 정의

Muon 옵티마이저는 각 레이어에서 가중치 업데이트를 직교화(orthogonalizing)하여 싱귤러 값(singular values)을 효과적으로 균등하게 만듦으로써, 대규모 언어 모델(LLM)의 사전 학습(pretraining)에서 탁월한 성능을 입증했습니다. 이러한 동작은 스펙트럼 노름(spectral norm) 하에서의 최단 강하(steepest descent)에 해당하며, 레이어를 선형 연산자로 볼 때 발생하는 최악의 변화를 제어합니다. 그러나 Muon은 매개변수 효율적 미세 조정(PEFT), 특히 저계수 적응(LoRA) 환경에서는 거의 사용되지 않습니다.

핵심적인 수학적 장애물은 Muon이 전체 가중치 행렬 레이어에 대한 하강 스텝을 직교화할 것을 요구한다는 점입니다. LoRA 설정에서는 가중치 업데이트가 저계수 매개변수화 $\Delta W = BA$ (여기서 $B \in \mathbb{R}^{d_1 \times r}$ 및 $A \in \mathbb{R}^{r \times d_2}$ )로 제한되므로, $\delta B$ 와 $\delta A$ 를 찾아 결과적인 $\delta \Delta W$ 가 Muon이 요구하는 풀-랭크(full-rank) 직교 하강 스텝과 완벽하게 일치하도록 만드는 것은 일반적으로 수학적으로 불가능합니다.

기존의 우회 방법들은 다음과 같습니다:

Per-factor Muon: $B$ 와 $A$ 를 각각 별도로 직교화하며, 이는 전체 레이어의 기하학적 구조를 무시합니다.
LoRA-Muon: 삼각 부등식을 사용하여 목적 함수를 완화합니다.
Riemannian: 저계수 다양체(low-rank manifold) 상의 리만 최적화(Riemannian optimization)로 문제를 정식화하지만, 이는 리트랙션(retraction) 단계에서 비용이 많이 드는 절단된 특이값 분해(truncated SVD)를 요구하여 상당한 계산 오버헤드를 초래합니다.

방법론

저자들은 저계수 설정 내에서 완화된 Muon 목적 함수의 해를 근사하는 방법인 sMuon(small Muon)을 제안합니다. 이 접근 방식은 선형화와 최소제곱 최적화를 활용하여 최적의 하강 스텝을 찾습니다.

1. 선형화 및 최소제곱 정식화

목표는 $\delta \Delta W \approx \delta B A + B \delta A$ 가 이상적인 Muon 스텝 $\delta \Delta W^* = -\eta \cdot \text{msign}(G_{\Delta W})$ 를 근사하도록 하는 $\delta B, \delta A$ 를 찾는 것입니다. (여기서 $\text{msign}(X) = UV^\top$ 는 SVD의 극 인자(polar factor)입니다.)

이상적인 스텝은 일반적으로 풀-랭크이며 저계수 인자들에 의해 정확히 구현될 수 없으므로, 저자들은 2차 항( $\delta B \delta A$ )을 제외하여 문제를 선형화하고 최소제곱 문제를 해결합니다:
$(\delta B^*, \delta A^*) = \arg \min_{\delta B, \delta A} \| \delta \Delta W^* - (\delta B A + B \delta A) \|_F^2$
이는 무어-펜로즈 역행렬( $B^\dagger, A^\dagger$ )을 포함하는 폐쇄형 해(closed-form solution)를 산출합니다. 결정적으로, 저자들은 최적의 업데이트가 인자 $(\delta B, \delta A)$ 측면에서는 유일하지 않지만, 결과적인 전체 레이어에 대한 1차 변화량 $\delta \Delta W$ 는 유일함을 보여줍니다.

2. 효율적인 Matmul 전용 계산

직접적인 폐쇄형 해 구현은 거대한 $d_1 \times d_2$ 행렬의 직교화를 요구하며, 이는 계산적으로 매우 부담스럽습니다. 이를 해결하기 위해 저자들은 대규모 행렬 직교화를 피하는 효율적인 정식화(Algorithm 2)를 도출했습니다:

저계수 그래디언트: 전체 그래디언트 $G_{\Delta W}$ 대신, 최적화 목적 함수에 대해 1차적으로 동등한 $H = P_B G_{\Delta W} + (I-P_B)G_{\Delta W} P_A$ 를 활용합니다.
차원 축소: 항등식 $\text{msign}(H) = U \text{msign}(U^\top H V) V^\top$ 를 활용합니다. 여기서 $U$ 와 $V$ 는 $H$ 의 열 공간과 행 공간을 생성하는 정규 직교 기저입니다. $H$ 의 랭크는 최대 $2r$ 이므로, 이를 통해 값비싼 $d_1 \times d_2$ 직교화 작업을 $2r \times 2r$ 직교화로 축소할 수 있습니다.
구현: 전체 알고리즘은 행렬 곱셈(matrix multiplication), 행렬 역제곱근(matrix inverse square roots), 그리고 작은 $r \times r$ 또는 $2r \times 2r$ 행렬의 극 분해(polar decomposition)만을 기반으로 합니다. 이는 Riemannion에서 요구되는 QR 및 SVD 분해의 오버헤드를 방지합니다.

3. 알고-리드믹 세부 사항

sMuon 옵티마이저는 다음을 포함합니다:

분할 가중치 감쇠 (Split Weight Decay): 풀-파라미터 학습의 역학을 맞추기 위해, 그래디언트 업데이트 전 인자 $A$ 와 $B$ 에 스케일링 인자 $s = \sqrt{1-\lambda\eta}$ 를 적용하여 가중치 감쇠를 수행합니다.
모멘텀 전송 (Momentum Transport): 매개변수 불변성(즉, 특정 인수화 $B, A$ 에 관계없이 옵티마이저가 일관되게 동작함)을 보장하기 위해, 모멘텀 버퍼를 매 단계마다 현재 인자 공간으로 투영합니다.

주요 기여

sMuon의 도출: 스펙트럼 노름 정규화된 최단 강하 목적 함수의 1차 근사를 취함으로써, LoRA 어댑터 학습을 위한 폐쇄형 방법을 도출했습니다.
효율적인 구현: 현대적 하드웨어에 효율적인 "matmul-only" 알고리즘을 제공하여 값비싼 분해 과정을 피합니다. 이는 Per-factor Muon 및 LoRA-Muon과 유사한 속도를 제공하면서도, 특히 더 큰 LoRA 랭크에 대해 Riemannion보다 훨씬 빠릅로 작동합니다.
실험적 평가: 4개의 베이스 모델(Muon으로 사전 학습된 모델 1개 및 AdamW로 사전 학습된 모델 3개 포함)과 11개의 벤치마크(상식 및 코딩)에 걸친 광범위한 실험을 수행했습니다.

결과

논문은 다음의 결과를 보고합니다:

성능: sMuon은 SFT 벤치마크에서 유리한 성능을 보입니다. 결과는 모델에 따라 달라질 수 있으나, Moonlight-16B-A3B(Muon으로 사전 학습됨) 모델을 미세 조정할 때 11개 작업 중 6개에서 최고의 정확도를 달는 sMuon이 가장 좋은 성과를 거두었습니다. 일반적으로 AdamW 및 Per-factor Muon보다 우수한 성능을 보입니다.
ReLoRA 사전 학습: ReLoRA(주기적으로 어댑터를 병합하고 갱신)를 사용하는 사전 학습 환경에서, sMuon은 Riemannion과 동등한 수준의 최저 최종 검증 손실을 달성하면서도 약 30% 더 빨랐습니다.
스펙트럼 특성: 가중치 업데이트 스펙트럼에 대한 경험적 분석 결과, sMuon은 Per-factor Muon 및 LoRA-Muon에 비해 더 균일한 스펙트럼(이상적인 Muon 동작에 더 가까운)을 유지함을 보여줍니다.
효율성: 표 4는 sMuon이 AdamW보다는 약간 느리지만, Riemannion보다는 현저히 빠르다는 것을 보여줍니다. 랭크 $r=64$ 일 때, Riemannion은 스텝 타임에 약 33%의 오버헤드를 추가하는 반면, sMuon은 단 ~0.8%만을 추가합니다.

의의 및 주장

본 논문은 sMuon이 저계수 어댑터에 Muon 최적화를 적용할 수 있는 원칙적이고 메모리 효율적이며 계산 가능한 방법을 제공한다고 주장합니다. s-rank의 기하학적 구조를 독립적인 행렬으로 취급하는 대신 전체 레이어의 기하학적 구조를 존중함으로써, sMuon은 특히 베이스 모델이 Muon으로 사전 학습되었을 때 중간 정도의 성능 향상을 달성합니다. 저자들은 자신들의 방법이 성능을 희생하지 않으면서도 Riemannian 최적화의 "비싼 다양체 제약"을 피한다는 점을 강조합니다. 이는 Muon 기반 미세 조정을 실제 응용 분야에서 더 쉽게 사용할 수 있게 만듭니다.

저자들은 결과가 "모델 및 평가에 따라 달라질 수 있음"을 언급하며 겸손한 태도를 유지하며, 주요 기여는 Muon의 이론적 이점과 PE로의 실질적 제약 사이의 간극을 메우는 깔끔하고 효율적인 근사법이라는 점을 명시합니다.