Tensorion: A Tensor-Aware Generalization of the Muon Optimizer

2026-06-24

AI Summary

This work addresses the limitation of mainstream first-order optimizers, which disregard the multilinear tensor structure inherent in model parameters, thereby constraining optimization efficiency. To overcome this, we propose Tensorion, the first optimizer that generalizes the matrix-based Muon algorithm to higher-order tensors. Tensorion achieves structure-aware optimization by constructing a computationally tractable linear minimization oracle (LMO) over the tensor spectral norm ball. The key innovation lies in extending steepest descent under spectral norm constraints to high-order tensors and employing an adaptive tensor unfolding strategy that tightly approximates the spectral norm; this formulation exactly recovers Muon in the second-order case. Empirical results demonstrate that Tensorion exhibits superior convergence and more stable gradient updates compared to Adam and existing tensor-aware optimizers on computer vision tasks with intrinsic tensor structures.

Abstract

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.

Problem

• tensor structure
• optimization
• first-order optimizer
• parameter structure
• machine learning models

Innovation

• tensor-aware optimization
• linear minimization oracle
• spectral norm constraint
• higher-order tensors
• steepest descent