Tensorion: A Tensor-Aware Generalization of the Muon Optimizer

Modern deep learning models, such as convolutional neural networks and transformers, rely on parameters that are inherently tensor-valued. Standard first-order optimizers like Adam treat these parameters as flattened vectors, ignoring the multilinear structure that defines these operations. This paper asks how to extend the success of the Muon optimizer—which exploits matrix structure via spectral norm constraints—to higher-order tensors without incurring the prohibitive computational costs of exact tensor spectral norm optimization.

The authors introduce Tensorion, an optimizer built around a Linear Minimization Oracle (LMO) over a tensor norm ball. The core challenge is that the exact tensor spectral norm is NP-hard to compute. To address this, the authors define a relaxed tensor spectral norm using the dual of the maximum nuclear norm across a set of tensor unfoldings. This construction provides a tight upper bound on the tensor spectral norm while ensuring the LMO remains computationally tractable. The authors also derive a heuristic for selecting the optimal unfolding strategy, which balances the row and column dimensions of the unfolding matrices to maximize the effective number of nonzero singular values.

Tensorion recovers the original Muon optimizer when applied to order-2 tensors (matrices). For higher-order tensors, the authors demonstrate that their dual-norm relaxation is theoretically sound and provides a more stable optimization trajectory than standard vector-based methods. Experiments on convolutional and transformer-based architectures for computer vision tasks show that Tensorion consistently outperforms conventional SGD and Adam-based baselines, offering improved convergence behavior and more stable gradient updates.