Saved in:
| Main Authors: | , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2504.05295 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866914037368881152 |
|---|---|
| author | Ahn, Kwangjun Xu, Byron Abreu, Natalie Fan, Ying Magakyan, Gagik Sharma, Pratyusha Zhan, Zheng Langford, John |
| author_facet | Ahn, Kwangjun Xu, Byron Abreu, Natalie Fan, Ying Magakyan, Gagik Sharma, Pratyusha Zhan, Zheng Langford, John |
| contents | Orthonormalized updates accelerate training, improve stability, and enable robust hyperparameter transfer, but existing methods like Muon rely on dense matrix operations that clash with sharded weights in large-scale LLM training, causing high compute and communication cost. We introduce Dion (Distributed Orthonormalization), a scalable and efficient update rule that replaces Newton-Schulz iteration with amortized power iteration on a momentum buffer, avoiding full-matrix reconstruction and integrating cleanly with weight sharding. The rank-fraction parameter with error feedback enables low-rank updates that balance quality with significant cost savings. On language models from 160M to 3B parameters, Dion retains the benefits of orthonormalized updates, while markedly reducing wall-clock time at scale, making it a practical optimizer for next-generation foundation models. Code is available at: https://github.com/microsoft/dion/ |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2504_05295 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | Dion: Distributed Orthonormalized Updates Ahn, Kwangjun Xu, Byron Abreu, Natalie Fan, Ying Magakyan, Gagik Sharma, Pratyusha Zhan, Zheng Langford, John Machine Learning Artificial Intelligence Optimization and Control Orthonormalized updates accelerate training, improve stability, and enable robust hyperparameter transfer, but existing methods like Muon rely on dense matrix operations that clash with sharded weights in large-scale LLM training, causing high compute and communication cost. We introduce Dion (Distributed Orthonormalization), a scalable and efficient update rule that replaces Newton-Schulz iteration with amortized power iteration on a momentum buffer, avoiding full-matrix reconstruction and integrating cleanly with weight sharding. The rank-fraction parameter with error feedback enables low-rank updates that balance quality with significant cost savings. On language models from 160M to 3B parameters, Dion retains the benefits of orthonormalized updates, while markedly reducing wall-clock time at scale, making it a practical optimizer for next-generation foundation models. Code is available at: https://github.com/microsoft/dion/ |
| title | Dion: Distributed Orthonormalized Updates |
| topic | Machine Learning Artificial Intelligence Optimization and Control |
| url | https://arxiv.org/abs/2504.05295 |