Dion2通过采样行列简化Muon优化器的正交化,提升可扩展性。
Dion2: A Simple Method to Shrink Matrix in Muon
- 每轮仅对部分行或列进行正交化,实现稀疏更新。
- 相比之前方法,计算与通信开销显著降低。
- 适合大规模模型训练,尤其关注效率的场景。
Muon优化器具有出色的实证性能和理论基础,但其正交化步骤的超线性成本随规模增长而加剧。为缓解此问题,已有研究尝试减小进入正交化步骤的矩阵规模。本文提出Dion2,一种比先前方法更简单的矩阵压缩方法。其核心思想是在每轮迭代中选取部分行或列进行正交化,该采样过程使更新稀疏,从而降低计算与通信开销,提升Muon的可扩展性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。