Scaling Muon for Diffusion Transformers

행렬 구조를 활용하는 옵티마이저 Muon을 13억~150억 파라미터 확산 트랜스포머까지 확장해 검증하고, 통신·계산 병목을 주기적 갱신으로 줄여 실제 훈련 시간까지 잡은 논문이다.

이 논문은 한마디로 Muon 옵티마이저를 대형 확산 모델 훈련에 실용화한 연구다. Muon은 행렬 구조를 활용해 AdamW보다 좋은 결과를 내지만, 매 스텝 하는 특수 연산(뉴턴-슐츠 반복)이 큰 모델에서 계산·통신 병목이 된다. 저자들은 1.3B~15B 확산 트랜스포머에서 Muon의 우위가 유지됨을 확인하고(생성 품질 12.9~19.1% 향상), 3스텝마다만 특수 연산을 하고 나머지는 가벼운 행 단위 정규화로 대체하는 방법을 제안했다. 품질은 그대로 유지하면서 옵티마이저 시간을 절반 가까이(46.9~54.3%) 줄였다.

핵심 요약

• 13억~150억 매개변수의 확산 트랜스포머에서 Muon이 AdamW보다 생성 품질 지표를 12.9~19.1% 개선하는 경향을 확인했다.
• 세 단계마다만 전체 스펙트럼 갱신을 하고 나머지는 저비용 행 단위 갱신으로 바꾸는 Periodic Row-wise Muon을 제안했다.
• 옵티마이저 시간은 46.9~54.3%, 전체 학습 단계 시간은 15.7~24.3%, 논리적 통신량은 66.7% 줄었다.
• 90억 매개변수 모델에서는 원래 Muon보다 생성 품질이 4.5% 더 좋아져 비용 절감과 품질 향상을 함께 보였다.
• 실험은 한 확산 모델 계열과 한 데이터셋·해상도, 32노드 H100 구성에 한정됐고 갱신 단계의 전체 모멘텀 통신은 여전히 남는다.

논문 링크

외부 연구를 정리한 자료입니다. HDATF가 발표한 논문이나 제품 성능을 측정한 결과는 아닙니다.