LiMuon: Light and Fast Muon Optimizer for Large Models (2509.14562v1)

Abstract: Large models recently are widely applied in artificial intelligence, so efficient training of large models has received widespread attention. More recently, a useful Muon optimizer is specifically designed for matrix-structured parameters of large models. Although some works have begun to studying Muon optimizer, the existing Muon and its variants still suffer from high sample complexity or high memory for large models. To fill this gap, we propose a light and fast Muon (LiMuon) optimizer for training large models, which builds on the momentum-based variance reduced technique and randomized Singular Value Decomposition (SVD). Our LiMuon optimizer has a lower memory than the current Muon and its variants. Moreover, we prove that our LiMuon has a lower sample complexity of $O(\epsilon^{-3})$ for finding an $\epsilon$-stationary solution of non-convex stochastic optimization under the smooth condition. Recently, the existing convergence analysis of Muon optimizer mainly relies on the strict Lipschitz smooth assumption, while some artificial intelligence tasks such as training LLMs do not satisfy this condition. We also proved that our LiMuon optimizer has a sample complexity of $O(\epsilon^{-3})$ under the generalized smooth condition. Numerical experimental results on training DistilGPT2 and ViT models verify efficiency of our LiMuon optimizer.

• The paper presents LiMuon, a memory-efficient optimizer that reduces sample complexity from O(ε⁻⁴) to O(ε⁻³) through momentum-based variance reduction and randomized SVD.
• It compresses momentum matrices using RSVD, lowering memory requirements from O(mn) to O((m+n)r̂) for large-scale models.
• Empirical evaluations on language and vision tasks demonstrate LiMuon’s superior performance and memory savings compared to AdamW, Muon, and other optimizers.

LiMuon: A Memory-Efficient Muon Optimizer for Large-Scale Models

Overview and Motivation

The LiMuon optimizer addresses the computational and memory bottlenecks inherent in training large-scale models with matrix-structured parameters. While adaptive gradient methods such as Adam and AdamW are widely used, recent advances in matrix-aware optimizers (e.g., Muon, SUMO) have demonstrated improved empirical performance by leveraging the structure of neural network parameters. However, existing Muon variants suffer from either high sample complexity ($O(\epsilon^{-4})$) or substantial memory overhead ($O(mn)$ for $m \times n$ matrices). LiMuon introduces a momentum-based variance reduction technique combined with randomized SVD (RSVD) to achieve both lower sample complexity ($O(\epsilon^{-3})$) and reduced memory requirements ($(m+n)\hat{r}$, with $\hat{r} \ll \min(m,n)$).

Algorithmic Design

LiMuon builds upon the Muon optimizer's orthogonalized momentum update, but innovates in two key areas:

1. Momentum-Based Variance Reduction: Inspired by STORM, LiMuon maintains a running estimate of the gradient using a recursive update, which reduces variance and improves convergence rates.
2. Low-Rank Compression via RSVD: Instead of storing the full momentum matrix $M_t \in \mathbb{R}^{m \times n}$, LiMuon compresses $M_t$ using RSVD, storing only the factors $\hat{U}_t \in \mathbb{R}^{m \times \hat{r}}$, $\hat{S}_t \in \mathbb{R}^{\hat{r} \times \hat{r}}$, and $\hat{V}_t \in \mathbb{R}^{n \times \hat{r}}$. This drastically reduces memory usage.

   Figure 1: Low Rank Compression of Momentum $M_t$.

The optimizer supports two options:

• Option #1: Standard momentum update, no compression.
• Option #2: RSVD-based compression, enabling memory savings.

The RSVD step is implemented as follows:

def rsvd(A, r_hat, s): l = r_hat + s Omega = np.random.randn(A.shape[1], l) Y = A @ Omega Q, _ = np.linalg.qr(Y) B = Q.T @ A U_tilde, S, V = np.linalg.svd(B, full_matrices=False) U = Q @ U_tilde return U, S, V

Theoretical Guarantees

LiMuon achieves a sample complexity of $O(\epsilon^{-3})$ for finding an $\epsilon$-stationary point in nonconvex stochastic optimization, under both standard Lipschitz and generalized smoothness conditions. This is a marked improvement over the $O(\epsilon^{-4})$ complexity of prior Muon variants, and matches the best known rates for matrix-aware optimizers, but without requiring large batch sizes.

Key theoretical results:

• Convergence under Lipschitz Smoothness: For batch size $b=1$, LiMuon finds an $\epsilon$-nuclear-norm stationary point in $O(\epsilon^{-3})$ samples.
• Convergence under Generalized Smoothness: The same sample complexity holds even when the strict Lipschitz condition is relaxed, which is critical for LLM training where such assumptions may not hold.

Empirical Evaluation

LiMuon was benchmarked against AdamW, Lion, SUMO, Muon, and Muon$^{++}$ on two tasks: LLMing with DistilGPT-2 on WikiText-103, and image classification with ViT on Tiny ImageNet.

DistilGPT-2 on WikiText-103

Figure 2: Performance comparison of training the DistilGPT-2 model on WikiText-103 dataset.

LiMuon(#1) achieves the lowest training and validation perplexity among all tested optimizers, with a significant reduction in memory usage compared to Muon and AdamW. LiMuon(#2) offers further memory savings with competitive performance.

Algorithm	Memory (GB)	Training PPL	Validation PPL
AdamW	7.34	717.22	998.83
Lion	7.30	315.29	491.07
SUMO	6.20	333.14	426.23
Muon	7.20	267.61	367.89
Muon$^{++}$	7.35	187.47	245.40
LiMuon(#1)	6.30	142.85	170.34
LiMuon(#2)	6.05	249.79	346.41

ViT on Tiny ImageNet

Figure 3: Performance comparison of training the ViT model on the Tiny ImageNet dataset (training loss and accuracy).

Figure 4: Performance comparison of training the ViT model on the Tiny ImageNet dataset (validation loss and accuracy).

LiMuon(#1) demonstrates faster convergence and higher final accuracy than all baselines, while LiMuon(#2) achieves the lowest memory footprint.

Algorithm	Memory (GB)	Train Acc (%)	Test Acc (%)
AdamW	4.40	20.31	17.23
Lion	4.40	23.44	20.06
SUMO	4.05	25.00	19.89
Muon	4.13	29.69	25.86
Muon$^{++}$	4.39	21.88	22.45
LiMuon(#1)	4.09	28.02	26.59
LiMuon(#2)	3.85	26.56	23.61

Implementation Considerations

• Computational Overhead: RSVD introduces additional computation per iteration, but the reduction in memory usage is substantial for large $m, n$.
• Hyperparameter Selection: The target rank $\hat{r}$ and oversampling parameter $s$ must be chosen to balance approximation quality and memory savings.
• Deployment: LiMuon is suitable for distributed and memory-constrained environments, such as multi-GPU or edge deployments for large models.
• Scalability: The optimizer is well-suited for models with millions to billions of parameters, as demonstrated in the experiments.

Implications and Future Directions

LiMuon advances the state of matrix-aware optimization for large-scale models by simultaneously reducing sample complexity and memory requirements. The theoretical analysis under generalized smoothness broadens its applicability to tasks where strict smoothness assumptions do not hold, such as LLM training. The empirical results suggest that LiMuon can be adopted as a drop-in replacement for existing optimizers in both NLP and vision domains, with immediate benefits in efficiency and scalability.

Potential future developments include:

• Extending LiMuon to tensor-structured parameters for convolutional architectures.
• Integrating adaptive rank selection for RSVD to further optimize memory-performance trade-offs.
• Exploring LiMuon in distributed and federated learning settings.

Conclusion

LiMuon introduces a principled approach to memory-efficient optimization for large models, leveraging low-rank momentum compression and variance reduction. It achieves lower sample complexity and memory usage than prior Muon variants, with strong empirical performance on both language and vision tasks. The optimizer is theoretically sound under both standard and generalized smoothness conditions, making it a robust choice for large-scale nonconvex stochastic optimization.