[Paper] TEON：张量化正交化 突破层级 Muon 的大语言模型预训练

Source: arXiv - 2601.23261v1

请提供您希望翻译的具体文本内容，我将按照要求保留源链接并进行简体中文翻译。

概述

本文介绍了 TEON（Tensorized Orthonormalization），这是一种新型优化器，它将成功的 Muon 技术从逐层矩阵正交化扩展到对整个网络的张量级梯度处理。通过此方式，TEON 在预训练参数规模从 60 M 到 1 B 的大语言模型（LLM）时，提供了更严格的收敛保证，并始终实现更佳的困惑度。

关键贡献

• 张量级正交化: 将 Muon 的逐层矩阵正交化推广到结构化的高阶张量，以捕获层间梯度关系。
• 理论改进: 提供比 Muon 更强的收敛界限，表明全局正交化更有效地降低梯度方差。
• 实用 TEON 算法: 通过近似 SVD（例如随机幂迭代）推导出计算上可行的实现，并展示其可在标准深度学习工具包中使用。
• 广泛的实证验证: 在 GPT‑style（130 M–774 M）和 LLaMA‑style（60 M–1 B）模型上进行基准测试，实现了所有规模下更低的训练/验证困惑度。
• 鲁棒性分析: 表明 TEON 在各种低秩 SVD 近似下仍然有效，适用于大规模分布式训练。

方法论

1. 梯度张量构建 – 而不是单独处理每层的梯度矩阵 $G_\ell$，TEON 将它们堆叠成一个 3‑D 张量 $\mathcal{G} \in \mathbb{R}^{L \times d_{\text{in}} \times d_{\text{out}}}$（层、输入维度、输出维度）。
2. 张量正交化 – TEON 为 $\mathcal{G}$ 寻找一个正交归一基 $\mathcal{Q}$，满足 $\mathcal{Q}^\top \mathcal{Q} = I$。这通过应用高阶奇异值分解（HOSVD）或更廉价的随机近似实现。
3. 更新规则 – 优化器将原始梯度（或动量）投影到正交归一基上，得到正交化梯度 $\tilde{\mathcal{G}} = \mathcal{Q}\mathcal{Q}^\top \mathcal{G}$。随后使用标准的类似 Adam 的步长。
4. 近似策略 – 为了控制成本，作者尝试了：
   • 随机幂迭代（少量迭代）用于张量的每个模式。
   • 低秩截断（仅保留前‑k 个奇异分量）。
   • 层级回退（当张量成本超过阈值时回退到 Muon）。

该算法仅带来适度的开销（≈ 5–10 % 额外计算），同时保持了典型 Adam‑style 优化器的内存占用。

结果与发现

• 持续提升： TEON 在所有模型规模上相较于 Muon 将困惑度提升约 0.5–0.8 点，较 Adam 提升约 1.0–1.5 点。
• 可扩展性： 随着模型规模增大，收益也随之提升，表明层间梯度相关性在更大的网络中更为显著。
• 鲁棒性： 使用不同的 SVD 近似（秩 k = 5、10、20）进行实验，性能损失几乎可以忽略，验证了廉价近似已足够。

实际意义

• 更快的 LLM 预训练收敛速度： 开发者可以在更少的训练步数下实现相同或更好的模型质量，从而在 GPU/TPU 集群上节省成本。
• 即插即用的替代方案： TEON 的 API 与 Adam/Muon 完全对应，因而只需在现有的 PyTorch 或 JAX 流程中添加几行代码即可集成。
• 低精度环境下更好的稳定性： 正交化步骤能够缓解梯度爆炸/消失，使混合精度（FP16/BF16）训练更加可靠。
• 下游微调的潜在优势： 由于 TEON 能够产生更好的初始化权重空间，下游在特定领域数据上的微调可能收敛更快并达到更高的准确率。

限制与未来工作

• 计算开销: 虽然适度，但额外的 5–10 % 计算仍可能在超大规模运行（数百亿参数）中显著。
• 梯度张量的内存占用: 将所有层的梯度堆叠可能会在非常深的模型上导致内存紧张；作者建议使用流式或块状正交化作为解决方案。
• 理论分析局限于凸代理: 收敛性证明假设局部凸近似；将保证扩展到 Transformer 的完整非凸空间仍是未解之题。

未来方向

• 探索自适应秩选择以进一步降低开销。
• 将 TEON 与其他二阶技巧（例如 K‑FAC）结合，以实现更快的收敛。
• 将张量正交化应用于其他领域，如视觉 Transformer 或扩散模型。

作者

• Ruijie Zhang
• Yequan Zhao
• Ziyue Liu
• Zhengyang Wang
• Dongyang Li
• Yupeng Su
• Sijia Liu
• Zheng Zhang

论文信息

• arXiv ID: 2601.23261v1
• 分类: cs.LG, cs.AI
• 出版日期: 2026年1月30日
• PDF: 下载 PDF