🤖 machine learning

Retraction-Free Optimization over the Stiefel Manifold for the LoRA Fine-Tuning

This paper proposes a retraction-free, penalty-parameter-free optimization algorithm with global convergence guarantees for the Stiefel manifold and applies it to develop Manifold-LoRA, a geometry-accelerated fine-tuning method for large language models that achieves efficient training and strong downstream performance.

Original authors: Yuan Zhang, Jiang Hu, Zhijian Lai, Lin Lin, Zaiwen Wen

Published 2026-07-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yuan Zhang, Jiang Hu, Zhijian Lai, Lin Lin, Zaiwen Wen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.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 how to speak a new language or solve a specific puzzle. The robot already knows a lot, but it's so huge that teaching it everything from scratch would take a lifetime and cost a fortune. So, instead of rewriting the whole robot's brain, you decide to just tweak a tiny, special part of it. This is the world of "fine-tuning" large language models.

But here's the tricky part: when you tweak these tiny parts, you have to follow very strict rules. Think of it like a dance. If you move your feet too far apart, you lose your balance and fall. In math, this "balance" is called staying on a specific shape called a "manifold." For a long time, the way computers kept this balance was like a clumsy dancer constantly checking a mirror, calculating exactly where to step, and then forcing their feet back into the perfect spot. This "mirror-checking" (called a retraction) is slow and computationally expensive, like trying to run a marathon while constantly stopping to tie your shoes.

Another group of dancers tried a different approach: they just danced freely and hoped they wouldn't fall too far, using a "penalty" system where they got a slap on the wrist if they stepped out of line. But this required a referee to constantly adjust how hard the slap was, which was annoying and hard to get right. The big question in this corner of computer science was: Can we teach the robot to dance perfectly without the slow mirror-checking and without needing a referee to constantly tweak the penalty?

This paper says, "Yes, we can." The authors, Yuan Zhang and their team, have developed a new way to train these AI models that skips the slow mirror-checking and doesn't need a referee to adjust the penalty. They call their new method Manifold-LoRA.

Here is how they did it. Instead of forcing the robot to stay perfectly balanced at every single step (which is slow), they let it take a few steps that might wobble it slightly off-balance. Then, they use a clever mathematical "magnet" that gently pulls it back to the perfect dance floor. The magic of their discovery is that they figured out exactly how strong this magnet needs to be. They proved that if you set the magnet to a specific, fixed strength (they found it works well at a value of 1/3), the robot will naturally glide back to the perfect spot on its own, without needing anyone to adjust the magnet's power.

They also realized that the way we usually teach these robots (using a method called LoRA) has a hidden redundancy, like having two people carry a heavy box when one could do it just as well. By treating the robot's adjustments as a dance on a specific geometric shape (the Stiefel manifold), they removed this extra weight. This means the robot can learn the same tasks with half the number of adjustable parts, or learn the same amount of parts twice as fast.

In their experiments, they tested this new dance routine on various tasks, from answering questions about reading comprehension to generating creative stories. The results were impressive. On several datasets, their method converged (learned the task) twice as fast as the standard methods. For example, on a dataset called SQuAD 2.0, their model reached the same level of performance in half the time. Even more surprisingly, they managed to get better results using only half the number of trainable parameters compared to the standard approach.

The authors didn't just guess this would work; they provided a rigorous mathematical proof showing that their method is guaranteed to land on the correct "dance floor" and find the best solution, whether they are using a small, fixed step size or a shrinking one. They showed that this approach is not only faster but also more stable, meaning the robot doesn't get confused or wobble as much as it learns.

In short, this paper offers a new, smoother way to teach giant AI models. By removing the need for slow, heavy calculations and unnecessary adjustments, it allows these models to learn faster and more efficiently. It's like teaching a dancer to find their balance by feeling the floor rather than constantly checking a mirror, resulting in a performance that is both quicker and more graceful.

Technical Summary: Retraction-Free Optimization over the Stiefel Manifold for LoRA Fine-Tuning

Problem Statement
Optimization over the Stiefel manifold ( $St(d, r) = \{X \in \mathbb{R}^{d \times r} : X^\top X = I\}$ ) is critical in machine learning but faces two primary challenges in existing methods. First, retraction-based algorithms (e.g., Riemannian gradient descent) require costly orthonormalization steps (such as QR decomposition or matrix inversion) at every iteration, which hinders scalability for large matrices. Second, retraction-free "landing" methods, which avoid orthonormalization by using penalty functions, typically rely on carefully tuned step sizes and penalty parameters, often requiring an unknown penalty parameter to ensure convergence.

Furthermore, in the context of Large Language Model (LLM) fine-tuning, Low-Rank Adaptation (LoRA) suffers from representation redundancy. The standard decomposition $\Delta W = BA$ is non-unique (i.e., $BA = (BC)(C^{-1}A)$ for any invertible $C$ ), meaning the optimization landscape contains redundant degrees of freedom that traditional methods fail to exploit, potentially limiting model performance.

Methodology
The authors propose a unified framework addressing both the optimization algorithm and the LoRA formulation:

Retraction-Free and Penalty-Parameter-Free Optimization:
The paper introduces an algorithm that directly "lands" on the Stiefel manifold without retraction operators. The update rule is:
$X_{k+1} = X_k - \alpha_k P_{T_{X_k}St(d,r)}(g_k) - \mu X_k(X_k^\top X_k - I)$
where $P_{T_{X_k}St(d,r)}$ is the projection onto the tangent space, $g_k$ is the stochastic gradient, and the second term is the gradient of a quadratic penalty function $\phi(X) = \frac{1}{4}\|X^\top X - I\|^2$ .
- Key Theoretical Insight: By leveraging the "proximal smoothness" of the Stiefel manifold and establishing a "restricted secant inequality" (RSI) for the penalty function, the authors prove that a specific, explicit penalty parameter ( $\mu = 1/3$ ) is sufficient to guarantee convergence. This eliminates the need for tuning $\mu$ .
- Step Size Strategy: The method employs a two-scale step size strategy: a decaying step size for the loss gradient and a constant step size for the penalty gradient. This is motivated by decentralized optimization literature.
Manifold-LoRA Formulation:
To address LoRA redundancy, the authors reformulate the fine-tuning problem as an optimization over the product of a Stiefel (or Oblique) manifold and a Euclidean space. Here, $B$ is constrained to the manifold (serving as the basis), and $A$ remains in Euclidean space (serving as coordinates).
- The algorithm, Manifold-LoRA, applies the retraction-free update to $B$ using projected gradients and standard updates to $A$ .
- It incorporates an adaptive step size strategy based on the $\ell_2$ norms of $A$ and $B$ to account for their different Lipschitz constants.

Key Contributions

Explicit Penalty Parameter: The paper proves the existence of an explicit penalty parameter ( $\mu = 1/3$ ) that ensures convergence, removing the need for heuristic tuning required by prior landing methods.
Convergence Guarantees: The authors establish global convergence with the best-known iteration complexities:
- Deterministic setting: Converges exactly to a stationary point with a constant step size (improving upon prior results that only guaranteed convergence to a neighborhood). Complexity is $O(1/K)$ .
- Stochastic setting: Achieves $O(\log K / \sqrt{K})$ complexity for the loss gradient and $O(\log K / K)$ for constraint violation, outperforming previous single-scale step size approaches.
Manifold-LoRA Algorithm: A new fine-tuning framework that exploits manifold geometry to eliminate LoRA redundancy. It uses the proposed landing technique with negligible computational overhead compared to standard methods.
Computational Efficiency: Unlike retraction-based methods requiring $O(dr^2 + r^3)$ operations for orthonormalization, the proposed method relies on scalable BLAS3 operations with $O(dr^2)$ complexity.

Experimental Results
The authors evaluated Manifold-LoRA on Natural Language Understanding (GLUE), Question Answering (SQuAD v1.1/v2.0), and Natural Language Generation (E2E) tasks, as well as scaling experiments on LLaMA models.

Convergence Speed: On datasets like SQuADv2.0 and CoLA, Manifold-LoRA achieved approximately 2x faster convergence in terms of training epochs compared to vanilla LoRA.
Performance: The method consistently outperformed vanilla LoRA and other Parameter-Efficient Fine-Tuning (PEFT) baselines (Adapter, BitFit). Notably, Manifold-LoRA with half the trainable parameters (rank $r=8$ ) often matched or exceeded the performance of vanilla LoRA with rank $r=16$ .
Scaling: In scaling experiments from 1B to 8B parameter models (LLaMA3), Manifold-LoRA maintained performance advantages over baselines on reading comprehension (SQuAD, QuAC) and mathematical reasoning (GSM8K, MATH) tasks.
Manifold Consistency: Visualizations of $B^\top B$ confirmed that the learned matrices adhered closely to the target manifold geometry.

Significance and Claims
The paper claims that its primary significance lies in bridging the gap between theoretical manifold optimization and practical LLM fine-tuning. By providing a retraction-free, penalty-parameter-free algorithm with rigorous convergence guarantees, it removes the computational bottleneck of orthonormalization and the tuning burden of penalty parameters.

The authors assert that reformulating LoRA as a manifold optimization problem effectively addresses the inherent redundancy in low-rank decompositions. This geometric approach not only accelerates training but also enhances downstream performance, suggesting that leveraging manifold constraints is a viable path for more efficient and effective parameter-efficient fine-tuning. The work positions itself as a step toward scalable, geometry-aware optimization for large-scale machine learning, while acknowledging that future work is needed to extend these theories to general manifolds and to systematically study the generalization behavior of such constrained adaptations.

Technical Summary: Retraction-Free Optimization over the Stiefel Manifold for LoRA Fine-Tuning

More like this