TUON: A Tensor-Aware Optimizer Utilizing Tensor-SVD and Polarization for Large Language Model Pretraining

Anda Tang; Xuanlin Yang; Junjie Wang; Zhouchen Lin

cs.LGOct 26, 2026

TUON: A Tensor-Aware Optimizer Utilizing Tensor-SVD and Polarization for Large Language Model Pretraining

Authors: Anda Tang, Xuanlin Yang, Junjie Wang, Zhouchen Lin

Organizations: the State Key Lab of General AI School of Intelligence Science and Technology Peking University Beijing, China · the Pazhou Laboratory (Huangpu) the State Key Lab of General AI School of Intelligence Science and Technology the Institute for Artificial Intelligence Peking University Beijing, China

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 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.

Figures & tables

Figure 1: Illustration of the TUON update

Optimizer	Per-iteration complexity	Dominant term
Muon	$O\!\left(K\,mn\cdot\min(m,n)\right)$	$K\,mn\cdot\min(m,n)$
TEON	$O\!\left(K\,mn\cdot\min(m,nK)\right)$	$Kmn\cdot\min(m,nK)$
TUON	$O\!\left(K\,mn\cdot\min(m,n)+K^{2}mn\right)$	$K\,mn\cdot\min(m,n)$

Table 1: Per-iteration computational complexity for a single tensor group of

K\times m\times n

Panel of Figure 2 — Figure 2: Training curves of Qwen models.

Optimizer	124M ${}_{\text{GPT-2}}$		350M ${}_{\text{GPT-2}}$		0.6B ${}_{\text{Qwen3}}$		4B ${}_{\text{Qwen3}}$
Optimizer	Val	$\Delta$	Val	$\Delta$	Val	$\Delta$	Val	$\Delta$
AdamW	2.9230	–	2.7307	–	2.9652	–	2.8108	–
Muon	3.0488	$+0.1258$	2.7404	$+0.0097$	2.9056	$-0.0596$	2.7636	$-0.0472$
TEON	2.9442	$+0.0212$	2.6935	$-0.0372$	2.8978	$-0.0674$	2.7604	$-0.0504$
TUON	2.9357	$\mathbf{+0.0127}$	2.6925	$\mathbf{-0.0382}$	2.8887	$\mathbf{-0.0765}$	2.7598	$\mathbf{-0.0510}$

Table 2: Validation loss on GPT-2 (124M, 350M; 50k iterations) and Qwen3 (0.6B, 4B; 20k iterations).

\Delta=\text{Val}-\text{Val}_{\text{AdamW}}

; negative

\Delta

indicates improvement. Best per column in bold. Lower is better.

Model	Muon	TEON ${}_{\text{mode-1}}$		TEON ${}_{\text{mode-2}}$		TUON ${}_{\text{DCT}}$		TUON ${}_{\text{DecorQ}}$
Model	Muon	K=2	K=full	K=2	K=full	K=2	K=full	K=2	K=full
124M ${}_{\text{GPT-2}}$	3.0583	3.0255	3.0513 $\left(\uparrow\right)$	3.0332	3.0789 $\left(\uparrow\right)$	3.0309	3.0295 $\left(\downarrow\right)$	3.0299	3.0233 $\left(\downarrow\right)$
350M ${}_{\text{GPT-2}}$	2.8427	2.8427	2.8617 $\left(\uparrow\right)$	2.8643	2.8762 $\left(\uparrow\right)$	2.8150	2.8218 $\left(\uparrow\right)$	2.8183	2.8143 $\left(\downarrow\right)$
0.6B ${}_{\text{Qwen3}}$	2.9056	2.8994	2.9323 $\left(\uparrow\right)$	2.8978	2.9583 $\left(\uparrow\right)$	2.9203	2.9210 $\left(\uparrow\right)$	2.8938	2.8887 $\left(\downarrow\right)$
4B ${}_{\text{Qwen3}}$	2.7636	2.7627	2.7649 $\left(\uparrow\right)$	2.7604	2.7640 $\left(\uparrow\right)$	2.7606	2.8100 $\left(\uparrow\right)$	2.7629	2.7598 $\left(\downarrow\right)$

Table 3: Effect of stacking size

K

\uparrow

\downarrow

denote

K=\text{full}

is worse/better than

K=2

. 20k iterations.

Model	Time per Iteration (ms)
Model	124M ${}_{\text{GPT-2}}$	350M ${}_{\text{GPT-2}}$	0.6B ${}_{\text{Qwen3}}$	4B ${}_{\text{Qwen3}}$
AdamW	152.74	379.95	3265.5	5466.9
Muon	160.88	404.81	3324.3	5794.7
TEON	161.81	408.30	3385.7	6081.5
TUON	162.12	411.64	3378.3	6150.7

Table 4: Time per iteration (ms) on four models. Lower values indicate faster training.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

optimizer	GPT-2 124M	GPT-2 350M	Qwen3 0.6B	Qwen3 4B
AdamW	$2.4\times 10^{-3}$	$2\times 10^{-3}$	$1\times 10^{-3}$	$3\times 10^{-4}$
Muon	$1\times 10^{-3}$	$1\times 10^{-3}$	$2\times 10^{-3}$	$1\times 10^{-3}$
TEON	$1\times 10^{-3}$	$1\times 10^{-3}$	$2\times 10^{-3}$	$1\times 10^{-3}$
TUON	$1\times 10^{-3}$	$1\times 10^{-3}$	$2\times 10^{-3}$	$1\times 10^{-3}$

Appendix

Table 5: Peak learning rates for the AdamW baseline on GPT-2 (124M, 350M) and Qwen3 (0.6B,4B), respectively.

Optimizer	124M ${}_{\text{GPT-2}}$		350M ${}_{\text{GPT-2}}$		0.6B ${}_{\text{Qwen3}}$		4B ${}_{\text{Qwen3}}$
Optimizer	Val	$\Delta$	Val	$\Delta$	Val	$\Delta$	Val	$\Delta$
AdamW	2.9230	–	2.7307	–	2.9652	–	2.8108	–
Muon	3.0488	$+0.1258$	2.7404	$+0.0097$	2.9056	$-0.0596$	2.7636	$-0.0472$
TEON	2.9442	$+0.0212$	2.6935	$-0.0372$	2.8978	$-0.0674$	2.7604	$-0.0504$
TUON	2.9357	$\mathbf{+0.0127}$	2.6925	$\mathbf{-0.0382}$	2.8887	$\mathbf{-0.0765}$	2.7598	$\mathbf{-0.0510}$

Appendix

Table 6: Validation loss on GPT-2 (124M, 350M; 50k iterations) and Qwen3 (0.6B, 4B; 20k iterations) with different optimizers.

\Delta=\text{Val}-\text{Val}_{\text{AdamW}}

within the same model and horizon; negative

\Delta

indicates improvement over AdamW. Best per column in bold. Lower is better.

Transform	2k	4k	8k	12k	16k	18k	20k	Final
DCT	3.5411	3.3646	3.1192	2.9444	2.8411	2.8189	2.8051	2.8100
Haar	3.4437	3.1413	2.9302	2.8321	2.7655	2.7540	2.7579	2.7635
DFT	3.3981	3.1002	2.9163	2.8223	2.7599	2.7503	2.7579	2.7631
DecorQ static	3.3999	3.0975	2.9145	2.8206	2.7582	2.7482	2.7569	2.7658
Residual-DecorQ	3.4100	3.1079	2.9204	2.8255	2.7615	2.7511	2.7585	2.7640
DecorQ	3.3899	3.0916	2.9109	2.8185	2.7564	2.7461	2.7545	2.7598

Appendix

Table 7: Effect of transform

\bm{Q}

on full layer stacking.

Model	Muon	TEON ${}_{\text{mode-1}}$		TEON ${}_{\text{mode-2}}$		TUON ${}_{\text{DCT}}$		TUON ${}_{\text{DecorQ}}$
Model	Muon	K=2	K=full	K=2	K=full	K=2	K=full	K=2	K=full
124M ${}_{\text{GPT-2}}$	3.0583	3.0255	3.0513 $\left(\uparrow\right)$	3.0332	3.0789 $\left(\uparrow\right)$	3.0309	3.0295 $\left(\downarrow\right)$	3.0299	3.0233 $\left(\downarrow\right)$
350M ${}_{\text{GPT-2}}$	2.8427	2.8427	2.8617 $\left(\uparrow\right)$	2.8643	2.8762 $\left(\uparrow\right)$	2.8150	2.8218 $\left(\uparrow\right)$	2.8183	2.8143 $\left(\downarrow\right)$
0.6B ${}_{\text{Qwen3}}$	2.9056	2.8994	2.9323 $\left(\uparrow\right)$	2.8978	2.9583 $\left(\uparrow\right)$	2.9203	2.9210 $\left(\uparrow\right)$	2.8938	2.8887 $\left(\downarrow\right)$
4B ${}_{\text{Qwen3}}$	2.7636	2.7627	2.7649 $\left(\uparrow\right)$	2.7604	2.7640 $\left(\uparrow\right)$	2.7606	2.8100 $\left(\uparrow\right)$	2.7629	2.7598 $\left(\downarrow\right)$

Appendix

Table 8: Effect of stacking size

K

\left(\uparrow\right)

\left(\downarrow\right)

denote whether the full-

K

setting is worse / better than

K=2

. 20k iterations.

Model	Muon	TEON ${}_{\text{mode-1}}$		TEON ${}_{\text{mode-2}}$		TUON ${}_{\text{DecorQ}}$
Model	Muon	Module ${}_{\text{QKV}}$	Module ${}_{\text{All}}$	Module ${}_{\text{QKV}}$	Module ${}_{\text{All}}$	Module ${}_{\text{QKV}}$	Module ${}_{\text{All}}$
124M ${}_{\text{GPT-2}}$	3.0583	3.0255	3.0531	3.0332	3.1027	3.0225	3.0233
350M ${}_{\text{GPT-2}}$	2.8258	2.8427	2.8637	2.8643	2.9144	2.8291	2.8143
0.6B ${}_{\text{Qwen3}}$	2.9056	2.8994	2.9023	2.8978	2.9018	2.8897	2.8887
4B ${}_{\text{Qwen3}}$	2.7636	2.7627	2.7644	2.7604	2.7639	2.7606	2.7598

Appendix

Table 9: Effect of layer type.

\text{Module}_{\text{QKV}}

mirrors the attention-only scope of existing tensor-aware optimizers, while

\text{Module}_{\text{All}}

extends the transform to QKVO and MLP. Lower is better; the best result for each model is highlighted in bold. 20k iterations.

Model	Muon	TEON ${}_{\text{mode-1}}$	TEON ${}_{\text{mode-2}}$
124M ${}_{\text{GPT-2}}$	3.0583	3.0255	3.0332
350M ${}_{\text{GPT-2}}$	2.8258	2.8427	2.8643
0.6B ${}_{\text{Qwen3}}$	2.9056	2.8994	2.8978
4B ${}_{\text{Qwen3}}$	2.7636	2.7627	2.7604

Appendix

Table 10: Effect of unfolding mode in TEON. The better mode is highlighted in bold for each model. Lower is better.

Stacking layer number	Method	Final val
2	Newton–Schulz ( $J=5$ )	2.7629
	PolarExpress	2.7638
Full	Newton–Schulz ( $J=5$ )	2.7598
	PolarExpress	2.7569

Appendix

Table 11: Comparison of standard Newton–Schulz and PolarExpress for polar approximation on Qwen3 4B .

Model	Time per Iteration (ms)
Model	124M ${}_{\text{GPT-2}}$	350M ${}_{\text{GPT-2}}$	0.6B ${}_{\text{Qwen3}}$	4B ${}_{\text{Qwen3}}$
AdamW	152.74	379.95	3265.5	5466.9
Muon	160.88	404.81	3324.3	5794.7
TEON	161.81	408.30	3385.7	6081.5
TUON	162.12	411.64	3378.3	6150.7

Appendix

Table 12: Time per iteration (ms) on four models. Lower values indicate faster training. Our method incurs only marginal overhead compared with AdamW, Muon, and TEON.

Language Model Pretraining

Language model pretraining teaches models language patterns and knowledge from large text corpora before they are adapted to specific tasks.

See all 448 papers

T-MUON: Tensorizing Muon via t-SVD for Large Language Model Pre-training
Oct 18, 2026Dongyang Li, Mingsong Yan, Zheng Zhang +1Muon Optimizer Language Model Pretraining
University of California, Santa Barbara Santa Barbara, CA 93106, USA
Tensorion: A Tensor-Aware Generalization of the Muon Optimizer
Jun 24, 2026Vladimir Bogachev, Vladimir Aletov, Alexander Molozhavenko +2Deep Learning Optimization Muon Optimizer
HSE University · Basic Research of Artificial Intelligence Laboratory (BRAIn Lab)
Optimizing Large Language Models with Chained LMOs
Oct 7, 2026Sungyoon Kim, Kaan Ozkara, Youngsuk ParkLMO Methods Deep Learning Optimization
Department of Electrical Engineering, Stanford University · Amazon Annapurna Labs

Oct 18, 2026J/K move · Enter open · S save

T-MUON: Tensorizing Muon via t-SVD for Large Language Model Pre-training

Dongyang Li, Mingsong Yan, Zheng Zhang, Sui Tang

University of California, Santa Barbara Santa Barbara, CA 93106, USA

Muon Optimizer Language Model Pretraining Transformer Training Dynamics

Open paper AskPDF arXiv

Muon has recently emerged as a strong optimizer for large language model pretraining by exploiting matrix structure through polar orthogonalization. However, it processes each parameter matrix independently, whereas recent tensor-aware extensions of Muon jointly update groups of eligible matrix parameters through matrix unfoldings. In contrast, we propose T-MUON, a t-SVD-based generalization that groups structurally related Transformer gradients into third-order tensors and defines the corresponding Muon-style update directly through the

M

-transform. We show that the exact T-MUON update is a steepest-descent direction under the

\star_M

-spectral norm and establish a convergence guarantee. Across GPT-2-style and LLaMA-style models from 51M to 1.3B parameters trained on FineWeb and FineWeb-Edu, head-wise T-MUON with the identity transform consistently achieves lower validation perplexity than AdamW, Muon, and variants using DFT, DCT, or DST transforms. These gains persist in additional experiments with longer context lengths and larger token budgets, although the preferred tensorization mode varies across settings.

Abstract

Language Model Pretraining

Explore similar work

T-MUON: Tensorizing Muon via t-SVD for Large Language Model Pre-training

Tensorion: A Tensor-Aware Generalization of the Muon Optimizer

Optimizing Large Language Models with Chained LMOs