LLM 자기증류 학습의 메모리 효율을 높인 뮤온 최적화 기법
원제 MUON-RT Reusing the Teacher's Trajectory for Momentum-Free Muon for LLM Self-Distillation
추천 20댓글 2
Key Point
메모리 오버헤드를 크게 줄이면서도 LLM 학습 성능을 개선하는 옵티마이저 개선으로, 제한된 리소스에서 대규모 모델을 학습하려는 개발자들에게 직접적인 영향이 있다.
핵심 요약
- 자기증류(Self-Distillation) 방식의 LLM 학습에서 지수이동평균(EMA) 기반 교사 모델이 이미 학생 업데이트 이력을 보유하고 있는데도 기존 뮤온 옵티마이저는 별도의 모멘텀 버퍼를 유지해 메모리를 낭비한다.
- 뮤온-알티(Muon-RT)는 교사-학생 간 변위(displacement)를 재사용해 숨겨진 가중치 행렬의 모멘텀 버퍼를 제거한다.
- 이 변위가 적용된 업데이트의 지수가중 이력을 인코딩하며 시간적으로 상관없는 노이즈를 평활한다는 것을 로컬 모델 가정 하에서 증명했다.
- 뮤온-알티는 변위의 방사형 성분을 제거하고 스펙트럼 정규화를 적용한 후 뉴턴-슐츠(Newton-Schulz) 직교화 이후 현재 그래디언트와 혼합한다.
- 명시적 정칙성과 정렬 가정 하에서 기댓값-목적함수 간격이 선형으로 수렴해 잔여 오차 수준에 도달함을 수학적으로 확립했다.
- 자기증류 파인튜닝(SDFT)으로 Qwen2.5-7B-Instruct를 학습할 때 뮤온-알티는 AdamW 대비 평균 정확도를 2.6포인트 개선하면서 지속 메모리를 42% 줄였으며, 기존 모멘텀 뮤온 대비 5.4포인트 개선하면서 27% 줄였다.
- 자기증류 정책 최적화(SDPO)로 올모-3-7B-인스트럭트(Olmo-3-7B-Instruct)를 학습할 때는 AdamW 대비 평균 정확도를 2.0포인트 개선했다.