NuMuon: Nuclear-Norm-Constrained Muon for Compressible LLM Training

이 논문은 Muon 옵티마이저로 학습된 모델도 낮은 랭크 구조를 보인다는 발견을 바탕으로, 업데이트 방향에 핵노름 제약을 추가해 가중치 압축성을 높이고 압축 후 모델 품질을 개선하는 NuMuon 옵티마이저를 제안합니다.

핵심

이 논문은 **"거대한 인공지능 (LLM) 을 더 작고 가볍게 만들면서도, 그 지능을 잃지 않는 새로운 훈련 방법"**을 소개합니다.

비유를 들어 쉽게 설명해 드리겠습니다.

1. 문제: "무거운 짐을 들고 가는 것"

지금까지의 거대 언어 모델 (LLM) 은 마치 수천 권의 책을 모두 외운 천재와 같습니다. 하지만 이 천재는 기억력이 너무 좋아서, 모든 정보를 머릿속에 저장하려다 보니 머리 (메모리) 가 너무 무겁고 비싸게 유지됩니다.

이 천재를 더 작은 방 (휴대폰이나 일반 서버) 에 데려가려면, **불필요한 정보를 버리고 핵심만 남기는 '압축'**이 필요합니다. 하지만 보통은 압축을 하면 지능이 떨어집니다. (예: 천재가 중요한 수학 공식까지 잊어버리는 것)

2. 기존 방법의 한계: "무작위 정리" vs "체계적인 정리"

• 기존 훈련 (AdamW 등): 천재가 정보를 배울 때, 책장 전체를 무작위로 훑어보며 중요한 것, 중요하지 않은 것을 구분합니다. 결과적으로 책장에는 중요한 정보도 있지만, 불필요한 잡동사니도 꽤 많이 섞여 있습니다. 그래서 압축하려고 하면, 중요한 정보까지 함께 버리게 되어 성능이 급격히 떨어집니다.
• 새로운 훈련 (Muon): 최근 등장한 'Muon'이라는 훈련법은 천재가 정보를 배울 때, 책장 전체를 한 번에 훑어보며 (Full-rank) 매우 효율적으로 학습합니다. 놀라운 점은, 이 방법으로 훈련된 천재도 실제로는 책장 대부분이 비어있고 핵심 정보만 모여 있다는 사실이 발견되었습니다. 하지만 이 '비어있는 상태'가 너무 약해서, 압축을 너무 강하게 가하면 바로 무너집니다.

3. 해결책: "NuMuon (뉴뮤온)" - "의도적으로 정리하는 훈련"

이 논문은 NuMuon이라는 새로운 훈련법을 제안합니다.

비유: "책장 정리를 미리 계획하는 것"

기존의 'Muon'이 무작위로 정보를 받아들이는 대신, NuMuon은 훈련 시작부터 **"이 책장은 최대 10 권의 책만 담을 수 있어!"**라고 미리 규칙을 정해둡니다.

• 핵심 아이디어: 천재가 정보를 배울 때, 중요한 정보 (핵심) 만 남기고 나머지는 아예 들어오지 못하게 막는 것입니다.
• 어떻게?: 수학적으로 '핵심 정보의 양 (핵심 노름, Nuclear Norm)'을 제한합니다. 마치 가방의 무게를 미리 정해두고, 그 무게 안에 가장 중요한 물건들만 넣으라고 지시하는 것과 같습니다.

4. NuMuon 의 놀라운 결과

이 방법으로 훈련된 천재 (모델) 는 다음과 같은 장점이 있습니다:

1. 압축에 강함: 처음부터 불필요한 잡동사니를 배제했기 때문에, 나중에 압축 (가방을 더 작게 만들기) 을 하더라도 핵심 지능이 거의 손상되지 않습니다.
2. 성능 유지: 압축률이 80% 에 달해도 (가방을 5 분의 1 로 줄여도) 원래 천재의 실력을 90% 이상 유지합니다. 기존 방법들은 이 정도 압축이면 천재가 멍청해지거나 아예 말을 못 했습니다.
3. 빠른 속도: 불필요한 정보가 없으므로, 이 천재가 말을 할 때 (추론) 훨씬 더 빠르고 가볍습니다.

5. 결론: "가볍고 빠른 미래"

이 연구는 **"무거운 AI 모델을 훈련할 때, 처음부터 가볍게 만들 수 있는 방법"**을 찾아냈습니다.

• 기존: 무거운 짐을 싣고 가서, 나중에 버릴 만한 것을 찾아서 버리다 보니 중요한 것도 함께 버림. (성능 저하)
• NuMuon: 처음부터 가방의 크기를 정해두고, 그 안에 꼭 필요한 것만 싣고 출발. (성능 유지 + 가벼움)

이 기술이 상용화되면, 고가의 서버 없이도 일반 스마트폰에서 매우 똑똑하고 빠른 AI 를 구동할 수 있게 되어, AI 기술이 우리 일상으로 훨씬 더 빠르게 퍼져나갈 것입니다.

기술 요약

NuMuon: 압축 가능한 LLM 훈련을 위한 핵-노름 제약 Muon

이 논문은 대규모 언어 모델 (LLM) 의 배포 비용과 메모리 제약 문제를 해결하기 위해, Muon 옵티마이저의 가중치 공간 구조를 분석하고 이를 개선한 새로운 옵티마이저 NuMuon (Nuclear-Norm-Constrained Muon) 을 제안합니다.

1. 문제 제기 (Problem)

• 배포 비용: 수십억 개의 파라미터를 가진 LLM 은 메모리, 저장 공간, 가속기 비용이 매우 높아 실제 배포에 어려움이 있습니다.
• 압축의 중요성: LLM 을 배포하기 위해 가중치 행렬의 구조 (예: 저랭크 구조, 희소성) 를 활용한 압축 기술이 활발히 연구되고 있습니다.
• 옵티마이저의 영향: AdamW 와 같은 기존 옵티마이저는 암묵적으로 저랭크 편향 (low-rank bias) 을 가지지만, 최근 제안된 Muon 옵티마이저는 풀랭크 (full-rank) 업데이트를 수행하여 수렴 속도를 높입니다.
• 핵심 질문: Muon 은 명시적인 랭크 제약을 두지 않고 풀랭크 업데이트를 수행하지만, 실제로 학습된 가중치가 저랭크 구조를 가지는지, 그리고 이것이 압축에 어떤 영향을 미치는지 명확히 규명되지 않았습니다.

2. 주요 발견 및 방법론 (Methodology)

2.1. Muon 의 놀라운 발견

저자들은 Muon 으로 학습된 모델 (Qwen3-0.6B 등) 의 가중치 행렬을 분석한 결과, 명시적인 랭크 제어 없이도 가중치 행렬이 훈련 과정에서 뚜렷한 저랭크 구조 (pronounced low-rank structure) 를 보인다는 것을 발견했습니다.

• 관측: Muon 은 풀랭크 직교화 업데이트를 수행하지만, 학습된 가중치의 안정된 랭크 (stable rank) 는 여전히 낮게 유지됩니다.
• 한계: Muon 기반 모델은 일반적인 저랭크 압축 파이프라인에서 어느 정도 압축이 가능하지만, 압축률이 높아질수록 (예: 80% 이상) 성능이 급격히 저하되는 취약점을 보입니다.

2.2. NuMuon 제안

이러한 관측을 바탕으로, 저자는 Muon 의 업데이트 방향에 핵-노름 (Nuclear Norm) 제약을 추가하여 명시적으로 랭크를 제어하는 NuMuon을 제안합니다.

• 핵심 아이디어: Muon 을 선형 최소화 오라클 (LMO, Linear Minimization Oracle) 의 관점에서 해석합니다. Muon 은 스펙트럼 노름 (spectral norm) 볼 (ball) 에서 LMO 를 수행합니다. NuMuon 은 여기에 핵-노름 예산 (nuclear-norm budget) 을 추가하여 업데이트 방향의 랭크를 제한합니다.
• 수학적 유도:
  • 제약 조건: $\|\Delta W\|_2 \le \rho$ (스펙트럼 노름) 및 $\|\Delta W\|_* \le \tau$ (핵-노름).
  • 이 문제는 선형 계획법 (Linear Program) 으로 축소되며, 최상위 $k$개의 특이 벡터 (top-$k$ singular vectors) 를 가진 닫힌 형식 (closed-form) 해를 가집니다.
  • 결과적으로 NuMuon 은 Muon 의 풀랭크 직교화 업데이트를 랭크 $k$의 저랭크 업데이트로 변환합니다.
• 실제 구현:
  • Top-$k$ SVD: 전체 SVD 는 계산 비용이 너무 크므로, 랜덤화 된 블록 크릴로프 (Randomized Block Krylov) 방법을 사용하여 효율적으로 상위 $k$개의 특이 벡터를 근사합니다.
  • 랭크 스케줄러 (Rank Scheduler): 훈련 초기에는 높은 랭크를 허용하고, 훈련이 진행됨에 따라 랭크를 점차 낮추는 (Cosine, Piecewise 등) 스케줄링 전략을 사용하여 수렴성과 압축성을 균형 있게 조절합니다.

3. 주요 기여 (Key Contributions)

1. Muon 의 암묵적 편향 규명: 풀랭크 업데이트를 수행하는 Muon 이 학습된 가중치에서 뚜렷한 저랭크 구조를 생성한다는 것을 실험적으로 증명했습니다.
2. NuMuon 알고리즘 개발: Muon 에 핵-노름 제약을 도입하여 업데이트 방향의 랭크를 명시적으로 제어하는 새로운 옵티마이저를 제안했습니다. 이는 상위 $k$개 특이 벡터에 대한 선형 계획법으로 효율적으로 해결됩니다.
3. 수렴성 보장: 비볼록 (non-convex) 가정 하에서 NuMuon 의 수렴성을 이론적으로 증명했습니다. (Ky Fan $k$-노름 기반의 정상성 bound 제공)
4. 압축성 향상: NuMuon 으로 학습된 모델은 기존 Muon 모델보다 훨씬 낮은 안정된 랭크를 가지며, 이는 후속 저랭크 압축 파이프라인에서 훨씬 우수한 성능을 유지함을 실험적으로 입증했습니다.

4. 실험 결과 (Results)

• 모델 범위: Qwen3-0.6B, Olmo2-1.4B, Llama3-1.8B 등 0.6B~1.8B 규모의 모델에서 AdamW, Muon, NuMuon 을 비교했습니다.
• 훈련 성능: NuMuon 은 Muon 과 유사한 수렴 속도와 최종 손실 (loss) 을 보이며, AdamW 보다 우수한 성능을 달성했습니다.
• 압축 성능 (핵심 결과):
  • 압축률 40~80% 구간: NuMuon 으로 학습된 모델은 ASVD, SVD-LLM, Dobi-SVD 등 최신 압축 기법을 적용했을 때, Muon 기반 모델보다 압축 - 품질 트레이드오프가 55.9% 까지 개선되었습니다.
  • 예시: Llama3-1.8B 모델을 40% 로 압축했을 때, Muon 은 검증 퍼플렉시티 (PPL) 가 급격히 증가한 반면, NuMuon 은 PPL 이 98.6% 감소하여 원본 모델에 가까운 성능을 유지했습니다.
  • 추론 속도: 동일한 퍼플렉시티 수준에서 NuMuon 모델은 더 높은 압축률을 허용하여, AdamW 나 Muon 모델보다 더 빠른 추론 처리량 (throughput) 을 달성했습니다.
• 구조적 분석: NuMuon 은 가중치 행렬의 주된 스펙트럼 부분 공간과 업데이트 방향이 더 잘 정렬 (alignment) 되어 있어, 저랭크 근사가 더 효과적으로 이루어짐을 확인했습니다.

5. 의의 및 결론 (Significance)

• 배포 최적화: NuMuon 은 LLM 의 훈련 단계에서부터 배포를 고려한 구조 (저랭크) 를 학습시킴으로써, 후처리 압축 없이도 높은 압축 효율을 달성할 수 있게 합니다. 이는 메모리 제약이 심한 환경이나 비용 민감한 배포 시나리오에 매우 유용합니다.
• 이론적 통찰: 옵티마이저의 동역학이 학습된 가중치의 기하학적 구조 (스펙트럼 분포) 에 어떻게 영향을 미치는지에 대한 이해를 깊게 했습니다.
• 확장성: 분산 훈련 환경 (대역폭 제약) 에서 저랭크 업데이트를 효율적으로 통신할 수 있는 잠재력을 가지며, 향후 분산 LLM 훈련 연구에도 기여할 수 있습니다.

요약하자면, NuMuon은 Muon 의 빠른 수렴 장점을 유지하면서, 명시적인 랭크 제어를 통해 모델이 압축에 더 친화적인 가중치 구조를 학습하도록 유도하여, 고압축 환경에서도 높은 성능을 유지하는 LLM 훈련 솔루션을 제공합니다.