Search papers, labs, and topics across Lattice.
This paper introduces Tensorion, a tensor-aware optimizer that generalizes the Muon optimizer by extending its constrained optimization framework from matrices to higher-order tensors. By leveraging a linear minimization oracle over a carefully chosen tensor norm ball, Tensorion achieves improved convergence and stability in gradient updates for tensor-based computer vision tasks. Experimental results demonstrate that Tensorion outperforms both Adam-based optimizers and existing tensor-aware methods, highlighting its effectiveness in optimizing complex models with multilinear structures.
Tensorion achieves superior convergence and stability in optimizing tensor-based models, outperforming traditional methods like Adam.
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.
