Tensorion: A Tensor-Aware Generalization of the Muon Optimizer

会议: NeurIPS 2026
arXiv: 2606.25975
OpenReview: https://openreview.net/forum?id=5PDUovNqbT
代码: https://github.com/MTML-LAB/Tensorion
领域: 其他
关键词: 优化算法、张量优化、Muon推广、线性极小化神谕、谱范数松弛

一句话总结

本文提出了一种感知张量多重线性结构的优化器 Tensorion，通过将 Muon 的谱范数约束优化视角从矩阵提升至高阶张量，基于对偶核范数构造出紧致且可解的松弛线性极小化神谕（LMO），在卷积核与高阶张量参数上实现了超越 SGD 和 AdamW 的收敛性与泛化表现。

研究背景与动机

以 Adam 为代表的传统一阶深度学习优化器，通常将模型内部的每组参数张量简单展平成无结构的一维向量。这种扁平化处理完全忽略了现代神经网络中天然存在的多线性（multilinear）权重几何结构。近期针对矩阵结构优化的研究表明，显式利用参数矩阵的二维代数属性能够显著改善优化动力学；其中具有代表性的 Muon 优化器在谱范数约束球下执行最陡下降（Steepest Descent），在多类模型架构上取得了极具竞争力的效率与稳定性。

然而，卷积神经网络中的高维卷积核、Transformer 中的多头注意力机制（MHA）以及混合专家网络（MoE）的参数在数学上本质都是高阶张量（order \(d \geq 3\)）。若将高阶张量粗暴展平成二维矩阵或一维向量，不仅强行混合了独立的张量模态，还割裂了多线性算子的结构信息。更棘手的是，若直接将 Muon 推广到高阶张量，面临着基础数学理论上的根本瓶颈：计算高阶张量的谱范数本身是 NP-hard 的，求解对应的张量谱范数球线性极小化神谕（Linear Minimization Oracle, LMO）在实际深度网络训练中在计算上完全不可行；而常见的单模态展开近似又存在严重破坏对称性或在退化模态下直接塌缩至 Frobenius 范数的病态缺陷。

为了打破高阶张量谱优化的计算壁垒，本文并未在原空间中对展开谱范数进行保守的经验聚合，而是将视角转换至凸共轭对偶空间，基于各模式展开矩阵的核范数最大值建立紧致下界，再通过对偶映射定义出一种可计算且紧贴张量谱范数的松弛范数。核心 idea：将 Muon 的矩阵谱范数约束 LMO 拓展至对偶构造的张量松弛范数球上，使得高阶张量的最陡下降能够严格等价于自适应选择核范数最大展开矩阵的奇异值正交化，既保留了多线性几何结构，又将计算复杂度完全约束在可高效执行的 Newton-Schulz 迭代内。

方法详解

整体框架

Tensorion 的核心计算流由静态形状决策与迭代更新两个阶段组成。在优化启动前，算法根据各层张量参数的具体维度分布，通过离线平衡启发式准则预先判定最优展开模态划分 \(\tau^{\mathrm{opt}}\)；在每个训练迭代步中，梯度与一阶动量保持高阶张量结构，经由所选展开操作转换为二维矩阵后送入 Newton-Schulz 正交化例程求解 LMO，最后通过逆展开（fold）恢复张量形状并按缩放后的自适应步长更新权重。

对于卷积层等典型 4 阶张量，其输入为动量张量 \(M_k \in \mathbb{R}^{n_1 \times \dots \times n_d}\)，输出为沿对偶张量范数球边界更新后的权重 \(W_{k+1}\)。整个流程消除了逐迭代步枚举所有模态展开的昂贵开销，完全在等效矩阵流形上并行计算。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["权重与梯度张量<br/>W, G ∈ ℝ^{n₁×…×n_d}"] --> B["动量累积更新<br/>M_k = G_k + β M_{k-1}"]
    B --> C["对偶核范数展开机制<br/>unfold(M_k, τ^opt)"]
    C --> D["Newton-Schulz 正交化求解 LMO<br/>NS(展开矩阵)"]
    D --> E["张量形状逆折叠<br/>X_k = fold(·, τ^opt)"]
    E --> F["维数自适应步长缩放与更新<br/>W_{k+1} = (1-γη_k)W_k - η̂_k X_k"]

关键设计

1. 对偶核范数松弛与张量谱范数下界界定：破解张量 LMO 的 NP-hard 困境

直接求解张量谱范数 \(\|X\|_\sigma = \sup_{\|u^{(k)}\|_2=1} |\langle X, u^{(1)} \otimes \dots \otimes u^{(d)} \rangle|\) 的 LMO 是 NP-hard 问题。直接在原空间对多种展开谱范数求最大值 \(\max_{\tau} \|X_{(\tau)}\|_2\) 会在包含单元素模态（如 \(1 \times m \times n\)）时直接退化为平庸的 Frobenius 范数 \(\|X\|_F\)；求最小值则不满足三角不等式。Tensorion 的切入点是在对偶空间定义松弛核范数，令指标集 \(\mathcal{T}\) 包含所有合法的模态划分，定义对偶范数为：

\[ \|M\|_\Sigma^\dagger := \max_{\tau \in \mathcal{T}} \|M_{(\tau)}\|_*, \quad \|X\|_\Sigma := (\|X\|_\Sigma^\dagger)^\dagger. \]

该构造从数学上证明了严格夹逼关系：\(\|X\|_\sigma \leq \|X\|_\Sigma \leq \min_{\tau \in \mathcal{T}} \|X_{(\tau)}\|_2\)。由此，针对松弛张量范数球的 LMO 求解问题 \(X^{\mathrm{opt}} \in \mathrm{Argmax}_{\|X\|_\Sigma \leq 1} \langle M, X \rangle\)，其解集严格对应于对偶范数在 \(M\) 处的次微分：

\[ X^{\mathrm{opt}} \in \partial \|M\|_\Sigma^\dagger = \mathrm{conv}_{\tau \in I} \left( \mathrm{fold}_\tau \left( \partial \|M_{(\tau)}\|_* \right) \right), \quad I = \left\{ \tau \;\middle|\; \|M_{(\tau)}\|_* = \|M\|_\Sigma^\dagger \right\}. \]

这使得张量优化方向的寻找精确简化为：挑选出具有最大矩阵核范数的展开矩阵，对其执行紧凑 SVD 分解 \(M_{(\tau)} = U \Sigma V^\top\) 后取 \(U V^\top\)，再将其 fold 回张量结构。当张量阶数 \(d=2\) 时，该构造精确退化回标准 Muon 优化器。

2. 维度平衡离线展开启发式选择：消除逐步在线搜索的计算冗余

虽然在线策略能在每个 step 动态评估所有展开的核范数并取最大者，但这需要并发计算多组奇异值分解或 Newton-Schulz 迭代，耗费大量算力。作者证明了在元素服从零均值次指数独立同分布假定下，张量展开矩阵 \(M_{(\tau)} \in \mathbb{R}^{m_\tau \times n_\tau}\)（设 \(m_\tau \leq n_\tau\) 且 \(m_\tau n_\tau = \prod n_i\)）的核范数高概率上界由行数主导并随维度平衡性单调递增：

\[ \|M_{(\tau)}\|_* \leq C \Bigl(m_\tau \sqrt{n_\tau} + \sqrt{m_\tau} \Bigl(\sqrt{\log\tfrac{2}{\delta}} + \log\tfrac{2}{\delta}\Bigr)\Bigr). \]

对于固定的总元素量，该理论上界在矩阵形状尽可能接近正方形（即最平衡展开）时达到最大。基于此，Tensorion 提出了脱机的确定性静态展开策略：

\[ \tau^{\mathrm{opt}} := \operatorname*{arg\,max}_{\tau \in \mathcal{T}} \left( \min \left\{ \prod_{i \in \tau} n_i, \prod_{i \notin \tau} n_i \right\} \right). \]

在训练开始前即可为各网络层一次性锁定展开指标 \(\tau^{\mathrm{opt}}\)。实测表明离线策略与昂贵的在线动态搜索在 LMO 目标值上的相对相对偏差（relative nuclear norm gap）均值仅为 0.39%，计算耗时与原版矩阵 Muon 完全处于同一量级。

3. 卷积算子雅可比的结构化保持与内近似：理论弥合卷积核与线性映射

对于标准卷积层，卷积操作实质是高维稀疏块循环 Toeplitz 算子 \(T_K \in \mathbb{R}^{c_{\mathrm{out}} n^2 \times c_{\mathrm{in}} n^2}\) 的线性映射 \(\mathrm{vec}(Y) = T_K \mathrm{vec}(X)\)。若直接对全尺寸卷积雅可比矩阵 \(T_K\) 施加 Muon 的谱约束，矩阵维度极其庞大且在反向传播中难以维系 Toeplitz 稀疏结构。论文利用卷积算子谱范数界定定理 \(\|K\|_\sigma \leq \|T_K\|_2 \leq \sqrt{hw} \|K\|_\sigma\)，导出了可行域球体的包含链：

\[ \left\{ H : \|H\|_\Sigma \leq \frac{1}{\sqrt{hw}} \right\} \subseteq \mathbb{B}_\sigma\left(\frac{1}{\sqrt{hw}}\right) \subseteq \mathbb{B}_T(1) \subseteq \mathbb{B}_\sigma(1). \]

这一推导证明了在缩放 \(1/\sqrt{hw}\) 的张量松弛谱范数球上对卷积核直接求解 LMO，恰好构成了对结构化全雅可比 Muon 问题的严谨内近似（Inner Approximation），从而为将 4 阶卷积核视为多重线性张量直接更新提供了坚实的算子理论支撑。

损失函数 / 训练策略

Tensorion 的核心更新律集成了解耦权重衰减（Weight Decay）与维度缩放步长调节。设第 \(k\) 步基础标称学习率为 \(\eta_k\)，权重衰减系数为 \(\gamma\)，动量衰减系数为 \(\beta\)。有效步长依据展开矩阵的最大维度自适应缩放：

\[ \hat{\eta}_k = \eta_k \cdot 0.2 \sqrt{\max(m_{\tau^{\mathrm{opt}}}, n_{\tau^{\mathrm{opt}}})}. \]

参数迭代公式为：

\[ W_{k+1} = (1 - \gamma \eta_k) W_k - \hat{\eta}_k X_k, \quad X_k = \mathrm{fold}\left(\mathrm{NS}\left(\mathrm{unfold}(M_k, \tau^{\mathrm{opt}})\right), \tau^{\mathrm{opt}}\right). \]

在混合训练实践中，Tensorion 仅作用于张量阶数 \(d \geq 3\) 的多维权重块（如卷积核、Patch Embedding 等）；对于 \(d \leq 2\) 的普通矩阵权重、偏置及标量归一化参数，则默认与标量/向量级优化器（如 SGD 或 AdamW）协同优化。

实验关键数据

主实验

论文在标准图像分类基准（CIFAR-10、CIFAR-100、Tiny ImageNet）上对 ResNet 系列进行了多组充分评测。以下汇总对比了使用 SGD 优化非张量参数时的主要性能表现（指标为均值 \(\pm\) 标准差）：

数据集	模型	优化器组合	测试损失 (Test Loss) ↓	测试准确率 (Test Acc) ↑	相对基线提升 (Gain)
CIFAR-10	ResNet-18	SGD-M	\(0.25 \pm 0.02\)	\(94.1 \pm 0.13\%\)	基准
CIFAR-10	ResNet-18	AdamW	\(0.44 \pm 0.02\)	\(93.7 \pm 0.30\%\)	-0.4%
CIFAR-10	ResNet-18	Tensorion + SGD	\(\mathbf{0.21 \pm 0.01}\)	\(\mathbf{95.3 \pm 0.10\%}\)	+1.2%
CIFAR-100	ResNet-34	SGD-M	\(\mathbf{0.89 \pm 0.13}\)	\(78.2 \pm 0.20\%\)	基准
CIFAR-100	ResNet-34	AdamW	\(1.33 \pm 0.10\)	\(74.2 \pm 0.10\%\)	-4.0%
CIFAR-100	ResNet-34	Tensorion + SGD	\(0.98 \pm 0.09\)	\(\mathbf{79.1 \pm 0.10\%}\)	+0.9%
Tiny ImageNet	ResNet-50	SGD-M	\(1.79 \pm 0.06\)	\(61.01 \pm 0.02\%\)	基准
Tiny ImageNet	ResNet-50	AdamW	\(\mathbf{1.33 \pm 0.05}\)	\(57.40 \pm 0.10\%\)	-3.61%
Tiny ImageNet	ResNet-50	Tensorion + SGD	\(1.65 \pm 0.06\)	\(\mathbf{62.70 \pm 0.10\%}\)	+1.69%

在 Vision Transformer（ViT）和 Swin Transformer 系列实验中，即使仅将 Tensorion 施加于最浅层的卷积 Patch Embedding 模块（阶数 \(d \geq 3\)），在 Imagenette 数据集微调中也观察到了显著增益，验证了多线性张量约束对输入表征层稳定性的关键作用。

消融实验

为剥离求解器近似误差并验证张量展开选取对优化轨迹的决定性影响，作者在 CIFAR-10 上使用 ResNet-18（固定 25 轮恒定学习率、精确 SVD 算子）对 4 阶卷积核的全部 7 种互补非对称展开方式进行了全面消融：

展开选取方案 (\(\tau\))	矩阵形状属性	测试准确率 (Test Acc) ↑	测试损失 (Test Loss) ↓	说明
\(\tau = \{1\}\)	等效经典 Muon 展开	\(90.67 \pm 0.32\%\)	\(0.424 \pm 0.026\)	偏置高精度群
\(\tau = \{2\}\)	沿输入通道展开	\(90.62 \pm 0.44\%\)	\(\mathbf{0.404 \pm 0.024}\)	偏置高精度群
\(\tau = \{3\}\)	沿空间高度展开（极度狭长）	\(87.40 \pm 1.00\%\)	\(0.485 \pm 0.051\)	表现较差群（失衡）
\(\tau = \{4\}\)	沿空间宽度展开（极度狭长）	\(88.05 \pm 0.69\%\)	\(0.423 \pm 0.028\)	表现较差群（失衡）
\(\tau = \{1, 2\}\)	通道-空间分离展开	\(88.83 \pm 0.39\%\)	\(0.430 \pm 0.022\)	表现较差群
\(\tau = \{1, 3\}\)	混合模态展开	\(90.30 \pm 0.51\%\)	\(0.426 \pm 0.031\)	偏置高精度群
\(\tau = \{2, 4\}\)	混合模态展开	\(90.77 \pm 0.45\%\)	\(0.432 \pm 0.038\)	偏置高精度群
\(\tau^{\mathrm{opt}}\) (离线启发式)	逐层形状最大平衡	\(\mathbf{90.44 \pm 0.36\%}\)	\(0.441 \pm 0.024\)	逼近最高群，无额外耗时
Online \(\tau\) (逐步动态搜索)	在线实时贪心核范数	\(\mathbf{91.10 \pm 0.39\%}\)	\(0.416 \pm 0.026\)	理论上限，计算代价高昂
SGD-M (基线)	无张量结构感知	\(86.10 \pm 1.82\%\)	\(0.527 \pm 0.095\)	泛化表现最差
AdamW (基线)	坐标方向独立缩放	\(86.04 \pm 1.05\%\)	\(0.507 \pm 0.041\)	恒定步长下较易震荡

关键发现

• 展开形状决定优化上限：张量展开方式直接决定了所诱导的优化几何。若展开为极度瘦长矩阵（如仅沿空间小卷积核尺寸 \(\tau=\{3\}\) 或 \(\{4\}\) 展开），非零奇异值上限受到微小维度的刚性限制，优化效果显著劣化，性能落后高达 3% 以上。
• 离线平衡策略极具性价比：提出的 \(\tau^{\mathrm{opt}}\) 离线启发式准则自动避开了所有劣质展开，落入最高性能区间，与昂贵的逐步在线搜索仅相差 0.66% 准确率，且实际训练耗时与普通 Muon 基本一致。
• 泛化过拟合抑制：在 Tiny ImageNet 训练深层 ResNet-50 时，AdamW 约在 50 epoch 后测试损失反弹出现严重过拟合；而 Tensorion + SGD 的测试损失始终维持平稳单调下降，体现出张量谱约束带来的内在正则化特性。

亮点与洞察

• 对偶松弛的化繁为简：直接近似张量谱范数极易陷入无解或退化困境，作者巧妙通过凸共轭将问题转换至对偶核范数极大值，成功把高维张量 LMO 映射为展开矩阵的 SVD/正交化问题。
• 矩阵正方形化直觉的严谨数学支撑：利用次指数向量的集中不等式（Hanson-Wright 不等式），严格论证了更均衡的展开形状能承载更高核范数上界，为工程上选择正方形展开提供了坚实的理论依据。
• 即插即用的模块协同性：针对卷积与 Transformer 混合架构，算法天然支持阶数拆分策略——仅对阶数 \(d \geq 3\) 的卷积核/嵌入层施加张量正交化，既控制了计算开销，又赋能了关键特征抽取层的表征能力。

局限与展望

• 实验尺度与领域受限：目前验证主要集中于计算机视觉领域的标准图像分类任务（CNN 与小/中规模 ViT），尚未在大规模视频生成、多模态时空预测或大语言模型全流程预训练中进行数十亿量级参数验证。
• 对非张量优化器具有依赖性：Tensorion 专门面向 \(d \geq 3\) 的张量块，其余权重仍需依赖基座优化器（SGD 或 AdamW），混合训练时的超参数耦合（如张量与非张量学习率配比）增加了调优空间。
• 复杂张量网络的进一步拓展：未来可进一步探索在显式低秩张量分解网络（如 Tensor-Train 或 Tucker 分解格式）中内生嵌入 Tensorion，直接在分解核心张量上开展谱约束梯度下降。

相关工作与启发

• vs Muon (Jordan et al., 2024): Muon 仅在二维参数矩阵上施加谱约束，处理卷积核时通常依赖生硬的手工展开规则。Tensorion 从理论上构建了张量谱范数的紧致松弛，并在 \(d=2\) 时完美兼容 Muon，为多维卷积核与张量参数提供了形式自洽的理论范式。
• vs Shampoo / SOAP (Gupta et al., 2018; Vyas et al., 2024): Shampoo 等张量二阶优化算法依赖维护多个模态的协方差预条件矩阵，存在较高显存占用与矩阵逆求幂开销。Tensorion 则属于一阶 LMO 视角，无需显式构造任何辅助协方差矩阵，开销显著更低。
• vs TEON (Zhang et al., 2026): TEON 尝试跨层堆叠多个参数矩阵构成外生张量进行联合正交化；Tensorion 则忠实于网络层自身的固有张量几何，仅作用于单层权重的内部多线性模态，更符合层内算子的局部几何性质。

评分

• 新颖性: ⭐⭐⭐⭐⭐ 成功攻克张量谱范数 LMO 不可解的经典难题，提出形式优雅且具备紧致保证的对偶松弛范数。
• 实验充分度: ⭐⭐⭐⭐☆ 覆盖了从理论合成展开消融到 ResNet、ViT、Swin 及拉普拉斯谱问题的多维验证，但缺乏百亿大模型预训练实测。
• 写作质量: ⭐⭐⭐⭐⭐ 数学推导严谨清晰，从原范数、对偶空间到卷积算子雅可比的理论脉络环环相扣。
• 价值: ⭐⭐⭐⭐☆ 为高阶张量与卷积核参数的高效几何优化提供了全新标准框架，具有极高的理论启发与算法迁移潜力。

相关论文

• [NeurIPS 2026] Move on Muon: A Hamiltonian Probability Gradient Flow Perspective of Muon Optimizer
• [NeurIPS 2026] Towards Understanding the Power and Limits of the Muon Optimizer: A River-Valley Perspective
• [NeurIPS 2026] Phases of Muon: When Muon Eclipses SignSGD
• [NeurIPS 2026] AMUSE: Anytime Muon with Stable Gradient Evaluation
• [ICML 2026] LiMuon: Light and Fast Muon Optimizer for Large Models