# TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for Large Language Model Pre-Training

arXiv: https://arxiv.org/abs/2601.23261 · HTML: https://mlsanity.com/paper/2601.23261/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 41 references that are in the map.
- First posted on arXiv 30 Jan 2026.

## Builds on (most cited first, 28 of 28)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale](https://mlsanity.com/paper/2406.17557/) (2024)
- [Large Batch Optimization for Deep Learning: Training BERT in 76 minutes](https://mlsanity.com/paper/1904.00962/) (2019)
- [Kimi K2: Open Agentic Intelligence](https://mlsanity.com/paper/2507.20534/) (2025)
- [GLM-4.5: Agentic, Reasoning, and Coding (ARC) Foundation Models](https://mlsanity.com/paper/2508.06471/) (2025)
- [GaLore: Memory-Efficient LLM Training by Gradient Low-Rank Projection](https://mlsanity.com/paper/2403.03507/) (2024)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Kimi-VL Technical Report](https://mlsanity.com/paper/2504.07491/) (2025)
- [Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training](https://mlsanity.com/paper/2305.14342/) (2023)
- [Gradient Descent Happens in a Tiny Subspace](https://mlsanity.com/paper/1812.04754/) (2018)
- [SOAP: Improving and Stabilizing Shampoo using Adam](https://mlsanity.com/paper/2409.11321/) (2024)
- [Training Deep Learning Models with Norm-Constrained LMOs](https://mlsanity.com/paper/2502.07529/) (2025)
- [The Polar Express: Optimal Matrix Sign Methods and Their Application to the Muon Algorithm](https://mlsanity.com/paper/2505.16932/) (2025)
- [Scaling Laws for Precision](https://mlsanity.com/paper/2411.04330/) (2024)
- [Understanding Gradient Orthogonalization for Deep Learning via Non-Euclidean Trust-Region Optimization](https://mlsanity.com/paper/2503.12645/) (2025)
- [NorMuon: Making Muon more efficient and scalable](https://mlsanity.com/paper/2510.05491/) (2025)
- [Preconditioned Stochastic Gradient Descent](https://mlsanity.com/paper/1512.04202/) (2015)
- [SLTrain: a sparse plus low-rank approach for parameter and memory efficient pretraining](https://mlsanity.com/paper/2406.02214/) (2024)
- [MuonBP: Faster Muon via Block-Periodic Orthogonalization](https://mlsanity.com/paper/2510.16981/) (2025)
- [CoLA: Compute-Efficient Pre-Training of LLMs via Low-Rank Activation](https://mlsanity.com/paper/2502.10940/) (2025)
- [Preconditioner on Matrix Lie Group for SGD](https://mlsanity.com/paper/1809.10232/) (2018)
- [Curvature-Informed SGD via General Purpose Lie-Group Preconditioners](https://mlsanity.com/paper/2402.04553/) (2024)
- [Black Box Lie Group Preconditioners for SGD](https://mlsanity.com/paper/2211.04422/) (2022)
- [LaX: Boosting Low-Rank Training of Foundation Models via Latent Crossing](https://mlsanity.com/paper/2505.21732/) (2025)
- [Stochastic Hessian Fittings with Lie Groups](https://mlsanity.com/paper/2402.11858/) (2024)

## Built on by (most cited first, 6 of 6)

- [Muon2: Boosting Muon via Adaptive Second-Moment Preconditioning](https://mlsanity.com/paper/2604.09967/) (2026)
- [MUON+: Towards More Effective Muon via One Additional Normalization Step for LLM Pre-training](https://mlsanity.com/paper/2602.21545/) (2026)
- [CacheMuon: Using Temporal Preconditioning To Approximate Polar Factor](https://mlsanity.com/paper/2606.16371/) (2026)
- [Hierarchical Muon: Tiled Newton-Schulz Updates for Efficient Muon Optimization](https://mlsanity.com/paper/2606.27216/) (2026)
- [A Note on Stability for Orthogonalized Matrix Momentum with Client Sampling](https://mlsanity.com/paper/2606.01720/) (2026)
- [Optimizing Large Language Models with Chained LMOs](https://mlsanity.com/paper/2610.10975/) (2026)

## Abstract (arXiv)

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.