English

Dion2: A Simple Method to Shrink Matrix in Muon

Machine Learning 2025-12-22 v1 Distributed, Parallel, and Cluster Computing

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.

Keywords

Cite

@article{arxiv.2512.16928,
  title  = {Dion2: A Simple Method to Shrink Matrix in Muon},
  author = {Kwangjun Ahn and Noah Amsel and John Langford},
  journal= {arXiv preprint arXiv:2512.16928},
  year   = {2025}
}

Comments

https://github.com/microsoft/dion/

R2 v1 2026-07-01T08:32:17.261Z