본 논문은 딥러닝 모델 훈련에 필요한 극 분해(polar decomposition)를 계산하기 위한 새로운 방법인 Polar Express를 소개합니다. 이 방법은 GPU 환경에서 효율적인 연산을 위해 행렬-행렬 곱셈만을 사용하며, 딥러닝의 특성상 높은 정밀도보다 높은 처리량에 중점을 둡니다. Polar Express는 Chen & Chow와 Nakatsukasa & Freund의 연구를 바탕으로, 각 반복에서 minimax 최적화 문제를 해결하여 업데이트 규칙을 조정합니다. 이를 통해 초기 반복과 점근적 단계에서 가능한 한 빠르게 수렴하도록 설계되었습니다. 또한, bfloat16과 같은 유한 정밀도 환경에서도 사용할 수 있도록 구현되었습니다. Muon 훈련 프레임워크에 Polar Express를 통합한 결과, FineWeb 데이터셋의 10억 개의 토큰을 사용하여 GPT-2 모델을 훈련할 때 다양한 학습률에서 기존 방법보다 일관된 검증 손실 감소를 보였습니다.
