跳到主要内容
返回时间线
arXiv来源发表:

大模型预训练优化器不再逐层正交化:TUON把层叠成张量,四档模型验证损失均低于Muon与TEON

核心概要

大模型预训练优化器从逐层处理矩阵转向把同一模块的多层动量堆成三阶张量。此前按层正交化的Muon忽略层间相关,而基于矩阵展开的张量优化器只能处理两层堆叠且局限于QKV。TUON沿层方向做酉变换、对每个变换后切片单独极化再逆变换,在GPT-124M/350M与Qwen-0.6B/4B上验证损失均低于Muon和TEON,每步耗时与二者相当。

Source-provided article image: TUON: A Tensor-Aware Optimizer Utilizing Tensor-SVD and Polarization for Large Language Model Pretraining
Figure 1 · arXiv · 第 5 页

深度剖析

大模型预训练优化器开始把Transformer同一模块的多层权重当作一个整体来更新,而不只是逐层正交化。 Muon这类矩阵感知优化器逐层处理,忽略层与层之间的相关性;TEON这类张量感知优化器虽跨层,却依赖把张量展开成矩阵的代数。 在GPT-2 124M/350M与Qwen3 0.6B/4B四个设置上,TUON的验证损失分别为2.9357、2.6925、2.8887、2.7598,均低于Muon的3.0488、2.7404、2.9056、2.7636和TEON的2.9442、2.6935、2.8978、2.7604。

做法是先把层方向混合、再逐片分离:沿层方向做一次酉变换,让各层信息在变换域里重新组合。 随后对每个变换后的切片单独取极因子,也就是把该矩阵的非零奇异值全部置一、只保留方向的正交化操作,再逆变换回原张量域。 默认变换DecorQ由动量张量的层Gram矩阵做SVD得到,训练中每200步刷新一次;在Qwen3-4B上,Newton–Schulz与PolarExpress两种极化近似在两层堆叠下分别为2.7629和2.7638,全层堆叠下分别为2.7598和2.7569。

层数堆得更多不再拖累性能,这是此前展开式张量优化器做不到的。 TEON在两层堆叠时最好,堆到全层后性能下降;TUON配合DecorQ在四个模型上都从两层到全层继续改善。 GPT-2 124M从3.0299降到3.0233,GPT-2 350M从2.8183降到2.8143,Qwen3 0.6B从2.8938降到2.8887,Qwen3 4B从2.7629降到2.7598。

启示与展望

结果覆盖GPT-2 124M/350M与Qwen3 0.6B/4B四个模型规模、OpenWebText与C4两个预训练语料,以及注意力投影与QKVO加MLP两类模块范围;所有跨方法比较都在同一模型家族和同一训练步数内进行。

从两层堆叠扩到全层堆叠带来的改善是真实的但有界,作者指出其仍低于理论上的最大收益,并把这归因于层间差异增大后共享结构逐渐失效。消融实验多数只训练20k步,与主实验的50k步GPT-2训练设置不同。