🤖 machine learning

Accelerating LMO-Based Optimization via Implicit Gradient Transport

This paper proposes LMO-IGT, a new class of stochastic optimization methods that leverages implicit gradient transport to achieve an improved $\mathcal{O}(\varepsilon^{-3.5})$ iteration complexity with only a single gradient evaluation per iteration, while introducing a unified framework and the regularized support function to bridge theoretical gaps between unconstrained and constrained LMO-based approaches.

Original authors: Won-Jun Jang, Si-Hyeon Lee

Published 2026-05-08

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Won-Jun Jang, Si-Hyeon Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the lowest point in a vast, foggy valley (the "loss landscape") to train a massive AI model. You can't see the whole valley, so you have to take steps based on the slope right under your feet. This is what optimization algorithms do.

For a long time, the standard way to do this was like taking a step in the direction the ground is sloping down, but adjusting your step size based on how steep it is. Recently, some new methods (like Lion and Muon) changed the game. Instead of just looking at the slope, they look at the average slope over time (momentum) and then "normalize" it. Think of it like a hiker who doesn't just walk downhill, but constantly checks their compass to ensure they are walking in the most efficient direction, regardless of how steep the hill is.

However, these new methods still have a problem: they can be a bit "laggy." Because they rely on the average of past steps, they sometimes react too slowly to sudden changes in the terrain.

The Problem: The Laggy Compass

The paper identifies that these "LMO-based" methods (Linear Minimization Oracle) are great, but they suffer from a delay. Imagine you are driving a car with a very heavy steering wheel. You turn the wheel, but the car takes a moment to actually change direction. In math terms, the "momentum" (the car's current direction) is based on old data, so it doesn't perfectly match where you actually need to go right now.

To fix this lag, previous researchers tried a technique called Variance Reduction. This is like sending a scout ahead to check the road, then coming back to tell you the direction. It works faster, but it's expensive: you have to send the scout out twice for every step you take (calculating gradients twice), which slows down the whole process and uses more computer power.

The Solution: The "Lookahead" Trick (IGT)

The authors propose a new method called LMO-IGT (Implicit Gradient Transport). They wanted to get the speed boost of the "scout" without the cost of sending two scouts.

Here is the creative analogy:
Imagine you are walking a dog on a leash.

Standard Method: You look at where the dog is right now, guess where it will go, and pull the leash. But the dog is already moving, so you are always a split-second behind.
Variance Reduction (Old Fix): You stop, run ahead to where the dog might be, check the terrain, run back, and then pull the leash. Accurate, but exhausting (two trips).
LMO-IGT (The New Fix): You don't stop or run ahead. Instead, you imagine a "ghost" version of yourself walking slightly ahead of you on the same path. You ask the ghost, "What does the ground feel like there?" and use that information to pull the leash. You only take one step, but you are using information from a spot slightly ahead of you.

This "ghost" is the transported point. By calculating the slope at this slightly advanced point, the algorithm corrects its momentum before it actually makes the mistake of following old data. It's like having a crystal ball that only shows you the next few inches of the path, allowing you to steer perfectly without extra effort.

The Unified Framework

The paper also builds a "universal translator" for these methods.

Some methods work best on open fields (unconstrained).
Some work best inside walled gardens (constrained).
Previously, scientists used different rulebooks to measure success for each.

The authors created a new measuring stick called the Regularized Support Function (RSF). Think of this as a universal ruler that can measure how close you are to the bottom of the valley, whether you are in an open field or a walled garden. This allows them to compare all these different methods fairly on a single scale.

The Results

Using this new "Lookahead" trick (IGT), the authors found:

Speed: Their new method converges (finds the bottom) faster than the standard methods.
Efficiency: Unlike the "scout" method (Variance Reduction), it doesn't require extra calculations. It keeps the "one step, one calculation" rule, so it runs just as fast as the standard methods but gets better results.
Performance: When they tested this on image recognition (CIFAR-10) and language models (writing text), their new version, called Muon-IGT, consistently beat the others. It reached higher accuracy in the same amount of time.

In Summary

The paper introduces a smarter way to navigate the complex terrain of AI training. Instead of getting stuck reacting to old information (lag) or paying a heavy price to check the road ahead (variance reduction), they use a clever "lookahead" trick to steer more accurately with the same amount of effort. This makes training large AI models faster and more efficient without needing more computer power.

Technical Summary: Accelerating LMO-Based Optimization via Implicit Gradient Transport

Problem Statement

Recent optimizers such as Lion and Muon have demonstrated strong empirical performance by normalizing gradient momentum via Linear Minimization Oracles (LMOs). While these methods are effective, their theoretical understanding remains fragmented between unconstrained and constrained formulations, often relying on disparate stationarity measures (gradient norms vs. Frank–Wolfe gaps). Furthermore, while variance reduction (VR) techniques have been explored to accelerate LMO-based methods, they typically incur substantial computational overhead by requiring additional gradient evaluations per iteration (e.g., computing gradients at both the current and previous iterates). This creates a trade-off: achieving faster theoretical convergence rates often degrades practical efficiency in large-scale settings due to increased wall-clock time and memory usage.

The paper addresses two central questions:

Can LMO-based methods achieve accelerated convergence while avoiding the substantial computational overhead associated with variance reduction?
Can stochastic LMO-based optimization be understood within a unified framework that bridges unconstrained and constrained settings?

Methodology

1. Implicit Gradient Transport (IGT) for LMOs

The authors propose LMO-IGT, a new class of stochastic LMO-based methods that extends the concept of Implicit Gradient Transport (IGT) to general LMO updates.

Core Mechanism: Unlike standard stochastic LMO, which evaluates gradients at the current iterate $w_t$ , LMO-IGT evaluates stochastic gradients at a "transported" point $x_t$ . This point serves as a lookahead approximation of the optimization trajectory.
Update Rule: The method maintains two momentum buffers with different time scales. The gradient momentum $g_t$ is updated using the gradient evaluated at $x_t$ , while the iterate $w_t$ is updated via the LMO direction derived from $g_t$ . The transported point $x_{t+1}$ is updated using the same LMO direction but potentially different step sizes.
Efficiency: Crucially, this mechanism preserves the single-gradient-per-iteration structure of standard stochastic LMO. It does not require the extra gradient evaluation at a previous iterate that variance-reduced methods (like Lion-VR or Muon-VR) necessitate.

2. Unified Framework and Regularized Support Function (RSF)

To analyze these methods uniformly, the authors introduce a unified framework encompassing three classes:

Stochastic LMO: Standard methods (e.g., Lion, Muon) with one gradient evaluation.
Variance-Reduced LMO (LMO-VR): Methods incorporating STORM-like correction terms with two gradient evaluations.
LMO-IGT: The proposed method with one gradient evaluation but using transported points.

A key theoretical contribution is the introduction of the Regularized Support Function (RSF), denoted as $\Psi_{C,\lambda}(w)$ .

Definition: $\Psi_{C,\lambda}(w) := \sup_{v \in C} \langle -\nabla F(w), v - \lambda w \rangle$ .
Unification:
- When $\lambda > 0$ (constrained setting), RSF scales to the Frank–Wolfe gap.
- When $\lambda = 0$ (unconstrained setting), RSF reduces to the support function, which coincides with the dual norm (e.g., gradient norm for Euclidean balls).
This measure allows for a common convergence analysis across both constrained and unconstrained formulations, bridging the gap between gradient-norm and Frank–Wolfe-gap notions of stationarity.

Key Contributions

LMO-IGT Algorithm: A new class of optimizers that accelerates LMO-based methods by evaluating gradients at transported points, achieving faster convergence without additional gradient evaluations.
Unified Theoretical Framework: A generalized update algorithm and the RSF stationarity measure that unify the analysis of Stochastic LMO, LMO-VR, and LMO-IGT.
Convergence Guarantees: The paper establishes the following iteration complexities to reach $\epsilon$ $ϵ$ -stationarity (measured by RSF):
- Stochastic LMO: $O(\epsilon^{-4})$ (Standard rate).
- LMO-VR: $O(\epsilon^{-3})$ (Requires 2 gradient evaluations per iteration).
- LMO-IGT: $O(\epsilon^{-3.5})$ (Requires only 1 gradient evaluation per iteration).
Practical Instantiations: The authors instantiate LMO-IGT with specific LMO sets, yielding Lion-IGT (using $\ell_\infty$ -norm balls) and Muon-IGT (using operator-norm balls).

Experimental Results

The authors evaluated the proposed methods on image classification (CIFAR-10 with ResNet-18) and language modeling (Shakespeare dataset with nanoGPT, and OpenWebText with a 124M parameter model).

Performance: Muon-IGT consistently achieved the strongest overall performance across evaluated settings, outperforming standard Lion/Muon, their variance-reduced counterparts (Lion-VR/Muon-VR), and baselines like AdamW and NIGT.
Efficiency: While variance-reduced methods (Lion-VR, Muon-VR) theoretically offer $O(\epsilon^{-3})$ rates, they incurred significantly higher wall-clock time due to the extra gradient evaluation. In contrast, LMO-IGT methods ran at nearly the same speed as their plain counterparts (single gradient eval) while achieving higher accuracy.
Robustness: Ablation studies and hyperparameter sensitivity tests indicated that Muon-IGT is robust to choices of learning rate and weight decay, performing well across a broad search space.

Significance and Claims

The paper claims that IGT provides an effective and practical acceleration mechanism for modern LMO-based optimization.

Bridging Theory and Practice: The work successfully bridges the gap between standard stochastic LMO and variance-reduced methods. It demonstrates that one can achieve convergence rates closer to the optimal $O(\epsilon^{-3})$ (specifically $O(\epsilon^{-3.5})$ ) without the prohibitive computational cost of variance reduction.
Unified Perspective: By introducing the RSF, the paper provides a common theoretical language to compare optimizers that were previously analyzed under incompatible criteria (gradient norms vs. Frank–Wolfe gaps).
Practical Impact: The empirical results suggest that for large-scale settings where computational efficiency is paramount, LMO-IGT offers a superior trade-off, delivering state-of-the-art accuracy with negligible overhead compared to standard LMO methods.

The authors note that their theoretical guarantees rely on standard smoothness and second-order regularity assumptions, which are idealized relative to practical deep learning, but the framework offers a principled basis for rate comparison. They do not claim to solve all optimization challenges but highlight that transported gradients extend beyond Euclidean normalization to general LMO-based updates.