Tensorion: A Tensor-Aware Generalization of the Muon Optimizer
Vladimir Bogachev, Vladimir Aletov, Alexander Molozhavenko, Sergei Kudriashov, Maxim Rakhuba
Read on arXiv →
Key claim
Tensorion improves optimization for tensor-based models significantly.
In plain English
Imagine you're training a complex machine learning model, like one that processes images or videos. You want it to learn quickly and accurately, but traditional optimizers like Adam treat all the model's parameters as if they were just a long list. This can lead to inefficiencies, especially when the model has a more intricate structure, like tensors, which are multi-dimensional arrays. When the optimizer doesn't account for this structure, it can struggle with convergence, meaning it takes longer to find the best solution or gets stuck in suboptimal states. This is what's called inefficient optimization dynamics.
To address these issues, Tensorion was developed. It’s an optimizer that recognizes the multi-dimensional nature of tensors and uses this understanding to improve the training process. By focusing on the tensor structure, Tensorion can optimize more effectively, leading to faster convergence and more stable updates. It does this by employing a linear minimization oracle that operates within a carefully defined tensor norm ball, ensuring that the optimization process remains tractable while still being effective.
Compared to existing methods, Tensorion shows promising results, particularly in tasks that involve tensor-based data, like computer vision. It not only matches the performance of previous optimizers but also offers improvements in how quickly and reliably models can be trained. For builders, this means that if you're working with complex models that utilize tensor structures, Tensorion could be a significant upgrade over traditional optimizers, helping you achieve better results in less time.
Tensorion introduces a new tensor-aware optimization method that extends existing techniques.
The experiments show improved convergence and stability compared to established optimizers, though more diverse evaluations could strengthen claims.
Deep reliability assessment
The paper’s methodology supports that a tensor-aware LMO with adaptively chosen unfoldings can be implemented efficiently enough to train CNNs on CIFAR-10 and Tiny ImageNet, recovering Muon when d=2 and showing better early-epoch accuracy than SGD/AdamW and a Muon-as-τ={1} baseline in the tested settings. Claims about tight spectral control, broader stability, and scalability beyond small-to-mid CV benchmarks are not yet substantiated by large-scale or non-vision experiments, and the practical tightness of the proposed bound on real learned tensors is not quantitatively validated.
Reproducibility
Code: yes (GitHub link provided). Datasets: standard (CIFAR-10, Tiny ImageNet). Details: seeds and mean±std reported; unfolding selection heuristic specified; some configs and additional experiments in Supplement; ablation used exact SVDs (vs approximate in practice), and large-scale settings are not covered.
Key figure
Figure 1 shows that for ResNet-18 on CIFAR-10 (25 epochs), different tensor unfoldings split into low- and high-accuracy regimes, with the online τ selection best but costly, and the offline τ_opt competitive and much cheaper, outperforming SGD-M and AdamW under the same constant-LR setup.