StelLA: Subspace Learning in Low-rank Adaptation using Stiefel Manifold
Zhizhong Li, Sina Sajadmanesh, Jingtao Li, Lingjuan Lyu
TL;DR
StelLA introduces a geometry-aware extension to Low-Rank Adaptation by factorizing the adapter as $\tilde{W} = W + \frac{\alpha}{r} U S V^\top$ with $U,V$ constrained on the Stiefel manifold, enabling joint subspace optimization during fine-tuning. By leveraging a modular Riemannian-optimization design and a three-factor representation, StelLA learns input/output subspaces while remaining compatible with existing optimizers, achieving state-of-the-art results across NLP, image classification, and text-to-image generation benchmarks. Across commonsense reasoning, math/code generation, and generative vision tasks, StelLA consistently outperforms strong LoRA variants, with notable gains in accuracy and FID reductions, while maintaining manageable parameter overhead and inference mergeability. The approach demonstrates that exploiting subspace geometry during training yields tangible benefits for parameter-efficient fine-tuning, and the modular implementation supports scalable adoption in large-scale models.
Abstract
Low-rank adaptation (LoRA) has been widely adopted as a parameter-efficient technique for fine-tuning large-scale pre-trained models. However, it still lags behind full fine-tuning in performance, partly due to its insufficient exploitation of the geometric structure underlying low-rank manifolds. In this paper, we propose a geometry-aware extension of LoRA that uses a three-factor decomposition $U\!SV^\top$. Analogous to the structure of singular value decomposition (SVD), it separates the adapter's input and output subspaces, $V$ and $U$, from the scaling factor $S$. Our method constrains $U$ and $V$ to lie on the Stiefel manifold, ensuring their orthonormality throughout the training. To optimize on the Stiefel manifold, we employ a flexible and modular geometric optimization design that converts any Euclidean optimizer to a Riemannian one. It enables efficient subspace learning while remaining compatible with existing fine-tuning pipelines. Empirical results across a wide range of downstream tasks, including commonsense reasoning, math and code generation, image classification, and image generation, demonstrate the superior performance of our approach against the recent state-of-the-art variants of LoRA. Code is available at https://github.com/SonyResearch/stella.