Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Ruijie Zhang, Yequan Zhao, Ziyue Liu, Zhengyang Wang, Dongyang Li, Yupeng Su, Sijia Liu, Zheng Zhang

TL;DR

TEON tackles cross-layer gradient correlations in large language model pre-training by generalizing Muon from layer-wise to tensor-wise orthogonalization. It models gradients as a structured tensor and applies mode-i orthogonalization within a Non-Euclidean Trust Region framework, yielding convergence benefits up to $\sqrt{K}$ over Muon under favorable alignment of singular vectors. Empirically, TEON improves training and validation perplexity across GPT- and LLaMA-style models and remains robust to approximate SVD methods, with ablations guiding practical choices (e.g., $K=2$, mode-1, QKV stacking). Overall, TEON offers a principled, scalable optimizer enhancement that leverages cross-layer gradient structure to improve pre-training efficiency for large language models.

Abstract

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.

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

TL;DR

TEON tackles cross-layer gradient correlations in large language model pre-training by generalizing Muon from layer-wise to tensor-wise orthogonalization. It models gradients as a structured tensor and applies mode-i orthogonalization within a Non-Euclidean Trust Region framework, yielding convergence benefits up to over Muon under favorable alignment of singular vectors. Empirically, TEON improves training and validation perplexity across GPT- and LLaMA-style models and remains robust to approximate SVD methods, with ablations guiding practical choices (e.g., , mode-1, QKV stacking). Overall, TEON offers a principled, scalable optimizer enhancement that leverages cross-layer gradient structure to improve pre-training efficiency for large language models.

Abstract

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.
Paper Structure (47 sections, 15 theorems, 56 equations, 6 figures, 16 tables, 1 algorithm)

This paper contains 47 sections, 15 theorems, 56 equations, 6 figures, 16 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Layer-wise Muon
  4. Tensor: notations and definitions
  5. Matricization and folding.
  6. The Teon Method
  7. From Layer-wise to Tensor-wise Modeling: The Teon Optimizer
  8. Implementation of Teon
  9. Theoretical Analysis of Teon
  10. Main Theoretical Results
  11. Definitions and Assumptions
  12. Key Results.
  13. Implication in Algorithm Design
  14. Sketch of Proofs
  15. Sketch of Proof for Theorem \ref{['thm:convergence_bound_main']}.
  16. ...and 32 more sections

Key Result

Theorem 4.5

Consider minimizing $f(\bm{\mathcal{W}}), \bm{\mathcal{W}}\in \mathbb{R}^{m\times n\times K}$ and suppose $f$ is lower bounded by $f^\star$. Define the best iterates of Muon and Teon as $\tau_{\textsc{muon}} := \arg\min_{0\le t < T}\|\nabla f(\bm{\mathcal{W}}_t)\|_{\textsc{muon},*}$ and $\tau_{\text The smoothness constants satisfy $L_{\textsc{teon}} \le L_{\textsc{muon}} \le K \, L_{\textsc{teon}

Figures (6)

  • Figure 1: Pre-training GPT-Small on 10 Billion FineWeb Tokens. We run 5 trials with different random seeds to estimate the standard deviations. Validation perplexity (PPL) is reported; lower is better. Our proposed method Teon, consistently outperforms Muon across different orthogonalization methods. Additional experimental results across various model configurations are presented, with more detailed analysis in Section \ref{['sec:experiment']}.
  • Figure 2: Mode-1 Matricization of a given tensor $\bm{\mathcal{G}}$. $\bm{\mathcal{G}}$ is first sliced into its column fibers, i.e., vectors obtained by fixing all indices except the first, and these fibers are then arranged as columns of a matrix to form the mode-1 unfolding.
  • Figure 3: Overview of Teon. (a) Transformer blocks: Gradients are collected from $K$ successive layers ($K=3$ in this figure). (b) Gradient tensors: Gradients of the same layer type are stacked to form structured gradient tensors. (c) Mode-$i$ orthogonalization: The gradient tensor is matricized along mode $i$, followed by Muon-style orthogonalization on the resulting matrix, which is then used to update the parameters accordingly.
  • Figure 4: Training curves of GPT-Small on 10B tokens from the FineWeb dataset, comparing Muon and Teon using the original Newton–Schulz iteration with 5 steps. (a) At the beginning of training, Teon exhibits faster convergence, achieving lower training loss than Muon. (b) In the later stages, Teon continues to outperform Muon, achieving better final training loss.
  • Figure 5: We perform SVD on the momentum terms $\mathbf{M}_t$ and plot the inner product between top singular vectors of consecutive layer pairs. For QKV matrices, the top right singular vectors are aligned while top left singular vectors remain almost orthogonal. We provide the plot for all matrices pairs in Appendix \ref{['app:Alignment of top singular vectors']}.
  • ...and 1 more figures

Theorems & Definitions (29)

  • Definition 4.1: Muon and Teon norms
  • Theorem 4.5: Convergence Bound
  • Proposition 4.6: Maximal Gain (Informal)
  • Theorem 4.7: NTR convergence under $\|\cdot\|$
  • Lemma 4.8: teon's maximal gain over muon. Proved in Appendix \ref{['app:teon_max_gain_modes']}
  • Definition 1.1: Teon norm family
  • Definition 1.2: Muon norm
  • Lemma 1.3: Dual of $\|\cdot\|_{\textsc{teon}-i}$
  • proof
  • Lemma 1.4: Dual of $\|\cdot\|_{\textsc{muon}}$
  • ...and 19 more