🤖 machine learning

Approximate Muon with low-rank adapters

This paper introduces sMuon, an efficient method that approximates the Muon optimizer for low-rank parameter-efficient fine-tuning by linearizing the objective and solving it via least-squares, thereby enabling Muon's performance benefits in settings like SFT and ReLoRA pretraining where it was previously incompatible with LoRA.

Original authors: Ben Anson, Conor Houghton, Edward Milsom

Published 2026-08-17

📖 4 min read☕ Coffee break read

CC0 1.0

Original authors: Ben Anson, Conor Houghton, Edward Milsom

Original paper dedicated to the public domain under CC0 1.0 (http://creativecommons.org/publicdomain/zero/1.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 teach a giant, super-smart robot a new skill, like writing poetry or solving math problems. You don't want to retrain the robot's entire brain from scratch because that takes forever and costs a fortune. Instead, you want to attach a small, lightweight "training module" to its existing brain. This is called Parameter-Efficient Fine-Tuning (PEFT). Think of it like putting a specialized headset on a general-purpose robot; the headset teaches it the new trick without changing the robot's core hardware.

One popular way to build these headsets is called LoRA (Low-Rank Adaptation). It's like building the headset out of two thin, flat sheets of paper that, when stacked, create the shape of the new skill. To teach the robot, we need an optimizer—a mathematical coach that tells the headset how to adjust its shape to get better. The most common coach is called AdamW, but there's a newer, more powerful coach called Muon. Muon is famous for being incredibly efficient when pretraining (teaching the robot its very first lessons from scratch). It works by making sure the robot's brain updates happen in a perfectly balanced, "orthogonal" way, like a dancer spinning perfectly on one foot without wobbling.

However, there's a problem. When you try to use Muon to train the small LoRA headset, the math breaks down. Muon wants to balance the entire layer of the robot's brain, but the LoRA headset is just a tiny, low-rank slice of that layer. It's like trying to use a conductor's baton to direct a single violinist in a way that makes sense for the whole orchestra; the geometry just doesn't fit. The standard workaround has been to ignore the orchestra and just tell the violinist to spin on their own, but that misses the magic of Muon. This paper asks: Can we find a way to make Muon work for these small headsets without breaking the math or slowing everything down?

The authors, Ben Anson, Conor Houghton, and Edward Milsom, say yes. They introduce a new method called sMuon (small Muon). Instead of trying to force the impossible, they use a clever mathematical trick: they "linearize" the problem. Imagine you are trying to walk a tightrope while holding a giant, wobbly umbrella. Instead of trying to balance the whole umbrella at once, you pretend the umbrella is just a straight stick for a split second, calculate the perfect step, and then adjust. By doing this, they find the best possible way to update the LoRA headset that approximates the perfect Muon dance.

The paper finds that this new sMuon method works surprisingly well. When they tested it on various language models, sMuon often performed better than the standard AdamW coach and other Muon attempts. It was particularly impressive when used on a robot that had already been pre-trained with Muon (called Moonlight), where sMuon achieved the top scores in six out of eleven different tasks. In a pre-training experiment, sMuon tied for the best results in lowering the error rate, but it did so about 30% faster than its closest rival, a method called Riemannion.

Crucially, the authors show that their method is not just a theoretical idea; it is also practical. While other advanced methods require heavy, slow calculations (like taking apart and reassembling complex matrices), sMuon relies almost entirely on simple matrix multiplications—the kind of math that modern computer chips are built to do at lightning speed. This means sMuon is fast enough to be used in real-world training without slowing down the process. The paper suggests that by respecting the geometry of the whole layer rather than treating the LoRA parts as separate, unrelated pieces, we can get a significant performance boost with very little extra cost. It's a reminder that sometimes, the best way to teach a giant robot a new trick is to give it a small, smart, and mathematically elegant nudge.

=== SUMMARY ===

Technical Summary: Approximating Muon with Low-Rank Adapters (sMuon)

Problem Statement

The Muon optimizer has demonstrated superior performance in pretraining large language models (LLMs) by orthogonalizing weight updates at each layer, effectively equalizing singular values. This behavior corresponds to steepest descent under the spectral norm, controlling the worst-case change to a layer as a linear operator rather than a flattened vector. However, Muon is rarely used for Parameter-Efficient Fine-Tuning (PEFT), particularly with Low-Rank Adaptation (LoRA).

The core mathematical obstruction is that Muon requires orthogonalizing the descent step for the entire weight matrix layer. In the LoRA setting, where the weight update is constrained to a low-rank parameterization $\Delta W = BA$ (with $B \in \mathbb{R}^{d_1 \times r}$ and $A \in \mathbb{R}^{r \times d_2}$ ), it is mathematically impossible to generally find updates $\delta B$ and $\delta A$ such that the resulting $\delta \Delta W$ perfectly matches the full-rank orthogonalized gradient required by Muon.

Existing workarounds include:

Per-factor Muon: Orthogonalizing $B$ and $A$ separately, which ignores the geometry of the full layer.
LoRA-Muon: Relaxing the objective using the triangle inequality.
Riemannion: Formulating the problem as Riemannian optimization on a low-rank manifold, which requires expensive truncated Singular Value Decompositions (SVDs) during the retraction step, leading to significant computational overhead.

Methodology

The authors propose sMuon (small Muon), a method that approximates the solution to a relaxed Muon objective within the low-rank setting. The approach relies on linearization and least-squares optimization to find the optimal descent step.

1. Linearization and Least-Squares Formulation

The goal is to find updates $\delta B, \delta A$ such that the implied update $\delta \Delta W \approx \delta B A + B \delta A$ approximates the ideal Muon step $\delta \Delta W^* = -\eta \cdot \text{msign}(G_{\Delta W})$ , where $\text{msign}(X) = UV^\top$ is the polar factor of the SVD.

Since the ideal step is generally full-rank and cannot be realized exactly by low-rank factors, the authors linearize the problem by dropping second-order terms ( $\delta B \delta A$ ) and solving a least-squares problem:
$(\delta B^*, \delta A^*) = \arg \min_{\delta B, \delta A} \| \delta \Delta W^* - (\delta B A + B \delta A) \|_F^2$
This yields a closed-form solution involving Moore-Penrose inverses ( $B^\dagger, A^\dagger$ ). Crucially, the authors show that the optimal update is not unique in terms of the factors $(\delta B, \delta A)$ , but the resulting first-order change to the full layer $\delta \Delta W$ is unique.

2. Efficient Matmul-Only Computation

A direct implementation of the closed-form solution requires orthogonalizing a large $d_1 \times d_2$ matrix, which is computationally prohibitive. To address this, the authors derive an efficient formulation (Algorithm 2) that avoids explicit large-matrix orthogonalization:

Low-Rank Gradient: Instead of using the full gradient $G_{\Delta W}$ , they utilize a low-rank projection $H = P_B G_{\Delta W} + (I-P_B)G_{\Delta W} P_A$ , which is first-order equivalent for the optimization objective.
Dimensionality Reduction: They exploit the identity $\text{msign}(H) = U \text{msign}(U^\top H V) V^\top$ , where $U$ and $V$ are orthonormal bases spanning the column and row spaces of $H$ . Since $H$ has rank at most $2r$ , this reduces the expensive $d_1 \times d_2$ orthogonalization to a $2r \times 2r$ orthogonalization.
Implementation: The entire algorithm relies solely on matrix multiplications, matrix inverse square roots, and polar decompositions of small $r \times r$ or $2r \times 2r$ matrices. This avoids the overhead of QR and SVD decompositions required by Riemannion.

3. Algorithmic Details

The sMuon optimizer incorporates:

Split Weight Decay: To match the dynamics of full-parameter training, weight decay is applied via a scaling factor $s = \sqrt{1-\lambda\eta}$ to the factors $A$ and $B$ before the gradient update.
Momentum Transport: To ensure parameterization invariance (i.e., the optimizer behaves consistently regardless of the specific factorization $B, A$ ), momentum buffers are projected onto the current factor spaces at every step.

Key Contributions

Derivation of sMuon: A closed-form method for training LoRA adapters with Muon by taking a first-order approximation of the spectral-norm-regularized steepest descent objective.
Efficient Implementation: A "matmul-only" algorithm that is computationally efficient for modern hardware, avoiding expensive decompositions. It offers comparable speed to Per-factor Muon and LoRA-Muon while being significantly faster than Riemannion, especially for larger LoRA ranks.
Empirical Evaluation: Extensive experiments across four base models (including one Muon-pretrained and three AdamW-pretrained) and eleven benchmarks (commonsense and coding).

Results

The paper reports the following findings:

Performance: sMuon performs favorably across SFT benchmarks. While results are model-dependent, sMuon achieved the best accuracy in six out of eleven tasks when fine-tuning the Moonlight-16B-A3B model (which was pretrained with Muon). It generally outperforms AdamW and Per-factor Muon.
ReLoRA Pretraining: In a pretraining setting using ReLoRA (periodically merging and refreshing adapters), sMuon achieved the lowest final validation loss, tying with Riemannion but being approximately 30% faster.
Spectral Properties: Empirical analysis of the weight update spectra shows that sMuon retains a more uniform spectrum (closer to the ideal Muon behavior) compared to Per-factor Muon and LoRA-Muon.
Efficiency: Table 4 demonstrates that while sMuon is slightly slower than AdamW, it is significantly faster than Riemannion. For a rank of $r=64$ , Riemannion adds ~33% overhead to the step time, whereas sMuon adds only ~0.8%.

Significance and Claims

The paper claims that sMuon provides a principled, memory-efficient, and computationally tractable way to apply Muon optimization to low-rank adapters. By respecting the geometry of the full layer rather than treating LoRA factors as independent matrices, sMuon achieves moderate performance improvements, particularly when the base model was also pretrained with Muon. The authors emphasize that their method avoids the "expensive manifold constraints" of Riemannian optimization without sacrificing performance, making Muon-based fine-tuning more accessible for practical applications.

The authors remain modest, noting that results are "model- and eval-dependent" and that the primary contribution is a clean, efficient approximation that bridges the gap between Muon's theoretical benefits and the practical constraints of PEFT.