PaperlayerView exact revision on arXiv

Source-linked AI summary

LoRA-TSD: Tangent-Space Spectral Descent for LoRA via Muon-Style Updates

Dmitrii Andriianov, Andrey Veprikov, Aleksandr Beznosikov

arXiv:2609.02734v1cs.LG

TL;DR

Independent LoRA factor updates can ignore the geometry of the low-rank weight change they induce. LoRA-TSD performs Muon-style spectral steepest descent in the tangent space of the fixed-rank manifold and retracts the result through the LoRA factors. Across the evaluated benchmarks, it outperforms competing LoRA optimizers and remains stable across adapter ranks.

  • Problem

    Independent optimization of LoRA’s two factors does not account for the geometry of the induced low-rank weight update.

  • Method

    LoRA-TSD computes a spectral-norm tangent-space update and maps it back to standard LoRA factors through a factor-induced retraction without materializing full weight matrices.

  • Results

    Across evaluated benchmarks, LoRA-TSD outperforms generic and LoRA-specific baselines and is the most stable method across adapter ranks.

  • Takeaways & Limitations

    The tangent-projected gradient provides a natural LoRA stationarity measure computable from factor gradients, with convergence guarantees for LoRA-TSD and LoRA-Pro.

  • Takeaways & Limitations

    The convergence analysis assumes that the full-space gradient is 1/λ-Lipschitz continuous with respect to the spectral norm.

Abstract

from arXiv · show

Low-rank adaptation (LoRA) is the standard way to fine-tune large models, yet when its two factors are trained independently, the update ignores the geometry of the low-rank weight change it induces. We introduce LoRA-TSD, an optimizer that treats every LoRA step as a tangent vector of the fixed-rank matrix manifold and takes the spectral-norm steepest-descent step of Muon inside that tangent space, mapping the result back to the factors through a retraction native to the LoRA parametrization. The step avoids expensive operations on full weight matrices, and its retraction is up to $2.8\times$ cheaper than the truncated-SVD retraction used by prior manifold methods. We prove that the Frobenius-norm version of our surrogate recovers LoRA-Pro, and we identify the tangent-projected gradient, the Riemannian gradient of the manifold, as the stationarity measure natural to LoRA training and computable from the factor gradients alone. Under this measure we give the first global convergence guarantees for both LoRA-Pro and LoRA-TSD, with rates that drive the factor-gradient norms to zero. Across six commonsense and natural-language-inference benchmarks with Llama-3.2-1B, Llama-3.1-8B and Qwen3-32B, LoRA-TSD outperforms every competing LoRA optimizer and stays robust to the adapter rank. Code is available at https://github.com/brain-lab-research/LoRA-TSD.

1 Introduction

LoRA-TSD treats LoRA updates as tangent-space steps on the fixed-rank manifold and maps spectral-norm Muon-style directions back through the LoRA factors. It provides convergence guarantees and strong, rank-robust benchmark performance.

  • Method: LoRA-TSD applies Muon’s spectral update within the tangent space of the fixed-rank LoRA manifold.The method computes the low-rank weight update and maps it back to the factors through a factor-induced retraction.
  • Efficiency: The factor-induced retraction operates through standard LoRA factors without materializing m × n matrices.Both the proposed and prior truncated-SVD retractions admit low-rank implementations, but LoRA-TSD avoids the truncated SVD.
  • Theory: LoRA-TSD has a stationary-point convergence guarantee in the L-smooth setting for momentum-free updates.The guarantee is stated for the tangent-space formulation of the LoRA problem.
  • Theory: The Frobenius-norm counterpart of LoRA-TSD recovers the LoRA-Pro direction and also has a stationarity convergence guarantee.This connects the proposed surrogate to an established LoRA-specific optimizer.
  • Results: Across evaluated benchmarks, LoRA-TSD outperforms generic and LoRA-specific baselines and is the most stable method across adapter ranks.The evaluation covers six commonsense and natural-language-inference benchmarks across three model scales.

2 Related Work

Prior LoRA optimizers address factorization-induced mismatch through reparameterization, update correction, or fixed-rank manifold geometry. LoRA-TSD instead uses spectral-norm tangent geometry while preserving transformation invariance and recovering LoRA-Pro under the Frobenius norm.

  • LoRA-specific optimizers: LoRA’s bilinear factor update can be misaligned with weight-space descent because it depends on factor norms and conditioning.This motivates optimizers that correct induced weight updates rather than independently optimizing the factors.
  • Update correction: LoRA-Pro corrects the induced update while LoFT projects Adam moments into the low-rank subspace.These methods retain the LoRA parametrization but modify optimization behavior.
  • LoRA-TSD: LoRA-TSD satisfies the same transformation invariance as LoRA-Rite while selecting directions through spectral-norm tangent geometry.In the Frobenius-norm case, it recovers the LoRA-Pro direction.
  • Spectral-norm methods: Muon orthogonalizes momentum updates and corresponds to steepest descent under the spectral norm, but assumes access to the full weight matrix.LoRA-TSD transfers this geometry into the tangent space of the fixed-rank manifold.
  • Manifold optimization: Riemannion, LoRA-Muon, and LoRA-TSD use spectral tangent updates, while their approximations and mappings back to LoRA factors differ.LoRA-TSD’s factor-induced retraction distinguishes it from prior manifold approaches.

3 Method

The method derives LoRA optimization in the tangent space of the fixed-rank manifold, computes spectral-norm descent from factor gradients, and retracts the result back to LoRA factors without full-matrix operations.

  • Fixed-rank geometry: LoRA residuals lie on a fixed-rank manifold, whose tangent space contains the first-order factor-induced updates.After dropping the second-order factor product, this tangent space is the exact first-order search space of LoRA.
  • Optimization surrogate: The local surrogate optimizes the first-order LoRA weight change directly because factor increments are non-unique.This removes dependence on a particular split of the same tangent weight update.
  • Tangent-space spectral descent: LoRA-TSD retains the spectral norm and uses a linear-minimization-oracle formulation to obtain a Muon-style tangent update.Alternating projections refine the tangent direction toward the spectral-norm solution.
  • Tangent-space spectral descent: The refined update nearly equalizes its first 2r singular values, providing the tangent-space analogue of Muon.Theorem 1 establishes convergence for this update under the stated assumptions.
  • Tangent projection: The tangent-projected gradient can be computed solely from the observable factor gradients, without forming the full-space gradient.The closed form uses the factor gradients and the factor Gram-matrix projectors.
  • Retraction: The resulting tangent update is split canonically into factor updates and applied through a LoRA-native retraction.The conceptual algorithm computes the projected gradient, refines it spectrally, and maps it back to LoRA factors.
  • Invariance: Transformation invariance follows because the tangent projector and spectral tangent step depend on the represented weight, not the chosen factors.LoRA-TSD is transformation invariant under invertible changes of the factorization.
  • Factorized implementation: Factored representations keep storage and multiplications in thin factors, avoiding expensive operations on m × n matrices.The tangent iterates use rank-at-most-2r factors, and the matrix-sign computation reduces to a 2r × 2r problem.

R Factorized projection

The factorized implementation preserves the tangent update while keeping storage and multiplications in thin factors.

  • R Factorized projection: The factorization trick leaves the underlying tangent update unchanged while avoiding expensive operations on m × n matrices.Storage and multiplication remain in the thin factors.

4 Experiments

The experiments evaluate LoRA-TSD across multiple models, benchmarks, adapter ranks, and projection-step settings. LoRA-TSD generally provides strong accuracy and stability, with a quality–cost trade-off controlled by the number of inner iterations.

  • Setup: The benchmark suite covers commonsense reasoning and natural-language inference using Llama-3.2-1B, Llama-3.1-8B, and Qwen3-32B.The evaluated tasks include BoolQ, PIQA, SIQA, OBQA, QNLI, and MultiNLI.
  • Results: Across 12 Llama model-dataset settings, LoRA-TSD is best or tied among LoRA optimizers and has the highest average at both model scales.It also outperforms Muon Full-FT in 11 of 12 settings.
  • Results: On Qwen3-32B, LoRA-TSD and LoRA-Rite are the two strongest methods on SIQA and OBQA.The evaluation uses transferred 8B hyperparameters with minimal tuning.
  • Inner iterations: τ = 3 improves average accuracy by 0.92 points over τ = 1, while additional increases to τ = 4 or 5 provide no noticeable gain.The higher accuracy costs two additional inner projections per training step.
  • Rank ablation: LoRA-TSD maintains the most stable high-accuracy performance across adapter ranks on SIQA with Llama-3.1-8B.The rank ablation is reported in Figure 1.

5 Theory

The theory characterizes LoRA optimization through tangent-space stationarity and establishes convergence guarantees for LoRA-TSD and LoRA-Pro under stated smoothness, boundedness, nondegeneracy, and alignment assumptions. It also explains how gauge rebalancing supports the assumptions and how the Frobenius formulation connects LoRA-TSD to LoRA-Pro.

  • Convergence guarantees: LoRA-TSD has a stationary-point convergence guarantee for the nonconvex fixed-rank LoRA problem expressed through the tangent-projected gradient.The analysis directly bounds this measure rather than only analyzing standard factor-space gradient methods.
  • Assumptions: The analysis assumes a lower-bounded objective, bounded gradients, uniformly full-rank factors, and a uniformly aligned approximate tangent oracle.These conditions correspond to assumptions on objective behavior, gradient size, factor nondegeneracy, and descent alignment.
  • Assumptions: Gauge transformations preserve BkAk and Wk while rebalancing factor norms to address scale imbalance from factorization ambiguity.Empirically, the smallest singular values of Ak and Bk remain above a positive δ during training, as required by (A3).
  • Stationarity: The natural stationarity criterion for fixed-rank LoRA is a vanishing tangent-projected gradient, not a zero full gradient generally unattainable under rank constraints.This is the criterion used in the convergence bounds.
  • LoRA-TSD: Under assumptions (A1)–(A4) and the stated step-size conditions, the momentum-free LoRA-TSD iterates satisfy the theorem’s convergence guarantee.The result is stated for any N ≥ 1, with an additional implication when factor spectral norms are uniformly bounded.
  • LoRA-Pro: LoRA-Pro is the Frobenius-norm counterpart of LoRA-TSD, and the induced updates coincide when η = λ.This directional equivalence transfers the convergence analysis to LoRA-Pro with the Frobenius norm replacing the spectral norm.

6 Conclusion

LoRA-TSD performs Muon-style spectral descent in the LoRA tangent space using a factor-induced retraction, with convergence guarantees and strong benchmark performance.

  • 6 Conclusion: LoRA-TSD performs Muon-style spectral descent in the LoRA tangent space through a factor-induced retraction.The method updates standard LoRA factors without materializing full m × n matrices.
  • 6 Conclusion: LoRA-TSD converges to projected stationarity, while the Frobenius counterpart recovers LoRA-Pro with an analogous guarantee.The tangent-projected gradient serves as the relevant stationarity measure.
  • 6 Conclusion: Across six benchmarks and model scales from 1B to 32B, LoRA-TSD outperforms generic and LoRA-specific optimizers and remains stable across adapter ranks.The conclusion reports results on six benchmarks spanning Llama-3.2-1B, Llama-3.1-8B, and Qwen3-32B.
  • 6 Conclusion: The approximate inner oracle lacks an optimality guarantee, while controlled bounds and generative evaluation remain future work.This is the stated limitation and scope for future evaluation.

A.3 Proof of Proposition 2

The Frobenius-norm tangent surrogate reduces to the auxiliary problem solved by LoRA-Pro, yielding the same induced weight-space direction up to step scaling.

  • A.3 Proof of Proposition 2: The factor-gradient linear terms combine into the inner product between the full-space gradient and the factor-induced weight increment.Specifically, the increment is ΔB A + B ΔA.
  • A.3 Proof of Proposition 2: The Frobenius surrogate becomes precisely the auxiliary optimization problem solved by LoRA-Pro.The reduction follows after rewriting the linear terms and dropping constants.
  • A.3 Proof of Proposition 2: LoRA-Pro’s induced weight-space direction is the orthogonal projection of the full-space gradient onto the tangent space.The projection is written as BA = P_T(G_W).
  • A.3 Proof of Proposition 2: The LoRA-TSD and LoRA-Pro directions agree for all positive η and λ, with equal increments when η = λ.Their step magnitudes differ by the factor η/λ unless the learning-rate parameters coincide.

A.4 Setup for Theorem 1

The theorem setup uses tangent-space projections and alternating projections while controlling factor rank, projector perturbations, and low-rank update norms.

  • A.4 Setup for Theorem 1: The analysis assumes full-rank LoRA factors with σ_r(A_k), σ_r(B_k) ≥ δ > 0 and imposes a descent-margin condition on the inner output.These are assumptions (A3) and (A4).
  • A.4 Setup for Theorem 1: Row-space projector perturbations are controlled through principal angles and singular-value separation between a factor and its update.The setup invokes Weyl-type singular-value control and a subspace perturbation bound.

A.5 Proof of Theorem 1

The proof establishes convergence of momentum-free LoRA-TSD under smoothness, boundedness, rank, and step-size assumptions, then connects projected stationarity to factor gradients.

  • A.5 Proof of Theorem 1: Under assumptions (A1)–(A4) and the step-size condition, LoRA-TSD’s projected-gradient norm converges to zero.The theorem is stated for the momentum-free algorithm; the momentum variant differs by adding momentum.
  • A.5 Proof of Theorem 1: The descent proof treats constant and diminishing step sizes by summing a descent inequality and telescoping the objective values.The constant-step analysis optimizes the bound over η, while the diminishing-step analysis uses η_k → 0.
  • A.5 Proof of Theorem 1: The chain rule converts vanishing projected stationarity into vanishing factor-gradient norms for the canonical factor updates.The convergence guarantee therefore concerns both the tangent-space measure and the LoRA factor gradients.
  • A.5 Proof of Theorem 1: The proof requires Frobenius smoothness, gradient boundedness, full-rank factors, and a step-size restriction; LoRA-Pro does not require the descent-margin condition because its inner problem is exact.The LoRA-Pro assumptions are stated separately from those for LoRA-TSD.
  • A.5 Proof of Theorem 1: The explicit LoRA-Pro inner solution includes a free matrix X representing gauge ambiguity, while the canonical choice fixes the factor-update representative.Different representatives induce the same first-order weight direction but can change the second-order factor product.

B Retraction Analysis

LoRA-TSD reconstructs tangent-space updates through the LoRA factors, yielding a smooth fixed-rank retraction without full-matrix materialization. Unlike truncated-SVD retraction, its second-order correction is factor-induced and its cost is less sensitive to layer size.

  • Tangent-space construction: LoRA-TSD represents tangent directions as ΔB A + B ΔA on the fixed-rank manifold.
  • Factor-induced map: The factor-induced map reconstructs factor increments using pseudoinverses and updates the product through the resulting factor changes.
  • Retraction validity: When B and A have full rank, the factor-induced map is a smooth retraction at W = BA.
  • Comparison with truncated SVD: LoRA-TSD and truncated-SVD retractions agree to first order but differ at second order through their respective curvature terms.
  • Computational cost: The factor-induced retraction uses a single product ΔB ΔA, whereas truncated-SVD retraction requires an SVD whose cost grows with layer size.

C Experimental Details and Additional Results

The experiments evaluate LoRA optimizers across commonsense and natural-language-inference benchmarks, using fixed cross-dataset tuning protocols. LoRA-TSD produces balanced update spectra, performs strongly across ranks, and reaches its best reported average accuracy with three inner iterations.

  • Datasets: Evaluations cover six English commonsense and natural-language-inference benchmarks, plus SIQA and OBQA for Qwen3-32B.
  • Hyperparameter tuning: Learning rates are selected on BoolQ and reused across the remaining benchmarks without per-dataset tuning.
  • Rank ablation: LoRA-TSD achieves the best accuracy in 11 of 14 rank settings and remains within a narrow accuracy band as rank varies.
  • Update spectrum: With rank r = 16, the smallest singular-value ratio reaches about 1.2 on average, while per-step spectra remain tightly clustered.
Loading 2609.02734v1…