Dion3: Full-Stack Orthogonal Updates

이 논문은 Muon 옵티마이저의 직교화 계산과 분산 통신을 함께 줄여, 모델 품질을 유지하면서 옵티마이저 단계를 최대 6배 빠르게 만든 Dion3를 제안했다.

이 논문은 한마디로 대규모 언어 모델 학습에서 Muon 옵티마이저의 느린 직교화 계산을 줄이는 방법을 제안한다. 직교화는 가중치 업데이트 방향을 고르게 만드는 계산이다. Dion3는 큰 직사각형 행렬 대신 더 작은 대칭 행렬에서 계산하고, 대칭성을 활용하는 GPU 커널과 여러 행렬을 한 번에 통신하는 방식을 결합한다. 또 매 단계에서 일부 행만 골라 직교화해 비용을 더 줄였다. 여러 모델 크기에서 Muon 계열의 손실을 유지하거나 낮추면서 옵티마이저 단계 시간을 단축했다.

핵심 요약

• Gram Newton-Schulz는 기존 방법을 수학적으로 같은 형태로 다시 쓰고, 큰 직사각형 행렬 대신 작은 대칭 행렬에서 반복 계산해 연산량을 줄인다.
• 대칭성을 활용하는 GPU 커널과 같은 모양의 행렬을 묶어 통신하는 방식을 추가했다. 10억 매개변수 모델의 8분할 설정에서는 묶음 통신만 바꿔 옵티마이저 단계 시간을 80.7밀리초에서 52.1밀리초로 줄였다.
• 매 단계에서 모멘텀 행렬의 일부 행만 골라 직교화해 계산과 통신을 줄였고, 이전 압축 방식인 Dion보다 속도와 손실이 모두 좋아졌다.
• 30억부터 140억 모델을 100억 토큰으로 학습한 실험에서 모두 검증 손실이 낮았고, 140억 모델은 NorMuon보다 손실이 0.027 낮고 평균 정확도는 0.7%포인트 높았다.
• 최대 6배 향상은 전체 학습이 아니라 옵티마이저 단계 기준이며, 저자 분석에서 이 단계가 전체 학습 시간에서 차지하는 비중은 설정에 따라 약 1%부터 17%였다.

논문 링크

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