TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for Large Language Model Pre-Training
Fuente:
arXiv
Guardado en:
| Autores principales: | , , , , , , , |
|---|---|
| Formato: | Preprint |
| Publicado: |
2026
|
| Materias: | |
| Acceso en línea: | |
| Etiquetas: |
Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
|
| _version_ | 1866908804598202368 |
|---|---|
| author | Zhang, Ruijie Zhao, Yequan Liu, Ziyue Wang, Zhengyang Li, Dongyang Su, Yupeng Liu, Sijia Zhang, Zheng |
| author_facet | Zhang, Ruijie Zhao, Yequan Liu, Ziyue Wang, Zhengyang Li, Dongyang Su, Yupeng Liu, Sijia Zhang, Zheng |
| contents | The Muon optimizer has demonstrated strong empirical performance in pre-training large language models by performing matrix-level gradient (or momentum) orthogonalization in each layer independently. In this work, we propose TEON, a principled generalization of Muon that extends orthogonalization beyond individual layers by modeling the gradients of a neural network as a structured higher-order tensor. We present TEON's improved convergence guarantee over layer-wise Muon, and further develop a practical instantiation of TEON based on the theoretical analysis with corresponding ablation. We evaluate our approach on two widely adopted architectures: GPT-style models, ranging from 130M to 774M parameters, and LLaMA-style models, ranging from 60M to 1B parameters. Experimental results show that TEON consistently improves training and validation perplexity across model scales and exhibits strong robustness under various approximate SVD schemes. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2601_23261 |
| institution | arXiv |
| publishDate | 2026 |
| record_format | arxiv |
| spellingShingle | TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for Large Language Model Pre-Training Zhang, Ruijie Zhao, Yequan Liu, Ziyue Wang, Zhengyang Li, Dongyang Su, Yupeng Liu, Sijia Zhang, Zheng Machine Learning Artificial Intelligence The Muon optimizer has demonstrated strong empirical performance in pre-training large language models by performing matrix-level gradient (or momentum) orthogonalization in each layer independently. In this work, we propose TEON, a principled generalization of Muon that extends orthogonalization beyond individual layers by modeling the gradients of a neural network as a structured higher-order tensor. We present TEON's improved convergence guarantee over layer-wise Muon, and further develop a practical instantiation of TEON based on the theoretical analysis with corresponding ablation. We evaluate our approach on two widely adopted architectures: GPT-style models, ranging from 130M to 774M parameters, and LLaMA-style models, ranging from 60M to 1B parameters. Experimental results show that TEON consistently improves training and validation perplexity across model scales and exhibits strong robustness under various approximate SVD schemes. |
| title | TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for Large Language Model Pre-Training |
| topic | Machine Learning Artificial Intelligence |
| url | https://arxiv.org/abs/2601.23261 |