기술 요약: Muon 은 Adam 으로 사전 학습된 모델을 미세 조정할 수 있는가?

문제 제기

Muon(Newton-Schulz 에 의해 직교화된 MomentUm) 은 대규모 언어 모델 (LLM) 사전 학습을 위한 매우 효율적인 최적화자로 부상하여, Adam 대비 약 2 배의 계산 효율성을 제공하고 2 차 모멘트 추정을 제거함으로써 메모리 사용량을 줄였습니다. 그러나 미세 조정 (fine-tuning) 에 대한 실용적 적용은 제한적입니다. 주요 장애물은최적화기 불일치 문제입니다: 대부분의 공개된 오픈 모델은 Adam 으로 사전 학습되었으며, 미세 조정 시 무작위로 Muon 으로 전환하면 성능이 저하됩니다. 반대로, Muon 으로 사전 학습된 모델을 Adam 으로 미세 조정해도 최적의 결과를 얻지 못합니다. 이러한 비호환성은 Muon 의 사전 학습 성공에도 불구하고, 지배적인 미세 조정 패러다임에서의 Muon 채택을 심각하게 제한합니다.

방법론 및 분석

저자들은 통제된 실험과 이론적 분석을 통해 Adam 과 Muon 의 고유한**암시적 편향 (implicit biases)**에 초점을 맞춰 이 불일치를 조사합니다.

1. 암시적 편향 특성화:

   • Adam은 1 차 및 2 차 모멘트에 기반한 요소별 (element-wise) 사전 조건화를 사용하여, **최대 노름 마진 ($\|W\|_{\max}$)**을 최대화하는 해로 수렴합니다.
   • Muon은 Newton-Schulz 직교화를 통한 행렬 수준 (matrix-level) 사전 조건화를 적용하여, **스펙트럼 노름 마진 ($\|W\|_2$)**을 최대화하는 해로 수렴합니다.
   • 단순화된 선형 회귀 문제에 대한 이론적 분석은 이러한 최적화기들이 근본적으로 다른 해 ($W^*_{\max}$ 대 $W^*_2$) 로 수렴함을 보여줍니다. 실험적으로 Muon 으로 학습된 가중치는 Adam 으로 학습된 가중치에 비해 더 높은 안정적 랭크 (stable rank) 와 더 균일한 특이값 분포를 보입니다.

2. 저하 메커니즘:

   • 불일치는 사전 학습된 지식을 방해합니다. 다른 암시적 편향을 가진 최적화기가 적용되면, 모델 가중치는 사전 학습 단계와 양립할 수 없는 구조적 방향으로 강제됩니다.
   • 실험 결과, 불일치된 미세 조정은 업데이트 강도에 대한 민감도를 증가시킵니다. 더 강력한 업데이트는 사전 학습된 구조를 더 크게 교란시켜 더 높은 퍼플렉시티 (perplexity) 와 더 심각한 재앙적 망각 (catastrophic forgetting) 을 초래합니다.

3. 제안된 완화 방안: LoRA:

   • 저자들은업데이트 범위를 제한함으로써 이 불일치를 완화할 수 있다고 가정합니다. 그들은 사전 학습된 가중치 ($W_0$) 를 고정하고 업데이트를 저랭크 부분 공간 ($\Delta W = BA$) 으로 제한하는**저랭크 적응 (Low-Rank Adaptation, LoRA)**을 활용합니다.
   • 이 접근법은 가중치 변화의 크기를 제한하여 최적화기가 사전 학습된 구조를 공격적으로 변경하는 것을 방지합니다. 이론적 분석에 따르면, "불일치 팽창 (mismatch inflation, 일치된 성능과 불일치된 성능 간의 격차)"은 LoRA 랭크 $r$에 의해 제한되며, $r \to 1$로 갈 때 사라집니다.

주요 기여

• 재현 및 분석: 이 논문은 접근 가능한 규모 (NanoChat, 561M 파라미터) 에서 최적화기 불일치 현상을 재현하고, 이를 Adam 과 Muon 의 고유한 암시적 편향 (최대 노름 대 스펙트럼 노름) 과 연결합니다.
• 해결책으로서의 LoRA: 저자들은 LoRA 가 불일치를 효과적으로 완화함을 입증합니다. 업데이트를 제한함으로써 LoRA 는 Muon 이 Adam 으로 사전 학습된 모델을 미세 조정할 때 LoRA-Adam 과 동등하거나 더 나은 성능을 발휘하도록 합니다.
• 종합적 검증: 이 발견들은 다음 분야에서 검증되었습니다:
  • 자연어 이해 (NLU): T5-Base 를 사용한 GLUE 벤치마크.
  • 자연어 생성 (NLG): 수학, 코드, 상식 추론에 대한 Llama 2-7B/13B 의 지시 미세 조정.
  • 컴퓨터 비전: 6 가지 이미지 분류 작업에 대한 CLIP ViT-B/32 미세 조정.
• 랭크 및 망각 연구:
  • 랭크 민감도: LoRA-Muon 의 이점은 저랭크에서 중랭크 구간에서 가장 두드러집니다. 랭크가 증가하여 전체 미세 조정에 가까워질수록 불일치가 재발하며, Adam 대비 성능이 저하됩니다.
  • 재앙적 망각: Adam 으로 사전 학습된 모델을 Muon 으로 전체 미세 조정하면 Adam 보다 사전 학습된 지식의 더 심각한 망각을 초래합니다. LoRA 는 이 망각을 크게 줄여 사전 학습된 구조를 보존합니다.
• LoRA 변형: 이 연구는 기존 LoRA 변형 (예: rsLoRA, PiSSA, LoRA-One) 을 평가합니다. LoRA 업데이트를 전체 미세 조정에 더 가깝게 만들도록 설계된 변형 (PiSSA 등) 은 Muon 에게 반드시 도움이 되지 않으며, 오히려 불일치를 악화시킬 수 있음을 발견하여 Muon 전용 적응의 필요성을 시사합니다.

결과

• 성능 격차 축소: 전체 미세 조정에서 Muon 은 Adam 으로 사전 학습된 모델에서 일관되게 Adam 보다 낮은 성능을 보입니다 (예: WikiText-2 에서 +0.009~+0.023 퍼플렉시티 격차). LoRA 하에서는 이 격차가 크게 축소됩니다 (Muon 사전 학습 모델은 39%, Adam 사전 학습 모델은 78% 축소).
• LoRA 내 우위: NLU, NLG, 비전 작업 전반에 걸쳐 LoRA-Muon은 LoRA-Adam과 동등하거나 더 나은 성능을 냅니다. 예를 들어, GLUE 벤치마크에서 LoRA-Muon-PE 는 평균 정확도 89.20% 를 기록하여 LoRA-Adam(88.93%) 과 전체 Adam(89.14%) 을 모두 능가했습니다.
• 효율성: Muon 은 Newton-Schulz 반복으로 인해 LoRA 설정에서 Adam 보다 1.1~1.2 배 느린 modest 한 단계당 오버헤드를 incur 하지만, 메모리 효율성 (Adam 의 두 개 모멘트 버퍼 대비 하나의 모멘트 버퍼만 저장) 은 여전히 실용적인 장점으로, 최적화기 상태 메모리의 약 50% 를 절약합니다.

중요성 및 주장

이 논문은 Adam 과 Muon 간의 최적화기 불일치 문제에 대한 최초의 심층 분석을 제공한다고 주장합니다. 그 주요 중요성은 방대한 Adam 사전 학습 모델 생태계에 대해 Muon 을 미세 조정에 실용적으로 사용할 수 있게 한다는 점에 있습니다.

저자들은 다음과 같이 결론 내립니다:

1. 최적화기 불일치는 구조적입니다: 업데이트 강도가 제한되지 않을 때 사전 학습된 지식을 방해하는 상충되는 암시적 편향에서 비롯됩니다.
2. 제한된 업데이트가 핵심입니다: 가중치 변화의 크기를 제한하는 LoRA 와 같은 기법은 이 불일치를 효과적으로 중화시킵니다.
3. 실용적 권장사항: 실무자들은 전체 미세 조정에서 관찰된 불일치 페널티 없이 메모리 효율성과 경쟁력 있는 성능을 얻기 위해 LoRA 미세 조정 시 Adam 의 대체제로 Muon 을 사용할 수 있습니다.

이 논문은 겸손하게도, LoRA 가 문제를 완화하지만 Adam 에 최적화된 기존 LoRA 변형이 자동으로 Muon 으로 이전되지 않으므로, Muon 전용 적응을 개발하기 위한 향후 연구가 필요하다고 지적합니다. 또한, 불일치의 이론적 특성화는 여전히 열린 질문으로 남아 있습니다.