想优化张量参数?这篇论文把Muon优雅地推广到高阶张量,实验比Adam更稳健。
Muon优化器通过谱范数约束执行最速下降,但仅适用于矩阵。Tensorion将这一方法扩展到高阶张量,基于线性最小化预言机(LMO)在张量范数球上进行优化。其LMO通过自适应选择展开矩阵可高效计算,且当限制为二阶张量时精确恢复Muon。在张量计算机视觉任务中,Tensorion相比Adam和现有张量感知基线展现出更优的收敛行为与更稳定的梯度更新。
Tensorion: A Tensor-Aware Generalization of the Muon Optimizer
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.