🤖 AI

LoRA-Muon: Spectral Steepest Descent on the Low-Rank Manifold

이 논문은 Muon의 스펙트럼 최급강하 규칙을 저계수 매니폴드에 적응시킨 메모리 효율적인 옵티마이저인 LoRA-Muon을 소개하며, QR 분해나 2차 모멘트 저장이 필요하지 않으면서도 학습률의 우수한 전이성과 종종 밀집(dense) 베이스라인을 능가하는 성능을 입증한다.

원저자: Franz Louis Cesista, Katherine Crowson, Cédric Simal, Stella Biderman

게시일 2026-06-12

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Franz Louis Cesista, Katherine Crowson, Cédric Simal, Stella Biderman

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

당신이 거대하고 믿을 수 없을 정도로 똑똑한 로봇(딥러닝 모델)에게 새로운 기술을 가르치려 한다고 상상해 보세요. 보통은 로봇의 뇌 전체를 새로 써야 하는데, 여기에는 엄청난 시간과 에너지가 소모됩니다. **LoRA (Low-Rank Adaptation)**는 로봇의 뇌를 새로 쓰는 대신, 작고 탈부착 가능한 공책을 주는 것과 같습니다. 훨씬 빠르고 저렴하지만, 한 가지 문제가 있습니다. 튜닝하기가 매우 까다롭다는 점입니다. 만약 "학습 조절 손잡이"(학습률)를 잘못 돌리면, 로봇은 혼란에 빠지고 공책의 페이지들(요인들)이 서로 어긋나게 됩니다.

이 논문은 그 손잡이를 돌리는 더 똑똑한 방법인 LoRA-Muon을 소개합니다. 이해를 돕기 위해 쉬운 비유를 들어 설명하겠습니다.

1. 문제점: "두 사람이 젓는 작은 배"

LoRA 공책을 하나의 덩어리가 아니라, 두 사람(요인 A와 요인 B)이 함께 노를 저어 배(가중치 행렬 $W$ )를 움직이는 조룻배라고 생각해 보세요.

기존 방식 (AdamW): 코치(옵티마이저)는 사람 A와 사람 B에게 각자 독립적으로 노를 저으라고 지시합니다. 만약 사람 A는 지치고 사람 B는 의욕이 넘친다면, 두 사람은 서로 다른 방향으로 노를 젓게 됩니다. 결국 배는 제자리에서 뱅글뱅글 돌게 되고, 팀은 실패합니다.
문제의 핵심: 코치가 명령하는 "최적의 속도"는 배의 크기나 두 사람의 크기에 따라 크게 달라집니다. 만약 배의 크기(rank)나 사람의 너비(width)를 변경한다면, 완벽한 속도를 찾기 위해 처음부터 다시 학습해야 합니다.

2. 해결책: "유령 배" (LoRA-Muon)

저자들은 두 사람을 따로 코칭하는 대신, 이들을 마치 풀 사이즈의 거대한 배를 타는 것처럼 코칭해야 한다는 것을 깨달았습니다.

비유: 로봇의 뇌를 나타내는 완벽한 풀 사이즈의 배가 물 위에 떠 있는 "유령 배"가 있다고 상상해 보세요. Muon 옵티마이저는 특별한 "스펙트럼" 규칙(가장 가파르고 효율적인 내리막길을 찾는 기하학적 방법)을 사용하여 이 유령 배를 조종하는 법을 아는 코치입니다.
마법 같은 효과: LoRA-Muon은 "만약 우리가 유령 배를 조종한다면 어떻게 할 것인가?"라고 묻습니다. 유령 배의 움직임을 계산한 뒤, 그 움직임을 사람 A와 사람 B에게 나누어 주어 두 사람이 완벽하게 보조를 맞추도록 합니다.
결과: 이들은 유령 배의 경로를 따르기 때문에 절대 서로 어긋나지 않습니다. 심지어 배의 크기나 사람의 크기를 바꾸더라도, "유령 배"는 항상 동일한 속도를 알려줍니다. 즉, **학습률이 다양한 크기와 형태 사이에서 완벽하게 전이(transfer)**됩니다.

3. "분할 가중치 감쇠(Split Weight Decay)" 기술

기존 방식에서는 배가 너무 무거워지는 것을 막기 위해(가중치 감쇠) 배를 약간 줄이려고 하면, 실수로 사람 A와 사람 B를 서로 다르게 줄여 배가 기울어지게 만들 수 있습니다.

LoRA-Muon의 해결책: 이들은 "분할 가중치 감쇠" 규칙을 발명했습니다. 이것은 마치 마법의 탄성 밴드와 같아서, 두 사람을 완벽한 조화를 이루며 함께 늘리고 줄임으로써 배가 수평을 유지하게 하고, 마치 풀 사이즈의 배를 다루는 것처럼 수학적으로 정확하게 작동하도록 만듭니다.

4. 경쟁 모델보다 뛰어난 이유

이 논문은 LoRA-Muon을 두 가지 다른 방법인 Spectron 및 LoRA-RITE와 비교합니다.

Spectron: 이 코치는 현재 사람 A와 사람 B가 얼마나 지쳤는지를 보고 속도를 결정합니다. 만약 실수로 사람 A를 사람 B보다 두 배 더 크게 보이게 만든다면("게이지 스케일링" 문제), Spectron은 혼란에 빠져 속도를 바꿔버립니다. 즉, 임의적인 선택에 민감합니다.
LoRA-RITE: 사실 이 방식은 LoRA-Muon과 동일한 일을 수행하지만, 매우 복잡하고 무거운 도구(QR 분해)를 사용합니다. 이는 마치 견과류를 깨기 위해 대형 망치를 사용하는 것과 같습니다.
LoRA-Muon: LoRA-Muon은 LoRA-RITE와 똑같이 똑똑한 조종을 수행하면서도, 더 가볍고 빠른 도구를 사용합니다. 무거운 망치를 사용할 필요가 없으므로 컴퓨터의 메모리를 적게 사용하며 더 빠릅니다.

5. 증명: Tiny Shakespeare

저자들은 이 모델을 작은 작업, 즉 로봇에게 "Tiny Shakespeare"(셰익스피어 희곡의 아주 작은 데이터셋)를 쓰는 법을 가르치는 데 테스트했습니다.

Rank 2 (작은 공책): 아주 작은 공책(Rank 2)을 사용하더라도, LoRA-Muon은 거대한 풀 사이즈 로봇과 정확히 동일한 완벽한 속도를 찾아냈습니다.
Rank 32 (중간 크기 공 notebook): 약간 더 큰 공책을 사용했을 때, LoRA-Muon은 실제로 풀 사이즈 로봇보다 더 나은 성능을 보이며 평균적으로 더 낮은 오차율을 달랐습니다.
"게이지(Gauge)" 테스트: 그들은 의도적으로 사람 A와 사람 B의 시작 크기를 엉망으로 만들었습니다. 기존 방식(Spectron)은 혼란에 빠져 성능이 떨어졌지만, LoRA-Muon은 그 혼란을 전혀 알아채지 못했고 계속해서 똑바로 노를 저어 나갔습니다.

결론

LoRA-Muon은 "공책" 버전의 모델을 "풀 사이즈 뇌" 버전처럼 취급하는 새로운 AI 모델 학습 방식입니다. 이를 통해 다음과 같은 효과를 얻습니다:

공책의 크기나 스케일에 상관없이 절대 혼란에 빠지지 않습니다.
아주 작은 공책에서도 거대한 공책과 동일한 속도 설정을 사용할 수 있습니다.
기존의 똑똑한 방법들보다 더 빠르게 실행되며 메모리를 적게 사용합니다.

이것은 작은 AI 공책을 튜닝하는 어려운 기술을, 그냥 작동할 수밖에 없는 단순하고 신뢰할 수 있는 과정으로 바꾸어 놓았습니다.

기술 요약: LoRA-Muon: 저계수 매니폴드 상의 스펙트럼 최급강하법 (Spectral Steepest Descent on the Low-Rank Manifold)

1. 문제 정의

저계수 적응(Low-Rank Adaptation, LoRA)은 훈련 FLOP 및 메모리 요구량을 줄이는 효율성 덕분에 대규모 딥러닝 모델의 미세 조정(fine-tuning)에 널리 채택되고 있습니다. 그러나 $A$ 와 $B$ 라는 개별적인 요소(factors)에 AdamW와 같은 표준 옵티마이저를 적용할 경우, 밀집 학습(dense training)에 비해 상당한 튜닝의 어려움이 발생합니다. 이러한 문제는 다음과 같습니다:

초기화에 대한 민감도: 성능이 초기화 선택에 따라 크게 좌우됩니다.
불량한 하이퍼파라미터 전이성: 최적의 학습률(learning rate)이 계수(rank), 너비(width), 깊이(depth) 또는 계수 스케일링 관례(factor scaling conventions)를 가로질러 잘 전이되지 않습니다.
발산하는 역학(Divergent Dynamics): 동일한 합성 가중치( $W = AB^T$ )에서 시작하더라도, 계수의 스케일이 서로 멀어지면 훈련 역학이 분리될 수 있습니다.
차선의 성능: 네이티브 저계수 사전 훈련(low-rank pretraining)은 추가적인 안정화 없이는 밀집 베이스라인(dense baselines)을 따라잡기 어렵습니다.

핵심 과제는 풀 랭크(full-rank) 옵티마이저(특히 Muon 및 Shampoo 계열)의 기하학적 특성을 유지하면서, 저계수 매니폴드 상에서 최적의 학습률이 계수, 너비, 깊이 및 계수 재스케일링 간에 전이되도록 하는 옵티마이저를 설계하는 것입니다.

2. 방법론

2.1 기하학적 기초

본 논문은 옵티마이저 설계를 특정 노름(norm) 하에서 옵티마이저가 수행하는 최급강하(steepest descent)를 식별하는 기하학적 문제로 재정의합니다.

Muon 옵티마이저: 주변 행렬 공간 $\mathbb{R}^{m \times n}$ 내의 스펙트럼 노름( $\|\cdot\|_{2\to2}$ ) 하에서 최급강하를 수행합니다. 업데이트 방향은 행렬 sign 함수인 $-\text{msign}(G)$ 로 주어집니다.
저계수 매니폴드(Low-Rank Manifold): LoRA는 가중치를 매니폴드 $\mathcal{M}_r = \{W \in \mathbb{R}^{m \times n} : \text{rank}(W) = r\}$ 로 제한합니다. 저자들은 $A$ 와 $B$ 를 독립적인 좌표로 최적화하는 대신, 이 매니폴드의 접공간(tangent space) $T_W \mathcal{M}_r$ 상에서 직접 최급강하 문제를 푸는 것을 제안합니다.

2.2 LoRA-Muon의 유도

저자들은 스펙트럼 노름 하의 저계수 접공간 상에서 제약된 선형 최소화 문제를 해결함으로써 LoRA-Muon을 유도합니다.

접공간 분해: $W$ 에 대한 업데이트는 $\Delta W = \Delta A B^T + A \Delta B^T$ 로 분해됩니다.
신뢰 영역 분할(Trust Region Splitting): 문제를 다루기 쉽게 만들기 위해, 신뢰 영역 예산 $\eta$ 를 두 개의 접 성분 사이에 균등하게 배분합니다. 이는 두 성분이 훈련 중에 빠르게 정렬되어 삼각 부등식이 거의 타이트해진다는 경험적 관찰에 의해 정당화됩니다.
투영체 공식화(Projector Formulation): 유니타리 불변성(unitary invariance)을 사용하여, 문제를 투영된 그래디언트에 대한 선형 최소화 오라클(Linear Minimization Oracles, LMOs)을 찾는 문제로 축소합니다. 이상적인 업데이트는 다음과 같습니다:
$\Delta W = \frac{\eta}{2} \text{LMO}(\nabla W f \cdot P_B) + \frac{\eta}{2} \text{LMO}(P_A \cdot \nabla W f)$
여기서 $P_A$ 와 $P_B$ 는 $A$ 와 $B$ 의 열 공간(column spaces)으로의 투영체입니다.
스펙트럼 특화: 스펙트럼 노름의 경우, LMO는 $-\text{msign}(\cdot)$ 입니다. 결과적인 계수 업데이트(식 14–15)는 현재 Gram 행렬( $S_A = A^TA, S_B = B^TB$ )에 의해 계수 그래디언트를 화이트닝(whitening)하고 스펙트럼 sign 함수를 적용하는 과정을 포함합니다:
$\Delta A^* = -\frac{\eta}{2} \text{msign}(\nabla_A f \cdot S_B^{-1/2}) S_B^{-1/2}$
$\Delta B^* = -\frac{\eta}{2} \text{msign}(\nabla_B f \cdot S_A^{-1/2}) S_A^{-1/2}$

2.3 분할 가중치 감쇠(Split Weight Decay)

표준적인 요소별 가중치 감쇠는 풀 랭크와 저계수 업데이트 역학 사이의 불일치를 초래합니다. 저자들은 감쇠 인자 $s = \sqrt{1 - \lambda \eta}$ 가 비대칭적으로 적용되는 분할 가중치 감쇠 규칙을 유도했습니다:
$A_{t+1} = s A_t + \frac{1}{s} \Delta A^*_t, \quad B_{t+1} = s B_t + \frac{1}{s} \Delta B^*_t$
이는 합성된 가중치 업데이트가 원하는 풀 랭크 역학을 일치시켜 하이퍼파라미터의 전이성을 보존하도록 보장합니다.

2.4 게이지 불변성(Gauge Invariance)

LoRA-Muon의 핵심적인 특성은 게이지 불변성입니다. 합성된 가중치 $W$ 에 대한 업데이트는 임의의 가역적 변환 $(A, B) \to (AR, BR^{-T})$ 하에서도 불변합니다. 이는 업데이트가 그 자체로 게이지 불변인 투영체 $P_A$ 와 $P_B$ 에만 의존하기 때문에 가능합니다.

3. 주요 기여

LoRA-Muon의 유도: 유니타리 불변 노름(스펙트럼 노름으로 특화됨) 하의 저계수 매니폴드 상에서 최급강하 문제를 해결하여 유도된 새로운 옵티마이저 계열.
분할 가중치 감쇠 규칙: 저계수 훈련의 역학을 풀 랭크 훈련과 일치시켜 하이퍼파라미터 전이를 가능하게 하는 특정 업데이트 규칙.
게이지 불변성 증명: LoRA-Muon 업데이트가 전체 게이지 그룹 작용에 대해 불변함을 이론적으로 증명하여, 옵티마이저가 임의의 계수 좌표가 아닌 유도된 가중치 업데이트에 의pon하도록 보장함.
Spectron과의 비교: Spectron은 LoRA-Muon과 달리 게이지 불변이 아니며 임의의 계수 재스케일링에 민감하다는 점을 입증.
LoRA-RITE와의 관계: LoRA-RITE의 단순화된 QR-좌표 코어가 대수적으로 LoRA-Muon의 스펙트럼 업데이트와 동일함을 증명하였으나, LoRA-Muon은 QR 분해나 2차 모멘트 저장의 계산 오버헤드를 피함.
경험적 검증: LoRA-Muon이 밀집(dense) Muon의 학습률 탐색을 위한 효과적인 프록시 역할을 한다는 증거 제시.

4. 실험 결과

실험은 계산량이 매칭된 밀집 및 저계수 베이스라인을 사용하여 TinyShakespeare 데이터셋 기반의 문자 수준 언어 모델링에 대해 수행되었습니다.

학습률 전이: 밀집 Muon의 최적 학습률($lr = 0.1$)이 계수(2, 4, 8, 16, 32), 너비, 깊이, 스칼라 인자 재스케일링의 변화에 대해 LoRA-Muon으로 완벽하게 전이되었습니다.
- Rank 2에서, LoRA-Muon은 밀집 베이스라인의 최적 테스트 학습률을 회복합니다.
- Rank 32에서, LoRA-Muon은 시드 평균 스윕(seed-averaged sweep)에서 밀집 베이스라인( $1.789 \pm 0.002$ )보다 낮은 평균 검증 손실( $1.776 \pm 0.002$ )을 달s성했습니다.
게이지 재균형(Gauge Rebalancing): 계수를 큰 스칼라 불균형(스케일 9)으로 초기화했을 때, 재균형 단계를 적용해도 LoRA-Muon은 미미한 영향(손실 변화 $\approx -5.1 \times 10^{-5}$ )만을 미쳐 게이지 불변성을 확인했습니다. 반면, Spectron의 성능은 유의미하게 변화( $2.0 \times 10^{-2}$ )하여 계수 스케일링에 대한 민감도를 드러냈습니다.
효율성: 랭크-2 프록시는 밀집 모델의 학습 가능한 파라미터의 약 4.3%만을 사용하면서도, 동일한 계산 제약 조건 하에서 약 9.2배 더 큰 옵티마이저 단계 예산을 지원합니다.

5. 의의 및 주장

본 논문은 LoRA-Muon이 LoRA의 튜닝 어려움에 대한 이론적으로 근거 있는 해결책을 제공한다고 주장합니다. 저자들은 LoTA를 단순한 계수 좌표계가 아닌 저계수 매니폴드 상의 최적화 문제로 취급함으로써 다음을 달성했습니다:

강건성(Robustness): 옵티마이저는 임의의 인수 분해 선택(게이지 대칭성)에 대해 불변합니다.
전이성(Transferability): 최적의 학습률(및 스케일링 법칙을 통한 다른 하이퍼파라미터)이 계수, 너비, 깊이, 계수 스케일 간에 원활하게 전이됩니다.
효율성(Efficiency): 값비싼 QR 분해나 2차 모멘트 저장 없이 스펙트럼 업데이트를 계산하여 가속기 친화적입니다.

저자들은 LoRA-Muon을 풀 랭크 Muon 및 Shampoo 계열 옵티마이저를 위한 "좋은 저계수 프록시"로 위치시킵니다. 이들은 저계수 파라미터화가 옵티마이저의 기하학을 결정해서는 안 되며, 대신 업데이트가 저계수 행렬 매니폴드 자체의 기하학에 의해 정의되어야 함을 강조합니다.

한계점: 저자들은 현재의 경험적 주장이 TinyShakespeare 규모의 모델에 구축되었으며, 대규모 사전 훈련이나 다운스트림 미세 조정 벤치마크에서는 아직 검증되지 않았음을 명시합니다. 본 연구는 전체 하이퍼파라미터 스윕보다는 학습률 전이에 초점을 맞추고 있습니다.