Tensor-aware optimizer improves large language model training efficiency
TUON: A Tensor-Aware Optimizer Utilizing Tensor-SVD and Polarization for Large Language Model Pretraining
Machine Learning
Summary
Training very large language models can be slowed down by inefficiencies in how the training adjusts the model's parameters, called optimization. The authors created a new optimizer named TUON that better captures relationships between layers of the model by using a special kind of tensor math. Unlike older methods that treat each layer separately or use simpler math, TUON mixes and separates layers in a way that keeps important information and works well with many layers. Tests show TUON speeds up training on several large language models compared to previous optimizers.
What this means in practice
- •For machine learning engineers: Optimize large language model training by implementing TUON to improve efficiency and training speed across diverse model layers and modules.
- •For deep learning framework developers: Integrate tensor-transform-based optimization methods like TUON to provide built-in support for scalable and efficient large language model training.
Authors
Anda Tang, Xuanlin Yang, Junjie Wang, Zhouchen Lin
Abstract
Large scale language model~(LLM) pretraining is increasingly bottlenecked by optimizer efficiency, driving optimizer design from vector-aware methods such as AdamW to matrix-aware methods. Matrix-aware optimizers such as Muon exploit matrix structure but treat layers independently, neglecting cross-layer correlations. Tensor-aware optimizers are therefore needed, yet existing ones rely on a matricization-based tensor algebra and remain restricted to shallow stacking and specific modules. In this paper, we propose \textbf{TUON}, a novel optimizer built on a transform-based tensor algebra derived from the Tensor-SVD (T-SVD). Unlike methods that polarize a globally unfolded matrix, TUON mixes layers via a unitary transform $\boldsymbol{Q}$ and then polarizes each transformed slice, and inverts the transform to recover the update. This mix-then-separate strategy captures cross-layer correlations while preserving per-layer diversity, avoiding the structural dimension imbalance and cross-layer misalignment of unfolding. Thus it successfully scales to more layers and diverse modules beyond QKV. Theoretically, we prove that TUON performs steepest descent under a $\boldsymbol{Q}$-transformed tensor spectral norm, and identify its update as the tensor polar factor of the T-SVD. We further establish transform non-redundancy, norm and smoothness comparisons with Muon and TEON, and a convergence guarantee. Empirically, TUON improves consistently over Muon and TEON on GPT-124M/350M and Qwen-0.6B/4B, with ablations confirming the transformation mechanism, cross-layer structure beyond two-layer stacking, adaptability beyond QKV, and robustness to the polarization method.