Muown은 암묵적으로 각도 단계 크기 감쇠 (Angular Step-size Decay)를 수행한다

Muon 및 Muown과 같은 행렬 인식 (Matrix-aware) 옵티마이저들은 최근 Transformer의 사전 학습 (pre-training)에서 강력한 경험적 성능을 보여주었습니다. 특히, Muown은 각 가중치 행렬을 행 크기 (row magnitudes)와 정규화되지 않은 방향 변수 (un-normalized direction variable)로 분리하며, 전자는 Adam으로, 후자는 Muon으로 업데이트합니다. 우리는 Muown의 방향 업데이트가 정규화된 방향에 대한 리만 단계 (Riemannian step)와 동일하며, 정규화되지 않은 파라미터화의 크기는 오직 각도 단계 크기 (angular step size)만을 조절한다는 것을 보여줍니다. 이는 Muown의 단계 크기 안정성을 설명하며, 각도 단계 크기를 명시적으로 만드는 것을 제안합니다. 그 결과물인 AngularMuown은 정규화된 방향에 대해 직접 최적화하며, 방사형 크기 (radial magnitude) 업데이트와 분리되어 스케줄링 가능한 각도 승수 (angular multiplier)를 사용합니다. AngularMuown은 Muown보다 성능을 개선하였으며, 이 글을 쓰는 시점에서 예비 버전은 변형된 nanoGPT 스피드러닝 경연 대회의 옵티마이저 카테고리에서 선두를 달리고 있습니다. Qwen2-0.5B 및 1.1B 파라미터 혼합 전문가 (Mixture-of-Experts, MoE) 모델에 대한 추가 실험은 이 알고리즘이 소규모 모델을 넘어 확장 가능하다는 것을 확인시켜 줍니다. 알고리즘의 구현체는 https://github.com/fhueb/angular-muown 에서 확인할 수 있습니다.