Hierarchical Muon: 효율적인 Muon 최적화를 위한 타일형 Newton-Schulz 업데이트

Muon 유형의 최적화 도구(optimizers)는 모멘텀-그래디언트(momentum-gradient) 행렬에 유한 Newton-Schulz 맵을 적용하여 밀집 신경망(dense neural-network) 가중치를 위한 업데이트 방향을 구축합니다. $H imes W$ 행렬의 경우, $r=\min{H,W}$ 및 $s=\max{H,W}$일 때, 전체 행렬 Newton-Schulz 업데이트를 $K$ 단계 수행하려면 $O(r^2 s K)$의 연산량이 필요하며, 반복적인 Gram 행렬 곱을 통해 모든 행과 열이 결합됩니다. 우리는 Muon 유형 최적화를 위한 타일형(tiled) Newton-Schulz 방식인 Hierarchical Muon (HiMuon)을 소개합니다. HiMuon은 각 모멘텀-그래디언트 행렬을 $T imes T$ 타일로 분할하고, 각 타일에 동일한 유한 Newton-Schulz 맵을 독립적으로 적용한 뒤 결과를 재조립합니다. 행렬 차원보다 작은 유한한 $T$에 대해, HiMuon은 전체 행렬 업데이트에 대한 수렴하는 근사치라기보다는 국소적 행렬-함수(matrix-function) 맵을 정의합니다. 즉, 스펙트럼 상호작용(spectral interactions)이 타일 내부에서는 보존되지만 타일 경계 사이에서는 무시됩니다. 고정된 유한한 $T$에 대해, 주요 Newton-Schulz 연산량은 $O(H W T K)$로 감소하며, 계산은 독립적인 작은 밀집 행렬 연산으로 분해됩니다. 이러한 구조는 타일 크기 의존적인 GPU 커널, 교차 레이어 배치(cross-layer batching), 메모리 제한 청킹(memory-bounded chunking), 그리고 런타임 타일 크기 스케줄링을 가능하게 합니다. Transformer 학습 및 제어된 행렬-함수 진단 실험을 통해, HiMuon은 테스트된 범위 내에서 전체 행렬 Muon과 유사한 학습 동작을 유지하면서도 최적화 단계(optimizer-step) 효율성을 개선함을 보여줍니다.