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

1. Problem Statement

Low-Rank Adaptation (LoRA) is widely adopted for fine-tuning large deep learning models due to its efficiency in reducing training FLOPs and memory requirements. However, standard optimizers like AdamW, when applied independently to LoRA factors ($A$ and $B$), introduce significant tuning difficulties compared to dense training. These issues include:

• Sensitivity to Initialization: Performance is highly dependent on initialization choices.
• Poor Hyperparameter Transfer: Optimal learning rates do not transfer well across different ranks, widths, depths, or factor scaling conventions.
• Divergent Dynamics: Even when starting from identical composed weights ($W = AB^T$), training dynamics often diverge if the factor scales drift apart.
• Suboptimal Performance: Native low-rank pretraining often struggles to match dense baselines without extra stabilization.

The core challenge is to design an optimizer for the low-rank manifold that maintains the geometric properties of full-rank optimizers (specifically Muon and the Shampoo family) while ensuring optimal learning rates transfer across rank, width, depth, and factor rescaling.

2. Methodology

2.1 Geometric Foundation

The paper reframes optimizer design as a geometric problem: identifying the norm under which an optimizer performs steepest descent.

• Muon Optimizer: Performs steepest descent under the spectral norm ($\|\cdot\|_{2\to2}$) in the ambient matrix space $\mathbb{R}^{m \times n}$. The update direction is given by the matrix sign function, $-\text{msign}(G)$.
• Low-Rank Manifold: LoRA restricts weights to a manifold $\mathcal{M}_r = \{W \in \mathbb{R}^{m \times n} : \text{rank}(W) = r\}$. The authors propose solving the steepest descent problem directly on the tangent space of this manifold, $T_W \mathcal{M}_r$, rather than optimizing factors $A$ and $B$ as independent coordinates.

2.2 Derivation of LoRA-Muon

The authors derive LoRA-Muon by solving the constrained linear minimization problem on the low-rank tangent space under the spectral norm.

1. Tangent Space Decomposition: Updates on $W$ are decomposed into $\Delta W = \Delta A B^T + A \Delta B^T$.
2. Trust Region Splitting: To make the problem tractable, the trust-region budget $\eta$ is split evenly between the two tangent components. This is justified by empirical observation that the two components rapidly align during training, making the triangle inequality nearly tight.
3. Projector Formulation: Using unitary invariance, the problem is reduced to finding Linear Minimization Oracles (LMOs) on projected gradients. The ideal update is:
   $$\Delta W = \frac{\eta}{2} \text{LMO}(\nabla W f \cdot P_B) + \frac{\eta}{2} \text{LMO}(P_A \cdot \nabla W f)$$
   where $P_A$ and $P_B$ are projectors onto the column spaces of $A$ and $B$.
4. Spectral Specialization: For the spectral norm, the LMO is $-\text{msign}(\cdot)$. The resulting factor updates (Equations 14–15) involve whitening the factor gradients by the current Gram matrices ($S_A = A^TA, S_B = B^TB$) and applying the spectral sign function:
   $$\Delta A^* = -\frac{\eta}{2} \text{msign}(\nabla_A f \cdot S_B^{-1/2}) S_B^{-1/2}$$
   $$\Delta B^* = -\frac{\eta}{2} \text{msign}(\nabla_B f \cdot S_A^{-1/2}) S_A^{-1/2}$$

2.3 Split Weight Decay

Standard factor-wise weight decay leads to a mismatch between the full-rank and low-rank update dynamics. The authors derive a split weight-decay rule where the decay factor $s = \sqrt{1 - \lambda \eta}$ is applied asymmetrically:
$$A_{t+1} = s A_t + \frac{1}{s} \Delta A^*_t, \quad B_{t+1} = s B_t + \frac{1}{s} \Delta B^*_t$$
This ensures the composed weight update matches the desired full-rank dynamics, preserving the transferability of hyperparameters.

2.4 Gauge Invariance

A critical property of LoRA-Muon is gauge invariance. The update on the composed weight $W$ remains invariant under arbitrary invertible transformations $(A, B) \to (AR, BR^{-T})$. This is achieved because the update depends only on the projectors $P_A$ and $P_B$, which are themselves gauge-invariant.

3. Key Contributions

1. Derivation of LoRA-Muon: A novel family of optimizers derived by solving the steepest descent problem on the low-rank manifold under unitary-invariant norms, specializing to the spectral norm.
2. Split Weight-Decay Rule: A specific update rule that aligns the dynamics of low-rank training with full-rank training, enabling hyperparameter transfer.
3. Proof of Gauge Invariance: Theoretical proof that the LoRA-Muon update is invariant under the full gauge group action, ensuring the optimizer depends on the induced weight update rather than arbitrary factor coordinates.
4. Comparison with Spectron: Demonstration that Spectron is not gauge-invariant and is sensitive to arbitrary factor rescalings, unlike LoRA-Muon.
5. Relation to LoRA-RITE: Proof that the simplified QR-coordinate core of LoRA-RITE is algebraically identical to the LoRA-Muon spectral update, though LoRA-Muon avoids the computational overhead of QR decompositions and second-moment storage.
6. Empirical Validation: Evidence that LoRA-Muon serves as an effective proxy for dense Muon learning rate search.

4. Experimental Results

Experiments were conducted on character-level language modeling using the TinyShakespeare dataset with compute-matched dense and low-rank baselines.

• Learning Rate Transfer: The optimal learning rate for dense Muon ($lr = 0.1$) transfers perfectly to LoRA-Muon across sweeps of rank (2, 4, 8, 16, 32), width, depth, and scalar factor rescaling.
  • At Rank 2, LoRA-Muon recovers the dense best-tested learning rate.
  • At Rank 32, LoRA-Muon achieves a lower mean validation loss ($1.776 \pm 0.002$) than the dense baseline ($1.789 \pm 0.002$) in the seed-averaged sweep.
• Gauge Rebalancing: When factors are initialized with a large scalar imbalance (scale 9), applying a rebalancing step has a negligible effect on LoRA-Muon (loss change $\approx -5.1 \times 10^{-5}$), confirming gauge invariance. In contrast, Spectron's performance changes significantly ($2.0 \times 10^{-2}$), confirming its sensitivity to factor scaling.
• Efficiency: A rank-2 proxy uses only ~4.3% of the trainable parameters of the dense model while supporting a ~9.2x larger optimizer-step budget under the same compute constraints.

5. Significance and Claims

The paper claims that LoRA-Muon provides a theoretically grounded solution to the tuning difficulties of LoRA. By treating LoRA as an optimization problem on the low-rank manifold rather than a factorized coordinate system, the authors achieve:

• Robustness: The optimizer is invariant to the arbitrary choice of factorization (gauge symmetry).
• Transferability: Optimal learning rates (and by extension, other hyperparameters via scaling laws) transfer seamlessly across rank, width, depth, and factor scales.
• Efficiency: It computes the spectral update without expensive QR decompositions or second-moment storage, making it accelerator-friendly.

The authors position LoRA-Muon as a "good low-rank proxy" for full-rank Muon and Shampoo-family optimizers. They emphasize that the low-rank parametrization should not dictate the optimizer's geometry; instead, the update should be defined by the geometry of the low-rank matrix manifold itself.

Limitations: The authors note that empirical claims are currently established on TinyShakespeare-scale models and have not yet been validated on large-scale pretraining or downstream fine-tuning benchmarks. The study focuses on learning rate transfer rather than a full hyperparameter sweep.