Accelerating LMO-Based Optimization via Implicit Gradient Transport

TL;DR

This paper introduces LMO-IGT, a stochastic LMO-based optimization method using implicit gradient transport, which accelerates convergence with minimal additional computational cost.

Contribution

The paper proposes a unified framework for stochastic LMO-based optimization with implicit gradient transport, achieving improved convergence rates and practical acceleration.

Findings

LMO-IGT achieves an iteration complexity of O(ε^{-3.5}) with a single stochastic gradient per iteration.

Empirical results show LMO-IGT outperforms stochastic LMO with negligible overhead.

Muon-IGT demonstrates strong overall performance across various settings.

Abstract

Recent optimizers such as Lion and Muon have demonstrated strong empirical performance by normalizing gradient momentum via linear minimization oracles (LMOs). While variance reduction has been explored to accelerate LMO-based methods, it typically incurs substantial computational overhead due to additional gradient evaluations. At the same time, the theoretical understanding of LMO-based methods remains fragmented across unconstrained and constrained formulations. Motivated by these limitations, we propose \emph{LMO-IGT}, a new class of stochastic LMO-based methods leveraging implicit gradient transport (IGT). We further introduce a unified framework for stochastic LMO-based optimization together with a new stationarity measure, the \emph{regularized support function} (RSF), which bridges gradient-norm and Frank--Wolfe-gap notions within a common framework. By evaluating stochastic gradients at transported points, LMO-IGT accelerates convergence while retaining the single-gradient-per-iteration structure of standard stochastic LMO. Our analysis establishes that stochastic LMO achieves an iteration complexity of $\mathcal{O}(\varepsilon^{-4})$, variance-reduced LMO achieves $\mathcal{O}(\varepsilon^{-3})$ at the cost of additional gradient evaluations, and LMO-IGT achieves $\mathcal{O}(\varepsilon^{-3.5})$ using only a single stochastic gradient per iteration. Empirically, LMO-IGT consistently improves over stochastic LMO counterparts with negligible overhead. Among its instantiations, Muon-IGT achieves the strongest overall performance across evaluated settings, demonstrating that IGT provides an effective and practical acceleration mechanism for modern LMO-based optimization.