让优化器理解张量结构,提升模型训练稳定性和收敛速度。
Tensorion: A Tensor-Aware Generalization of the Muon Optimizer
- 基于张量范数球的线性最小化查询,显式利用权重的多线性结构。
- 在视觉任务中收敛更快,梯度更新更稳定,优于Adam和现有张量优化器。
- 可无缝退化为矩阵优化器Muon,适用于高阶张量建模场景。
常见的一阶优化器(如Adam)将每个参数块视为无结构向量,忽略了现代模型中普遍存在的多线性权重结构。近期研究表明,利用矩阵结构可改善优化动态。以Muon为例,其在谱范数约束下执行最速下降。本文进一步提出Tensorion,一种面向高阶张量的优化器,将穆昂(Muon)的约束优化视角从矩阵扩展至张量。Tensorion围绕张量范数球上的线性最小化查询(LMO)构建,所选范数兼顾紧密约束张量谱范数与保持计算可行性。该LMO可计算,因其可简化为自适应选择的展开矩阵上的操作。值得注意的是,当限制为二阶张量(即矩阵)时,Tensorion精确恢复Muon。在基于张量的计算机视觉任务实验中,结果表明,相比基于Adam及现有张量感知基线,Tensorion展现出更优的收敛行为与更稳定的梯度更新。
原文摘要 · Abstract (English)
Common first-order optimizers, such as Adam, implicitly treat each parameter block as an unstructured vector, which disregards the multilinear weight structure present in many modern machine learning models. Recent work has shown that exploiting matrix structure can improve optimization dynamics. A notable example is Muon, which performs steepest descent under the spectral norm constraint. We take the next step and introduce Tensorion, a tensor-aware optimizer that extends Muon's constrained optimization perspective from matrices to higher-order tensors. Tensorion is built around a linear minimization oracle (LMO) over a tensor norm ball. The norm is carefully chosen to balance two objectives: tightly bounding the tensor spectral norm, while still keeping the LMO tractable. This LMO becomes computable because it reduces to operations on adaptively selected unfolding matrices. Notably, when restricted to order-2 tensors (i.e., matrices), Tensorion recovers Muon exactly. Experiments on tensor-based computer vision problems suggest that Tensorion can offer improved convergence behavior and more stable gradient updates compared with Adam-based and existing tensor-aware baselines in the evaluated settings.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。