MUON-RT Reusing the Teacher's Trajectory for Momentum-Free Muon for LLM Self-Distillation
EMA 교사 모델의 궤적을 재사용하여 모멘텀 버퍼 없이도 성능을 높이고 메모리 사용량을 절감하는 효율적인 LLM 자기 증류 기법
논문이 다루는 내용
LLM의 자기 증류(Self-distillation) 과정에서 EMA 교사 모델을 사용하더라도 별도의 모멘텀 버퍼를 유지해야 하므로 지속적인 학습 메모리 부담이 발생합니다. 이를 해결하기 위해 교사와 학생 모델 사이의 변위(displacement)를 모멘텀 대신 재사용하는 Muon-RT 방식을 제안합니다. 이 방식은 교사-학생 간의 차이가 업데이트 이력을 인코딩하고 노이즈를 평활화한다는 점을 활용하여, Newton-Schulz 직교화 전후로 이를 현재 그래디언트와 결합합니다. 실험 결과, Qwen2.5 및 Olmo-3 모델을 대상으로 한 자기 증류 작업에서 기존 AdamW 및 모멘텀 기반 Muon 대비 정확도와 메모리 효율성을 동시에 개선했습니다. 결과적으로 모멘텀 버퍼를 제거하면서도 더 높은 성능을 달리는 것을 입증했습니다.
핵심 결과
-
EMA 교사 모델의 궤적을 활용하여 별도의 모멘텀 버퍼 없이 업데이트 이력을 확보
-
Newton-Schulz 직교화와 결합된 새로운 그래디언트 혼합 방식 제안
-
자기 증류(Self-distillation) 환경에서 메모리 사용량 감소와 성능 향상을 동시에 달성
실무에서 볼 만한 점
LLM 학습 시 발생하는 막대한 메모리 점유 문제를 해결하면서도 성능 저하 없이 효율적인 파인튜닝이 가능함을 보여줍니다.
읽을 때 확인할 점
-
EMA 기반 자기 증류 파이프라인에 Muon-RT 로직 적용 테스트
-
기존 AdamW/Momentum 기반 최적화 기법과 메모리/성능 벤치마크 비교
-
Newton-Schulz 직교화 연산 비용이 전체 학습 시간에 미치는 영향 분석