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

Abstract: The Muon optimizer has demonstrated strong empirical performance in pre-training LLMs 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.

• The paper introduces TEON, which stacks layer gradients into tensors to enable joint cross-layer optimization.
• It demonstrates faster convergence and lower perplexity compared to Muon, validated on GPT and LLaMA models.
• TEON achieves up to a √K tighter convergence bound by exploiting explicit gradient alignment through mode-1 matricization.

Tensorized Orthonormalization for Optimizing LLM Pre-Training: An In-Depth Analysis of TEON

Introduction

The efficient pre-training of LLMs remains a critical challenge as architectural and data scale push computational and memory requirements to their limits. While innovations such as AdamW and more recently Muon have improved training dynamics via adaptive and layer-wise orthogonalization strategies, these approaches traditionally treat network layers as independent entities, neglecting inter-layer correlations that can propagate significant optimization structure. "TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for LLM Pre-Training" (2601.23261) addresses this by presenting TEON, a natural generalization of Muon. TEON introduces tensor-level gradient orthogonalization, which stacks gradients from multiple layers into structured tensors, enabling optimization in a parameter space that explicitly encodes cross-layer relationships.

Methodology and Theoretical Framework

Layer-Wise Muon vs. Tensor-Wise TEON

The core innovation of TEON is its structured stacking of layer gradients into higher-order tensors and subsequent orthogonalization along a chosen mode. In contrast to Muon, which orthogonalizes gradients independently per layer as matrices (thus only capturing intra-layer rank structure), TEON enables joint optimization over multiple layers, allowing flow of second-order information and gradient geometry across spatial (layer) dimensions.

Gradient stacking is implemented as follows: for $K$ consecutive layers, gradients $\{\mathbf{G}^{(k)}\}_{k=1}^K$ (e.g., for attention matrices in a Transformer block) are aggregated to form an order-3 tensor ${G} \in \mathbb{R}^{m \times n \times K}$. This tensor is then matricized along a designated mode (e.g., mode-1), producing a matrix of size $m \times (nK)$ (see Figure 1).

Figure 1: Mode-1 matricization where gradient slices along the first index are concatenated into a two-dimensional matrix suitable for matrix orthogonalization.

On this matrix, standard orthogonalization (e.g., via SVD or Newton–Schulz iterations, as in Muon) is performed, and the update direction is folded back to the original tensor shape. This procedure can be viewed as a structured block-wise extension of the classic Muon step.

Figure 2: Schematic of TEON, from stacking gradients over $K$ layers and matricizing them, to applying Muon-like orthogonalization and reconstructing updated tensors for parameter adjustment.

Convergence Guarantees and Norm Hierarchy

TEON's theoretical analysis is anchored in the Non-Euclidean Trust Region (NTR) framework. Specifically, by analyzing the dual-norm structure induced by tensor stacking and matricization, the paper proves that TEON can achieve the same or up to a $\sqrt{K}$-fold tighter convergence bound compared to Muon, as a direct consequence of the norm and smoothness constant comparisons. Crucially, these results only hold when the top singular vectors of the stacked gradient matrices are well aligned; otherwise, the benefit gracefully degrades. The analysis provides explicit guidance for (a) which layer types to stack (preference for those with high cross-layer subspace alignment, e.g., Q/K/V attention blocks), (b) how many layers to group (balancing alignment versus potential maximum gain), and (c) which matricization mode yields the best improvement (mode-1 for strong right singular vector alignment).

Empirical Evaluation

Pre-Training Results on GPT and LLaMA Models

TEON is benchmarking on both GPT-style and LLaMA-style transformers, over a model size range from 60M to 1B parameters and dataset sizes up to 13B tokens (FineWeb). The comparison includes AdamW, Muon (with various fast SVD approximations), and TEON.

Figure 3: Validation perplexity (PPL) for GPT-Small on FineWeb across five runs; TEON outperforms Muon consistently, regardless of orthogonalization method.

Across all configurations, TEON consistently achieves lower final perplexity and converges faster than Muon—both at early and late stages of training (Figure 4). This improvement is robust to the choice of SVD approximation (You, Jordan, PolarExpress) and is most pronounced with higher-quality orthogonalization implementations.

Figure 4: Training loss curves demonstrate TEON’s faster convergence and lower final loss compared to Muon, for the GPT-Small configuration.

On strong baselines (AdamW and Muon), TEON with PolarExpress achieves up to one perplexity point advantage on validation sets in LLaMA pre-training, and similar gaps on GPT-2 derivatives.

Analysis of Mode and Layer Selection

Ablation studies confirm that stacking Q/K/V projections from consecutive layers (which experimentally have high top-right singular vector alignment) yields the strongest gains, validating the theoretical results. In contrast, stacking gradients from diverse components (e.g., MLP blocks) or increasing the group size $K$ beyond 2–4 rapidly degrades gains, as alignment properties deteriorate and SVD approximations become unstable.

Figure 5: Cross-layer singular vector alignment—particularly for Q, K, V matrices—empirically supports mode-1 matricization for TEON.

Momentum and Low-Rank Structure

Additional analyses of gradient momentum tensors (Figure 6) reveal that the structural correlations propagate through training, supporting the modeling assumptions underlying the tensorized approach.

Figure 6: Example visualization of the momentum gradients for Q matrices, illustrating persistent cross-layer structure.

Implications, Potential Limitations, and Future Directions

TEON establishes that principled exploitation of cross-layer gradient structure can enhance optimization landscapes for LLM pre-training beyond what is attainable with purely layer-local adaptations. The practical result is a plug-and-play optimizer that operates at layer-group level, with runtime and memory nearly indistinguishable from Muon but with consistently stronger convergence and downstream perplexity.

The theoretical guidance on group size and mode selection—directed by empirical singular vector alignment—enables practitioners to systematically trade-off between computational burden, alignment, and convergence speed. However, as the number of stacked layers grows, TEON's marginal benefit diminishes, largely due to decreasing gradient subspace alignment and the limits of SVD approximation accuracy for highly unbalanced matrices.

Looking forward, TEON’s structured tensor approach suggests promising extensions in several directions:

• Extension to network blocks beyond transformers: While demonstrated for attention layers, similar strategies could be applied to convolutional or MLP blocks in vision architectures.
• Adaptive alignment-based optimization: Online monitoring of singular vector alignment could dynamically adjust tensor stacking configuration during training.
• Integration with global second-order approximations: Bridging tensorized orthogonalization with quasi-Newton or full Fisher methods for further gains.

Additionally, as resource constraints tighten and model scales increase, structured optimizers like TEON can contribute to reducing energy and training time footprints, facilitating more sustainable model development.

Conclusion

TEON introduces a rigorous, theoretically justified tensor-level generalization of Muon that effectively capitalizes on intrinsic cross-layer correlations in LLM gradients. The algorithm delivers improved convergence and validation perplexity at parity computational cost, validated across architectures and SVD implementations. The study illustrates the ongoing benefit of structural adaptation in optimizers for deep networks and highlights empirically grounded strategies for advancing the state of optimization for large-scale and resource-intensive neural pre-training scenarios.