PaperlayerView exact revision on arXiv

Source-linked AI summary

Approximate Muon with low-rank adapters

Ben Anson, Conor Houghton, Edward Milsom

arXiv:2608.14492v1cs.LG

TL;DR

Muon’s whole-layer orthogonalization does not directly fit LoRA’s low-rank updates. The paper introduces sMuon, which approximates the low-rank Muon objective efficiently, and finds moderate, model- and evaluation-dependent improvements.

  • Problem

    Muon does not straightforwardly combine with LoRA because low-rank factor updates cannot generally realize whole-layer orthogonalization.

  • Method

    sMuon linearizes the low-rank Muon objective and derives an efficient matmul-only update within the expressible low-rank subspace.

  • Results

    Across eleven SFT benchmarks, sMuon was best on six tasks for the Muon-pretrained model, while AdamW was best on one.

  • Takeaways & Limitations

    Muon-based low-rank fine-tuning can provide moderate performance improvements, especially on Muon-pretrained models, while remaining relatively inexpensive.

  • Takeaways & Limitations

    The method’s inverse-root computations require jitter and semi-orthonormal initialization for numerical stability.

Abstract

from arXiv · show

The Muon optimizer shows clear benefits versus alternatives when pretraining neural networks. However, it is used less frequently for parameter-efficient fine-tuning (PEFT). One potential reason is that the most common PEFT method, LoRA, does not naturally combine with Muon since it is not mathematically possible to orthogonalize the weight update given by a low-rank parameterization. In this paper, we address this issue by approximating the solution to a relaxed Muon objective in the low-rank setting via linearization and then least-squares. We provide an efficient implementation that uses matmul operations only, as opposed to more complex linear algebra decomposition routines. Our method, sMuon (small Muon), performs favourably across SFT and a ReLoRA pretraining experiment. While results are model- and eval-dependent, we find overall that using Muon for low-rank fine-tuning provides moderate performance improvements.

1 Introduction

Muon orthogonalizes layer updates and has performed well in LLM pretraining, but its use in PEFT is limited because it does not straightforwardly combine with LoRA. This paper develops sMuon, a low-rank Muon method based on linearization, least-squares optimization, and efficient matrix multiplications.

  • Background: Muon orthogonalizes updates at each layer, equalizing singular values, and has often outperformed AdamW in LLM pretraining.It can be interpreted as steepest descent under the spectral norm, controlling worst-case changes to a layer as a linear operator.
  • Motivation: Muon is less commonly used for PEFT because it does not straightforwardly combine with LoRA, the most common PEFT method.Muon fine-tuning has also been shown to perform better with Muon-pretrained base models, which are currently less common than AdamW-pretrained models.
  • Method: The paper linearizes the low-rank Muon objective to find the optimal descent step within the available parameter space.The resulting solution orthogonalizes the full layer’s gradient and projects it onto the appropriate low-rank subspace.
  • Contributions: The method retains a more uniform weight-update spectrum than per-factor Muon and LoRA-Muon, while avoiding the expensive matrix decompositions used by Riemannion.The matmul-only implementation has comparable speed to per-factor Muon and LoRA-Muon and is significantly faster than Riemannion for larger LoRA rank.
  • Contributions: sMuon derives a closed-form LoRA training method from a first-order approximation of the spectral-norm-regularized steepest-descent objective.It requires less optimizer state than AdamW and uses a matmul-only algorithm designed for efficient modern hardware.

2 Related Work

Prior work establishes Muon as a spectral-steepest-descent optimizer with strong pretraining efficiency, while highlighting optimizer mismatch in fine-tuning. Related methods approximate Muon-compatible low-rank optimization for LoRA or efficient pretraining, but differ in formulation and setting.

  • Muon and spectral steepest descent: Muon orthogonalizes each weight-matrix gradient step, equalizing update singular values and making the update steepest descent under the spectral norm.It has scaled to billion- and trillion-parameter LLM pretraining with efficiency gains over AdamW.
  • Muon for fine-tuning and optimizer mismatch: Muon is used less for fine-tuning, and switching between AdamW and Muon during fine-tuning can degrade performance.Prior work attributes this to differing implicit biases, while LoRA-style constraints may narrow the gap.
  • Geometry-aware and Muon-based LoRA optimizers: Cesista et al. and Bogachev et al. approximate low-rank solutions to spectral-regularized gradient descent, with LoRA-Muon using a triangle-inequality relaxation and first-order minimization.The paper likewise linearizes the objective, as indicated in the supplied passage.
  • Low-rank optimization beyond LoRA fine-tuning: Other methods target efficient low-rank pretraining through stabilization or optimizer-state compression, while low-rank Muon orthogonalization operates on dense weights rather than LoRA adapters.Examples include Spectron, LORO, GaLore, LoRA-Pre, and He et al.’s low-rank orthogonalization routine.

3 Background

Muon chooses weight updates by regularized spectral-norm minimization, yielding an orthogonalized gradient or momentum direction. LoRA instead represents fine-tuning updates as low-rank adapters, motivating an analogue of Muon’s whitened update in that setting.

  • Muon: Muon selects the weight step direction through regularized minimization under the spectral norm.The spectral norm is the largest singular value.
  • Muon: For full-rank gradients, Muon’s solution is proportional to −msign(GW).Here, msign(X) = UV^⊤ for the singular value decomposition X = UΣV^⊤.
  • Muon: The Muon direction orthogonalizes the gradient or, in practice, the momentum.For low-rank gradients, the solution is non-unique; the stated convention promotes non-zero singular values and leaves the rest zero.
  • LoRA: LoRA adds low-rank adapters ∆W = BA to a pretrained weight W0, with rank r much smaller than both layer dimensions.The factors B and A are usually optimized with AdamW, while the paper seeks a Muon analogue for this whitened update.

4 Method

sMuon adapts Muon to low-rank LoRA updates by linearizing the factorized objective, solving a least-squares approximation, and replacing the full gradient with a low-rank equivalent. An efficient implementation reduces orthogonalization to small matrix operations and combines the updates with split weight decay and momentum projection.

  • Problem formulation: LoRA factor updates cannot generally realize the Muon target because their induced update has rank at most 2r, whereas msign(G∆W) is usually full rank.Simply orthogonalizing δB and δA does not make the induced adapter update δ∆W orthogonal.
  • Linearized least squares: The method linearizes δ∆W = δB A + B δA + δB δA by dropping second-order terms, then solves a least-squares problem for the factor updates.The original factorized optimization is nonlinear and has no known closed-form solution.
  • Low-rank gradient: The low-rank gradient H projects G∆W onto updates expressible by the current LoRA adapter and is first-order equivalent to G∆W for factor perturbations.This permits using δ∆W⋆ = −η msign(H) when linearizing the Muon objective.
  • Efficient implementation: The efficient construction reduces a d1 × d2 orthogonalization to a 2r × 2r orthogonalization using bases spanning the column and row spaces of H.Because H = BX + YA, its column and row spaces are contained in sums of the corresponding spaces of B, Y, A, and X.
  • Efficient implementation: All components of the factor updates can be computed with matmuls, inverse roots, and msign, yielding the matmul-only sMuon algorithm for LoRA adapters.The optimizer also incorporates split weight decay and momentum projection for first-order weight-decay matching and reparameterization invariance.

5 Experiments

Section 5 evaluates sMuon against several LoRA optimizer baselines on supervised fine-tuning benchmarks and ReLoRA pretraining. Results are model- and evaluation-dependent but generally favor sMuon, particularly with the Muon-pretrained Moonlight-16B-A3B model, while its matmul-only implementation remains computationally practical.

  • SFT comparisons: sMuon is compared with LoRA plus AdamW, per-factor Muon, LoRA-Muon, and Riemannion across commonsense and coding benchmarks after SFT.The evaluation covers four models, though the supplied passage truncates before listing them.
  • Interpretation: The authors attribute gains over AdamW and per-factor Muon to sMuon's more accurate approximation of the Muon objective.This explanation is presented as an attribution rather than a separately established causal result.
  • SFT comparisons: Across eleven aggregated SFT tasks, AdamW is best on one task, whereas sMuon is best on six, with particularly strong performance on Moonlight-16B-A3B.The results remain model- and evaluation-dependent, and per-factor Muon can be markedly weaker on individual tasks.
  • ReLoRA pretraining: In ReLoRA pretraining, every Muon-based method achieves lower validation loss than LoRA+AdamW, and sMuon attains the joint-best validation loss.The experiment trains a 162M-parameter transformer from scratch for 3000 steps, merging and reinitializing rank-64 adapters every 500 steps.
  • Efficiency: AdamW is fastest, while all Muon optimizers are slightly slower; Riemannion adds huge overhead at larger rank, whereas sMuon's extra matmuls remain negligible relative to total step time.Timing uses a 12-layer GPT model on an A100, averaging 25 optimizer steps across ranks r ∈{24, 25, 26, 27}.

6 Conclusion · A Other algorithmic details

The paper presents sMuon as a cheap, parameterization-invariant approximate Muon update for low-rank adapters. It also notes that sMuon incorporates pre-existing techniques discussed in the algorithmic-details appendix.

  • 6 Conclusion: sMuon optimizes low-rank adapters using an approximate form of Muon.
  • 6 Conclusion: The resulting optimization step is available in closed form.
  • 6 Conclusion: The step is invariant to the adapter reparameterization (B, A) 7→(BS, S−1A).
  • 6 Conclusion: Its computational cost is O(r2(d1 + d2) + r3) time and O(r max(d1, d2)) memory.The method uses less optimizer state than LoRA+AdamW.
  • 6 Conclusion: sMuon has empirically flatter spectra than per-factor Muon.
  • 6 Conclusion: sMuon out-performs other methods on the Moonlight Muon-pretrained model.
  • A Other algorithmic details: The paper discusses pre-existing techniques incorporated into sMuon.

A.1 Match AdamW RMS · A.2 Split weight decay · A.3 Invariance to parameterization

The method adjusts LoRA learning rates to match AdamW RMS behavior, adopts split weight decay to preserve first-order dynamics, and transports momentum to maintain invariance to adapter parameterization. With momentum projection into the buffers at every step, differently parameterized runs orthogonalize the same matrix and realize equivalent updates.

  • A.1 Match AdamW RMS: The full-parameter learning-rate scaling is 0.2 p.This adjustment is adopted following Liu et al. [2025].
  • A.1 Match AdamW RMS: The overall update, including weight decay, is matched in the LoRA setting through a corresponding adjustment.The passage introduces the scaling expression involving max(din, dout).
  • A.1 Match AdamW RMS: Orthogonalized LoRA weight updates have rank at most 2r because they can be written as the difference of two rank-r matrices.This provides the stated heuristic justification for the learning-rate adjustment.
  • A.2 Split weight decay: Applying weight decay separately to B and A produces dynamics different from full-parameter training.The paper identifies this as a mismatch requiring a separate treatment.
  • A.2 Split weight decay: Split weight decay is adopted to preserve first-order weight-decay dynamics, with s = √1 −λη.The resulting factor updates are identified with decoupled weight decay.
  • A.3 Invariance to parameterization: Momentum is accumulated separately for the adapter matrices and substituted into the factorization equation to reconstruct effective momentum.The natural accumulations are MA ←βMA + (1 −β)GA and MB ←βMB + (1 −β)GB.
  • A.3 Invariance to parameterization: Transporting momentum buffers fixes parameterization dependence by re-deriving them from ¯H against the current factors and projecting ¯H into the buffers at every step.The projection ensures that two differently parameterized runs orthogonalize the same matrix and therefore realize agreeing steps.

B Alternative stable algorithm

Algorithm 3 provides a mathematically equivalent but more stable sMuon procedure by using two orthogonalization steps and relative jitter. Additional inverse-root interventions offer only small practical benefits despite being slightly faster for small r.

  • Stable algorithm: Algorithm 3 is a mathematically equivalent, more stable alternative to Algorithm 2, using two orthogonalization steps to preserve orthogonality to machine precision.It also uses slightly different U2 and V2 definitions.
  • Stable algorithm: Relative jitter improves stability for larger r, with Algorithm 3 defaulting to ϵ = 10^-4.The algorithm takes jitter ϵ > 0 as an input.
  • Stable algorithm: The alternative formulation is valid because col(Y) = col(˜Y), and similarly for Z and ˜Z, when SB and SA are invertible.Invertibility of SB and SA is assumed.
  • Stable algorithm: Replacing msign(X) with X invroot(X⊤X) is slightly faster for small r, but its practical benefits were small.The speed advantage arises because X⊤X is an r × r matrix.

C Comparison of Muon-based LoRA optimizers

Muon-based LoRA optimizers share a structural obstruction: Muon’s ideal update generally lies outside the realizable LoRA step set, so methods must use different approximations. The paper compares these methods under a common update convention in Table 5.

  • Shared obstruction: Muon’s ideal step msign(G∆W ) is generally not realizable by a LoRA adapter because LoRA gradient steps are constrained to T = {Z : (I −PB)Z(I −PA) = 0}.The obstruction arises because the ideal Muon direction is not generally contained in the LoRA-realizable step set.
  • Method comparison: Different Muon-based LoRA methods address the obstruction through different approximations.The comparison is motivated by the fact that the ideal Muon update cannot generally be used directly within the LoRA constraint.
  • Method comparison: Table 5 compares the methods using the convention ∆W = BA, with ΠT (·) denoting the orthogonal projector onto realizable step directions.The table standardizes notation for the LoRA update and its projection onto the realizable subspace.

D Optimality of the sMuon step

Lemma 1 establishes that the chosen low-rank adapter updates minimize, to first order, the Frobenius distance to the target Muon step. Although the minimizing factor updates are non-unique, every first-order minimizer produces the same overall linear-layer update.

  • First-order optimality: The choices δA = −η B† msign(H) and δB = −η (I −PB) msign(H) A† minimize the target-step error to first order.The minimized quantity is ∥δ∆W⋆−δ∆W∥F, with the approximation valid up to O(η²).
  • Non-uniqueness: The first-order minimizer is not unique in the factor updates (δB, δA).If (δB, δA) is minimizing, then (δB − BS, δA + SA) is also minimizing for S = O(η).
  • Non-uniqueness: Every first-order minimizer realizes the same overall update to the linear layer.The transformed pair preserves δB A + BδA, and therefore has the same leading coefficient in the first-order error.
  • Proof strategy: The proof reduces the optimization to minimizing the leading coefficient ∥M + bA + Ba∥F after writing δB = ηb and δA = ηa.The Frobenius-orthogonal block decomposition shows the chosen updates attain equality and hence minimize the first-order error.

E Learning rate curves

This section examines learning-rate curves for commonsense runs across optimizers and base models, selecting main-text accuracy results by held-out validation loss.

  • Learning rate selection: Optimal learning rates vary between optimizers, so the study sweeps learning rates separately for each optimizer.Main-text accuracy uses the learning rate selected by loss on a held-out validation set.
  • Accuracy curves: Figure 2 reports overall commonsense benchmark accuracy at the final training step across learning rates for each base model.Each curve represents one optimizer.
  • Validation-loss curves: Figure 3 reports validation loss against learning rate for each base model over the same runs as Figure 2.These validation-loss curves support the learning-rate selection procedure.
Loading 2608.14492v1…