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
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 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 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(Kmn⋅min(m,n))
Kmn⋅min(m,n)
TEON
O(Kmn⋅min(m,nK))
Kmn⋅min(m,nK)
TUON
O(Kmn⋅min(m,n)+K2mn)
Kmn⋅min(m,n)
Table 1: Per-iteration computational complexity for a single tensor group of K×m×n .
Figure 2: Training curves of Qwen models.
Optimizer
124M GPT-2
350M GPT-2
0.6B Qwen3
4B Qwen3
Val
Δ
Val
Δ
Val
Δ
Val
Δ
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
+0.0127
2.6925
−0.0382
2.8887
−0.0765
2.7598
−0.0510
Table 2: Validation loss on GPT-2 (124M, 350M; 50k iterations) and Qwen3 (0.6B, 4B; 20k iterations). Δ=Val−ValAdamW ; negative Δ indicates improvement. Best per column in bold. Lower is better.
Model
Muon
TEON mode-1
TEON mode-2
TUON DCT
TUON DecorQ
K=2
K=full
K=2
K=full
K=2
K=full
K=2
K=full
124M GPT-2
3.0583
3.0255
3.0513 (↑)
3.0332
3.0789 (↑)
3.0309
3.0295 (↓)
3.0299
3.0233 (↓)
350M GPT-2
2.8427
2.8427
2.8617 (↑)
2.8643
2.8762 (↑)
2.8150
2.8218 (↑)
2.8183
2.8143 (↓)
0.6B Qwen3
2.9056
2.8994
2.9323 (↑)
2.8978
2.9583 (↑)
2.9203
2.9210 (↑)
2.8938
2.8887 (↓)
4B Qwen3
2.7636
2.7627
2.7649 (↑)
2.7604
2.7640 (↑)
2.7606
2.8100 (↑)
2.7629
2.7598 (↓)
Table 3: Effect of stacking size K . ↑ / ↓ denote K=full is worse/better than K=2 . 20k iterations.
Model
Time per Iteration (ms)
124M GPT-2
350M GPT-2
0.6B Qwen3
4B 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
Figure 3: Illustrations of T-SVD
optimizer
GPT-2 124M
GPT-2 350M
Qwen3 0.6B
Qwen3 4B
AdamW
2.4×10−3
2×10−3
1×10−3
3×10−4
Muon
1×10−3
1×10−3
2×10−3
1×10−3
TEON
1×10−3
1×10−3
2×10−3
1×10−3
TUON
1×10−3
1×10−3
2×10−3
1×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 GPT-2
350M GPT-2
0.6B Qwen3
4B Qwen3
Val
Δ
Val
Δ
Val
Δ
Val
Δ
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
+0.0127
2.6925
−0.0382
2.8887
−0.0765
2.7598
−0.0510
Appendix
Table 6: Validation loss on GPT-2 (124M, 350M; 50k iterations) and Qwen3 (0.6B, 4B; 20k iterations) with different optimizers. Δ=Val−ValAdamW within the same model and horizon; negative Δ 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 Q on full layer stacking.
Model
Muon
TEON mode-1
TEON mode-2
TUON DCT
TUON DecorQ
K=2
K=full
K=2
K=full
K=2
K=full
K=2
K=full
124M GPT-2
3.0583
3.0255
3.0513 (↑)
3.0332
3.0789 (↑)
3.0309
3.0295 (↓)
3.0299
3.0233 (↓)
350M GPT-2
2.8427
2.8427
2.8617 (↑)
2.8643
2.8762 (↑)
2.8150
2.8218 (↑)
2.8183
2.8143 (↓)
0.6B Qwen3
2.9056
2.8994
2.9323 (↑)
2.8978
2.9583 (↑)
2.9203
2.9210 (↑)
2.8938
2.8887 (↓)
4B Qwen3
2.7636
2.7627
2.7649 (↑)
2.7604
2.7640 (↑)
2.7606
2.8100 (↑)
2.7629
2.7598 (↓)
Appendix
Table 8: Effect of stacking size K . (↑) / (↓) denote whether the full- K setting is worse / better than K=2 . 20k iterations.
Model
Muon
TEON mode-1
TEON mode-2
TUON DecorQ
Module QKV
Module All
Module QKV
Module All
Module QKV
Module All
124M GPT-2
3.0583
3.0255
3.0531
3.0332
3.1027
3.0225
3.0233
350M GPT-2
2.8258
2.8427
2.8637
2.8643
2.9144
2.8291
2.8143
0.6B Qwen3
2.9056
2.8994
2.9023
2.8978
2.9018
2.8897
2.8887
4B Qwen3
2.7636
2.7627
2.7644
2.7604
2.7639
2.7606
2.7598
Appendix
Table 9: Effect of layer type. ModuleQKV mirrors the attention-only scope of existing tensor-aware optimizers, while ModuleAll 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 mode-1
TEON mode-2
124M GPT-2
3.0583
3.0255
3.0332
350M GPT-2
2.8258
2.8427
2.8643
0.6B Qwen3
2.9056
2.8994
2.8978
4B 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)
124M GPT-2
350M GPT-2
0.6B Qwen3
4B 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.
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 ⋆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.
Dongyang Li, Mingsong Yan, Zheng Zhang +1
University of California, Santa Barbara Santa Barbara, CA 93106, USA
Common first-order optimizers, such as Adam, implicitly treat each parameter block as an unstructured vector, which disregards the multilinear weight structure present in many modern machine learning models. Recent work has shown that exploiting matrix structure can improve optimization dynamics. A notable example is Muon, which performs steepest descent under the spectral norm constraint. We take the next step and introduce Tensorion, a tensor-aware optimizer that extends Muon's constrained optimization perspective from matrices to higher-order tensors. Tensorion is built around a linear minimization oracle (LMO) over a tensor norm ball. The norm is carefully chosen to balance two objectives: tightly bounding the tensor spectral norm, while still keeping the LMO tractable. This LMO becomes computable because it reduces to operations on adaptively selected unfolding matrices. Notably, when restricted to order-2 tensors (i.e., matrices), Tensorion recovers Muon exactly. Experiments on tensor-based computer vision problems suggest that Tensorion can offer improved convergence behavior and more stable gradient updates compared with Adam-based and existing tensor-aware baselines in the evaluated settings.
Vladimir Bogachev, Vladimir Aletov, Alexander Molozhavenko +2
HSE University · Basic Research of Artificial Intelligence Laboratory (BRAIn Lab)
Muon has motivated a growing family of optimizers that compose multiple matrix normalizations, but these methods remain fragmented and lack a unified perspective. We introduce chained linear minimization oracles (chained LMOs), which cast these methods as compositions of LMOs. Despite their empirical success, many chains fall outside the standard LMO framework and can diverge on smooth convex objectives. To explain why composition can nevertheless help, we turn to linear associative memory and show that chaining can improve over Muon under anisotropic embeddings. Empirically, we propose TensorChain, a novel optimizer within the framework that stacks compatible weight matrices across different layers and normalizes the 3d tensor across its axes. In Qwen3 0.6B and 1.7B pretraining, TensorChain outperforms all chained baselines in average token efficiency, with average token savings of 9.6% over Muon at matched validation loss.
Sungyoon Kim, Kaan Ozkara, Youngsuk Park
Department of Electrical Engineering, Stanford University · Amazon Annapurna Labs