LoRA-Muon: Spectral Steepest Descent on the Low-Rank Manifold

Abstract: Low-Rank Adaptation (LoRA) significantly reduces compute and memory costs for finetuning Deep Learning models but is often harder to tune than dense training: when using factor-wise optimizers such as AdamW, it is sensitive to initialization choices, its optimal learning rates transfer poorly across ranks, and it often fails to beat dense baselines. We derive LoRA-Muon by applying the Muon optimizer's spectral steepest-descent rule to the low-rank setting. Along with our split weight-decay rule, our main claim is that LoRA-Muon is a good low-rank proxy for full-rank Muon and Shampoo-family optimizers. Its optimal learning rates transfer across rank, width, depth, and factor-rescaling. In our compute-matched TinyShakespeare study, a rank-$2$ proxy recovers the dense best tested learning rate, and a rank-$32$ LoRA-Muon run attains lower mean validation loss than the dense baseline in the seed-averaged sweep. We further show that the Spectron optimizer depends on arbitrary factor scaling, so it would likely be a poor fit when finetuning starts from badly imbalanced factors, and that LoRA-RITE's simplified QR-coordinate core implements the same spectral update. LoRA-Muon computes that update without QR-decomposition and avoids storing second moments, making it more accelerator-friendly and memory-efficient.

• The paper introduces a gauge-invariant spectral steepest descent update on the low-rank manifold, mitigating tuning sensitivity and ensuring robust hyperparameter transfer.
• The paper leverages a tangent space projection with split weight decay to achieve efficient, QR-free updates that closely match dense optimizer performance.
• The paper validates its approach on TinyShakespeare-scale models, demonstrating improved validation loss and computational savings compared to dense methods.

LoRA-Muon: Spectral Steepest Descent on the Low-Rank Manifold

Introduction

"LoRA-Muon: Spectral Steepest Descent on the Low-Rank Manifold" (2606.12921) fundamentally re-examines optimization in low-rank adaptation (LoRA) by leveraging a geometric reduction of Muon's spectral steepest-descent update to the low-rank manifold. The primary objective is to alleviate tuning sensitivity and initialization fragility that arise when popular optimizers like AdamW are naïvely applied in factor space, with the broader aim of creating optimizers whose key hyperparameters—especially learning rate—transfer robustly across ranks, widths, depths, and arbitrary factor rescalings. The theoretical centerpiece is a gauge-invariant update derived from first principles that eschews coordinate artifacts, and an empirical validation that LoRA-Muon matches and, in certain regimes, surpasses dense Muon performance with significant computational savings.

Geometric Derivation of LoRA-Muon

The authors recast LoRA as a problem of steepest descent restricted to the rank-$r$ manifold, generalizing Muon's spectral norm update from ambient matrix space to tangent spaces of the low-rank manifold. Given a parameterization $W = AB^T$, with $A \in \mathbb{R}^{m \times r}$ and $B \in \mathbb{R}^{n \times r}$, the optimal update is a sum of two tangent components, which are shown to quickly become highly aligned in practice.

Figure 1: Alignment between the two tangent update components in the TinyShakespeare study, validating the near-tightness of the triangle-inequality split in the trust-region formulation.

By splitting the trust-region budget evenly between the $A$ and $B$ directions and exploiting projector-based decompositions, the construction naturally yields updates in Gram-matrix coordinates, directly connecting to the matrix sign function and removing the necessity of QR factorizations. This approach confers crucial gauge invariance, namely that updates are independent of arbitrary invertible reparameterizations of $A$ and $B$. LoRA-Muon's spectral update inherits all the favorable geometry of Muon in full space, as depicted below.

Figure 2: The LoRA-Muon update geometrically represents the projection of the ambient Muon update onto the tangent space of the low-rank manifold, inheriting symmetry under gauge transformation.

Split Weight Decay and Optimizer Implementation

A critical correction is introduced for the application of weight decay. Directly applying factor-wise decay yields parameter norms and dynamics mismatched with dense training, leading to impaired hyperparameter transfer. LoRA-Muon integrates a split weight decay update ensuring that the weight decay on $W = AB^T$ matches that of the full-rank optimizer, up to negligible higher-order terms. This facilitates direct transfer of scaling laws and other hyperparameter schedules from dense to low-rank training regimes. The implementation abides by the efficient use of first moments and Gram matrix inverse-square-roots, avoiding second moments and ill-conditioned triangle solves, which makes LoRA-Muon friendly to accelerator deployment and low-precision arithmetic.

Gauge Invariance and Comparison with Baselines

A distinguishing feature of LoRA-Muon is its full invariance to arbitrary invertible "gauge" transformations of the LoRA factors; i.e., $(A, B) \mapsto (AR, BR^{-T})$ does not affect the update on $W = AB^T$0. This symmetry extends to both the main step and the split weight decay, as proven constructively in the main text and appendices. This property is empirically confirmed in the gauge-rebalancing ablation, which exhibits negligible change in validation loss for LoRA-Muon under extreme factor rescalings.

Figure 3: Validation loss is invariant for LoRA-Muon under scalar gauge rebalancing, while the same is not true for Spectron, emphasizing the importance of gauge invariance for robust optimization.

The analysis reveals that popular alternatives such as Spectron do not enjoy gauge invariance and in fact exhibit strong dependence on arbitrary initialization artifacts, as both the trust region and effective update scale are functions of factor norms. LoRA-RITE, in contrast, is shown to realize the same geometric update as LoRA-Muon in QR coordinates, albeit with greater implementational complexity and higher sensitivity to numerical issues due to additional second-moment and mass-escape state.

Empirical Results

Comprehensive experiments on TinyShakespeare-scale LLMs validate the main claims:

• Learning Rate Transferability: LoRA-Muon matches dense Muon's optimal learning rate across sweeps in LoRA rank, model width, depth, and across a wide range of scalar rescalings of $W = AB^T$1 and $W = AB^T$2. Already at rank $W = AB^T$3, LoRA-Muon recovers the dense Muon optimum; at rank $W = AB^T$4, it obtains a lower mean validation loss than the dense baseline.

Figure 4: Validation loss vs learning rate across four critical sweeps (rank, width, depth, gauge scale); LoRA-Muon tracks dense Muon's optimum in all regimes.

• Compute Efficiency: LoRA-Muon with low rank matches or surpasses dense Muon with a fraction of the parameters and allows for a substantially larger number of optimizer steps under equal compute budgets.
• Robustness to Gauge Rebalancing: Validation loss statistics for LoRA-Muon are insensitive to scalar factor rescalings and associated rebalancing, while Spectron's loss curve shifts substantially under such interventions.

Theoretical and Practical Implications

The theoretical framework underscores the necessity of defining optimization geometry on the quotient of the low-rank parameterization, rather than factor space. By constructing updates invariant to coordinate transformations, LoRA-Muon avoids pitfalls of standard factor-wise optimizers and establishes a bridge to scaling laws and hyperparameter transfer in both practical low-rank adaptation and settings requiring native low-rank pretraining.

Practically, LoRA-Muon can be used as a lightweight proxy for Muon and other Shampoo-family optimizers for both learning-rate search and downstream LoRA finetuning. Its low-memory, QR-free, and second-moment-free core make it well-suited for accelerator implementations and mixed-precision environments. The framework's extension to arbitrary unitary-invariant norms and compatibility with modern hyperparameter scaling laws further enhances its utility for practical deployment and scaling research.

Conclusion

This work provides a rigorous geometric foundation for low-rank optimization by formulating and realizing a spectral, gauge-invariant steepest descent update on the low-rank manifold. Empirical studies demonstrate that LoRA-Muon efficiently recovers dense hyperparameter optima and remains robust to initialization and coordinate choices, significantly outperforming alternative baselines such as Spectron when it comes to gauge stability. These results advocate for viewing LoRA optimization as a quotient geometry problem and point towards future work in extending these geometric optimizer constructions for large-scale settings, non-spectral geometries, and broader classes of low-rank and structured adaptation methods.