Muon으로 Adam-Pretrained 모델의 파인튜닝 가능성에 관한 연구

대형 언어 모델의 사전학습에 Muon의 효율성이 주목되지만, 대부분의 공개 모델은 Adam으로 사전학습되어 파인튜닝 시 성능 저하가 발생한다. LoRA를 통해 업데이트를 제한하면 Muon과 Adam 간의 implicit bias 차이를 완화하고, 다수의 태스크에서 Muon의 파인튜닝 효과를 실현 가능하게 한다.

1. 출발점: Adam과 Muon은 gradient의 업데이트 방향에 대해 서로 다른 preconditioning를 사용한다. Adam은 각 매개변수에 대해 서로 다른 스케일로 학습률을 조정하고, Muon은 행렬 차원의 preconditioning으로 업데이트 방향의 구조를 균질화한다. 이로 인해 두 옵티마이저는 서로 다른 implicit bias를 형성하고, pretrained 가중치의 구조에 서로 다른 영향을 준다.
2. 해결 원리: LoRA를 적용해 pretrained 가중치를 고정하고 업데이트를 저랭크 어댑터로 제한하면, Muon과 Adam 간의 편향 차이가 줄어들어 미스매치를 완화한다. 이로써 Muon의 파인튜닝 효율을 유지하면서 pretrained 지식을 보존할 수 있다.
3. 달라지는 점: 실험적으로 LoRA는 Muon-Pretrained에서 성능 격차를 크게 좁히고, 특정 랭크 구간에서 Muon이 Adam과 거의 동등하거나 더 나은 성능을 낸다. 랭크가 증가하면 업데이트의학적 폭이 확대되어 미스매치가 다시 존재할 수 있지만, mild한 미스매치에서는 여전히 강건하다.

• 초기 설정: NanoChat(561M 파라미터) 두 모델을 Muon과 Adam으로 각각 프리트레이닝하고, WikiText-2에서 전체 파인튜닝 및 LoRA(r=8, α=16)로 비교한다. 3시드 평균을 사용한다.
• 이론적 분석: SignGD를 Adam의 프록시로 두고 Muon의 임의적 바이어스를 Toy 모델에서 분석한다. Wt+1 = Wt − ηt ortho(∇L(Wt))의 극한 해를 통해 두 최적해가 서로 다른 노름 제약을 따른다는 것을 보인다.
• LoRA의 효과 분석: LoRA가 update budget을 r에 따라 제약하고, worst-case mismatch inflation이 r에 비례해 조정됨을 이론적으로 보이고 실험으로 확인한다.
• 벤치마크 및 랭크 연구: GLUE/NLG/Image Classification에서 full fine-tuning과 LoRA 조합 비교. LoRA Variant(PtSSA, rsLoRA, LoRA-One, AdaLoRA 등)의 Muon 적용 시 효과를 비교한다.

NLU(GLUE, T5-Base): Full-Adam 89.14, Full-Muon 88.77, LoRA-Adam 88.93, LoRA-Muon 88.97, LoRA-Muon-PE 89.20(최고)으로 LoRA-Muon-PE가 최고 성능이다. Full 파인튜닝에서 Muon이 Adam 대비 소폭 열세이나 LoRA를 사용하면 차이가 줄어들고, LoRA-Muon-PE가 최고치를 기록한다. NLU(NLG): LLama 2-7B의 Math/Code/Commonsense에서 Full-Adam 54.92, Full-Muon 53.10, LoRA-Adam 51.62, LoRA-Muon 52.15, LoRA-Muon-PE 51.82. LoRA를 통한 Muon의 미스매치는 줄고, LoRA-Muon이 Math에서 LoRA-Adam과 근접, Code/Commonsense에서 우수한 측면을 보인다. Code-Feedback에서 13B 규모 HumanEval에서 LoRA-Muon이 LoRA-Adam과 비등하거나 우수한 경향을 보였다. CV(CLIP ViT-B/32): 평균 성능은 Full-Adam 대비 Full-Muon의 차이가 작고, LoRA 설정에서 Muon 계열이 Adam 대비 우수하며, LoRA-Muon-PE가 평균적으로 최고를 기록한다. 메타 분석으로 Muon의 LoRA 기반 격차 감소 효과의 집계값은 Muon 0.72% (95% CI: [0.41, 1.04], p<0.001), Muon-PE 0.83% (95% CI: [0.45, 1.20], p<0.001)로 통계적으로 유의하게 감소한다. LoRA 랭크 연구: 메타Math/Code-Feedback/StanfordCars에서 LoRA-Muon은 랭크 2–64에서 LoRA-Adam 대비 견고한 우위를 보이기도 하고, 랭크 증가 시 차이가 감소하거나 역전되는 경향이 있다. StanFordCars에서 LoRA-Muon이 대체로 우수하다. Catastrophic Forgetting: LoRA-Muon은 LoRA-Adam보다 Forgetting이 작아지는 경향을 보이며, 고랭크에서 Forgetting이 증가하는 반면 LoRA-Muon은 더 느리게 증가한다. LoRA-Granularity: LoRA 매트릭스의 스펙트럼 특성에서 Muon은 더 높은 stable rank와 엔트로피를 유지하며, 이는 pretrained 지식 보존에 유리하다. 컴퓨팅 효율: LoRA 하에서 Muon은 LoRA-Adam 대비 1.1–1.2× 느려지나 메모리 사용은 절감되며, CLIP에서 1.0–1.2× 이다. Muon의 옵티마이저 상태 메모리 절감은 약 50% 수준이다.

• Muon 아키텍처: 2D 가중치 행렬에 대해 NS(Mt)로 orthogonalization을 수행하고 Ot = NS(Mt), Mt = β Mt−1 + Gt, Wt+1 = Wt − η Ot를 사용한다. 이로써 업데이트 방향의 singular values가 거의 균등해지도록 한다. PE(POLAR EXPRESS) 버전은 계수들을 적응적으로 조정한다.
• LoRA: W0에 대해 W = W0 + ∆W = W0 + ˜αBA를 도입, B ∈ Rm×r, A ∈ Rr×n, r ≪ min(m,n). α˜은 보통 α/r 또는 α/√r로 설정한다. LoRA에서는 W0는 고정되고 A와 B만 학습한다.
• 이론적 분석: SignGD를 Adam의 프록시로 두고 Muon의 임의 편향을 toy linear 회귀에서 분석한다. Wt가 점진적으로 W에 수렴하며, W은 max-norm 최소해 또는 spectral-norm 최소해를 달성한다. LoRA는 고정 서브스페이스 가정에서 업데이트 예산 내에서 exact-fit에 도달하도록 한다.
• LoRA 변형: rsLoRA, LoRA-One, PiSSA 등은 LoRA의 업데이트 규모, 초기화, 또는 특성 차원을 다르게 적용하나 Muon-PE가 전체에서 가장 일관된 성능 상승을 보인다.

이론적 분석은 Toy 모델과 라인-설정에 한정되며, 대형 모델의 일반화에 대한 확장된 이론은 여전히 연구 중이다. 또한 Muon 구현은 분산 training 프레임워크(FSDP/DeepSpeed Zero)에서의 동일성 보장이 제한적이며, 실제 대형 모델에서의 효율성은 구현 상황에 따라 달라질 수 있다.