Lions and Muons: Optimization via Stochastic Frank-Wolfe
Maria-Eleni Sfyraki, Jun-Kun Wang
TL;DR
The paper addresses optimization in deep learning by unifying Lion and Muon within a Stochastic Frank-Wolfe framework, establishing convergence via the FW gap $\mathcal{G}(\mathbf{x})$ and showing that gap convergence implies convergence to a KKT point under the associated norm constraints. It extends the FW analysis to heavy-tailed stochastic gradients by introducing clipping-based variants and variance-reduction enhancements, achieving high-probability rates of $O\left( \log(T/\delta) / T^{(p-1)/(3p-2)} \right)$ or $O\left( \log(T/\delta) / T^{(2p-1)^{-1}(p-1)} \right)$ depending on the regime, with $p\in(1,2]$. The work also demonstrates practical gains via experiments on nanoGPT (Shakespeare) and CIFAR-10 (ResNet-18), where Lion+ and Muon+ improve convergence speed and test performance over their baselines. Overall, this unifying Stochastic FW perspective broadens projection-free optimization to nonconvex deep-learning tasks with heavy-tailed noise, enabling robust, efficient training under general convex constraints.
Abstract
Stochastic Frank-Wolfe is a classical optimization method for solving constrained optimization problems. On the other hand, recent optimizers such as Lion and Muon have gained quite significant popularity in deep learning. In this work, building on recent initiatives, we provide a unifying perspective by interpreting these seemingly disparate methods through the lens of Stochastic Frank-Wolfe. Specifically, we show that Lion and Muon with weight decay can be viewed as special instances of a Stochastic Frank-Wolfe, and we establish their convergence guarantees in terms of the Frank-Wolfe gap, a standard stationarity measure in non-convex optimization for Frank-Wolfe methods. We further find that convergence to this gap implies convergence to a KKT point of the original problem under a norm constraint for Lion and Muon. Moreover, motivated by recent empirical findings that stochastic gradients in modern machine learning tasks often exhibit heavy-tailed distributions, we extend Stochastic Frank-Wolfe to settings with heavy-tailed noise by developing two robust variants with strong theoretical guarantees that hold for general compact convex sets without the need for a large batch size, filling the gap in the literature on Stochastic Frank-Wolfe for non-convex optimization. Our contributions in the later part of this work, in turn, yield new variants of Lion and Muon, that better accommodate heavy-tailed gradient noise, thereby enhancing their practical scope.