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

📄 arXiv: 2610.35130v1 📥 PDF

作者: Anda Tang, Xuanlin Yang, Junjie Wang, Zhouchen Lin

分类: cs.LG

发布日期: 2026-10-26

💡 一句话要点

提出TUON优化器以解决大规模语言模型预训练效率问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 优化器 张量感知 大规模语言模型 预训练 跨层相关性 性能提升 机器学习

📋 核心要点

  1. 现有的优化器在处理大规模语言模型时效率不足，尤其是未能有效利用层间的相关性。
  2. TUON优化器通过单位变换混合层并进行极化，能够捕捉层间的相关性，同时保持每层的多样性。
  3. TUON在多个模型（如GPT-124M/350M和Qwen-0.6B/4B）上表现出显著的性能提升，验证了其有效性。

📝 摘要（中文）

大规模语言模型（LLM）预训练的效率逐渐受到优化器性能的制约，促使优化器设计从向量感知方法（如AdamW）转向矩阵感知方法。现有的矩阵感知优化器（如Muon）虽然利用了矩阵结构，但未能考虑层间的相关性。因此，论文提出了一种新的张量感知优化器TUON，基于张量奇异值分解（T-SVD）构建。TUON通过单位变换混合层并对每个变换切片进行极化，从而捕捉层间相关性，避免了展开过程中的结构维度不平衡和层间错位。实验结果表明，TUON在多个基准模型上均优于现有方法。

🔬 方法详解

问题定义：论文旨在解决现有优化器在大规模语言模型预训练中效率不足的问题，尤其是未能有效利用层间相关性的问题。现有的矩阵感知优化器往往独立处理各层，忽视了层间的相互影响。

核心思路：TUON优化器的核心思路是通过单位变换将各层混合，然后对每个变换后的切片进行极化，从而捕捉层间的相关性，同时保持每层的多样性。这种混合-分离的策略避免了展开过程中的结构维度不平衡和层间错位。

技术框架：TUON的整体架构包括三个主要模块：首先，通过单位变换对输入张量进行混合；其次，对混合后的切片进行极化处理；最后，逆变换以恢复更新。这一流程确保了层间信息的有效利用。

关键创新：TUON的最重要创新在于其基于张量奇异值分解的变换方法，区别于传统的矩阵展开方法，能够有效捕捉跨层的相关性并提升优化器的性能。

关键设计：TUON的设计包括对单位变换矩阵的选择、极化方法的实现以及更新的计算方式。通过理论证明，TUON在$oldsymbol{Q}$-变换的张量谱范数下实现最陡下降，并提供了收敛性保证。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

TUON在GPT-124M/350M和Qwen-0.6B/4B模型上均表现出显著的性能提升，相较于Muon和TEON，TUON在多个实验中均取得了更好的收敛速度和最终性能，验证了其有效性和适应性。

🎯 应用场景

TUON优化器的潜在应用场景包括大规模语言模型的预训练和微调，尤其是在需要高效处理多层网络结构的任务中。其设计理念和性能提升为未来的优化器研究提供了新的思路，可能推动更复杂模型的训练效率提升。

📄 摘要（原文）

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.