vix.ing · top · new · best · stats · spec

Dion2: A Simple Method to Shrink Matrix in Muon

2025/12/01 by Kwangjun Ahn, Noah Amsel, Ahn, Kwangjun +3 · 1 citation
Engineering · Physics and Astronomy · #Muon and positron interactions and applications #Particle physics theoretical and experimental studies #Neutrino Physics Research

paper · doi:10.48550/arxiv.2512.16928

Abstract

The Muon optimizer enjoys strong empirical performance and theoretical grounding. However, the super-linear cost of its orthonormalization step introduces increasing overhead with scale. To alleviate this cost, several works have attempted to reduce the size of the matrix entering the orthonormalization step. We introduce Dion2, a much simpler method for shrinking the matrix involved in Muon's computation compared to prior approaches. At a high level, Dion2 selects a fraction of rows or columns at each iteration and orthonormalizes only those. This sampling procedure makes the update sparse, reducing both computation and communication costs which in turn improves the scalability of Muon.

Citations

Cited by

Related