提出TEON优化器，让大模型训练更稳定高效
TEON: Tensorized Orthonormalization Beyond Layer-Wise Muon for Large Language Model Pre-Training
- 将梯度正交化从单层扩展到全网络张量结构
- 在多种模型规模下均降低困惑度，提升训练稳定性
- 适合追求训练收敛性与鲁棒性的大模型研究者
Muon优化器通过逐层矩阵级梯度正交化，在大语言模型预训练中表现优异。本文提出TEON，一种对Muon的严格推广，将神经网络梯度建模为高阶张量结构，实现跨层正交化。理论分析表明TEON具有更强的收敛保证，并基于此设计了实用版本，进行了消融实验验证。在GPT风格（130M~774M参数）和LLaMA风格（60M~1B参数）模型上测试，结果表明TEON在不同模型规模下均持续降低训练与验证困惑度，在多种近似SVD方案下也表现出强鲁棒性。
原文摘要 · Abstract (English)
The Muon optimizer has demonstrated strong empirical performance in pre-training large language models 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。