StelLA: Subspace Learning in Low-rank Adaptation using Stiefel Manifold

Thanks for the in-depth review and constructive comments. We have addressed all the comments and suggestions in the revised version of the manuscript. Below we clarify the questions and concerns in detail.

Optimization methods that yield U,V from a Stiefel manifold is beneficial for certain downstream analyses, e.g. adversarial robustness [1,2]. These can be included in the introduction.

Thanks for the suggestion on potential applications. We will add these references to the introduction section in the revised version. Specifically, we will add the following sentences to L50:

"Besides, maintaining orthogonality of the low-rank matrices $U$ and $V$ during training is also beneficial for certain downstream analyses, e.g. adversarial robustness [1,2]."

Although the authors give empirical evidence for improved evaluation metrics for a given rank budget, it is not clear to me why the manifold constraint should yield this. There is evidence and analyses that the plain LoRA does not preserve first-order optimality on the manifold of rank r matrices, see [3]. [3] Constructs a method that also preserves orthogonality of U,V - as a consequence first order optimality conditions and convergence criteria on the rank r manifold is available. However in the proposed work, the authors consider the Stiefel manifolds of U and V separately, which gives a different view on optimality and convergence. This should be discussed.

Thanks for the pointer to the work of GeoLoRA [3]. we will add GeoLoRA to the related work section and discuss the differences between GeoLoRA and StelLA.

Regarding the convergence and optimality of StelLA, if we consider the tuple $(U, S, V)$ as a point on the product Riemannian manifold $\mathrm{St}(r, m) \times \mathbb{R}^{r \times r} \times \mathrm{St}(r, n)$, then the convergence can be analyzed using the Riemannian optimization theory, which is well-established in the literature. For example, we refer the readers to [a] and [b] for the convergence analysis of Riemannian SGD and Riemannian Adam optimization algorithms. Because of the product structure, the components $U$ and $V$ can be optimized separately, i.e., $U$ and $V$ are optimized on their own respective Stiefel manifolds, leading to our Alg. 1. This treatment is mathematically equivalent to optimizing on the product manifold directly, thus preserving the convergence properties. To avoid confusion, we would like to emphasize that all $U$s and $V$s in StelLA are optimized jointly in the same forward and backward pass, as described in Alg. 1.

Furthermore, we would like to clarify that the product Riemannian manifold of StelLA is not the same as the manifold of rank-r matrices in GeoLoRA because each rank-r matrix has infinitely many representations in the product manifold due to the equivalence relationship by orthogonal groups (L287-289). Note that this does not affect the optimality because an optimal solution in the product manifold also corresponds to an optimal rank-r matrix.

In the gradient scaling section: Why is the "learning speed" an issue?

Thanks for the question on the motivation of gradient scaling, we will clarify this detail in the main text of the revised version. Previous works such as LoRA+ [c] demonstrated that using the same learning rate for $A$ and $B$ matrices in LoRA does not allow efficient feature learning. They showed that setting different learning rates for $A$ and $B$ with a fixed ratio can improve performance. The core idea is to ensure that both $A$ and $B$ are efficiently updated. However, $U$ and $V$ in StelLA are constrained to be orthogonal matrices, while $A$ and $B$ in LoRA+ are unconstrained matrices, so the technique in LoRA+ cannot be directly applied to StelLA. Motivated by their insights, in StelLA, we ensure that both $U$ and $V$ are efficiently learned by balancing their effective learning speed via gradient scaling, achieving the same goal as LoRA+ but in a different way.

It would be interesting to compare the convergence plots of LoRA, the proposed method and full-rank training using the same optimizer (e.g. Adam) with the same learning rates and hyperparameters.

Thanks for the suggestion on the comparison. We will add the convergence plots of StelLA, LoRA, and full-rank training in the appendix of the revised version.

A simple test case could be compressing (not finetuning) the fully-connected layers (everything but key,query,value) of Vit-b/16 on Cifar100. This should train to convergence in about 500epochs, if you just use the pretrained imagenet weights, but converges slowly enough so that the difference between LoRa and your "proper" method should be visible. Finetuning of low-rank adapters may converge too fast to see the difference. Rank 50 should be a good initial guess.

Thanks for the suggestion. Unfortunately, we do not have experience in the experiment setting of compressing and didn't find reference code for benchmarking. Due to the limited time, we are unable to implement this experiment. We would like to kindly note that we have already conducted extensive experiments (and confirmed by Reviewer Qyxg) on the commonsense reasoning benchmark, math and code reasoning, image classification, and image generation benchmarks, which cover a wide range of tasks and models.

There are methods e.g. AdaLoRa or [3] that are rank adaptive. Can your method handle adaptiveness?

Thanks for the question regarding rank adaptiveness. Yes, StelLA is compatible with AdaLoRA with minimal modifications. In AdaLoRA, the adapter is represented as $U\Lambda V^T$, where $U$ and $V$ are orthonormal matrices and $\Lambda$ is a diagonal matrix. Values in $\Lambda$ are used to control the rank allocation of the adapter. In StelLA, we can also constrain the matrix $S$ to be diagonal by setting all off-diagonal elements of $S$ to be zero during training. With this modification, StelLA can be used with AdaLoRA to adaptively adjust the rank of the adapter during training. However, it is hard to apply the rank adaptation approach in [3] to StelLA, because [3] needs expanding the rank of the adapter during training, which is not supported in StelLA (there is no rank augmentation step in StelLA).

Can your method be extended to tensor-valued layers, e.g. using low-rank Tucker format, see e.g. [4]

Thanks for the question regarding the extension of StelLA to tensor-valued layers. Yes, StelLA can be extended to tensor-valued layers straightforwardly as follows. The Tucker decomposition of a rank-$(r_1,\dots r_d)$ tensor is written as $C\times_1 U_1\cdots \times_d U_d$, where $U_i$ are column-orthonormal matrices. Therefore, the tuple $(C, U_1, \dots, U_d)$ lies in the product Riemannian manifold of $\mathbb{R}^{r_1\times\cdots r_d}\times \mathrm{St}(r_1, n_1)\times \cdots \times \mathrm{St}(r_d, n_d)$. We can use the same approach as in StelLA to optimize the low-rank Tucker decomposition, because each component $U_i$ can be optimized in its own Stiefel manifold thanks to the product manifold structure.

• [a] Bonnabel S. Stochastic gradient descent on Riemannian manifolds. IEEE Transactions on Automatic Control, 2013.
• [b] Becigneul G, Ganea O E. Riemannian Adaptive Optimization Methods. ICLR 2019.
• [c] Hayou S, Ghosh N, Yu B. Lora+: Efficient low rank adaptation of large models. ICML 2024.

We would like to thank the reviewer for the constructive comments and suggestions. We have addressed all the comments and suggestions in the revised version of the manuscript. Below we clarify the questions and concerns in detail.

A single initialization strategy is used (no combinations with Pissa, etc). This might additionally boost the performance.

Thanks for the suggestion. We indeed have experimented with different initialization strategies, including the combination of StelLA with PiSSA and MiLoRA in the ablation studies. In Table 5 of the main paper, the row marked as "SVD-major" is in fact StelLA + PiSSA, and the row marked as "SVD-Minor" is StelLA + MiLoRA. We will clarify this matter in the revised version to avoid confusion.

For easier reference, we summarize the results in the table below.

Table: Comparison different initializations on the commonsense benchmark using LLaMA3-8B.

The results show that StelLA is robust to different initialization strategies. The subspace learning powered by Riemannian optimization can effectively learn the optimal subspace regardless of the initialization strategy (L312-313).

The method implies using additional computation overheads in the retraction process (which is common for all Riemannian approaches).

Thanks for raising the concern. The retraction step is indeed computationally expensive, but with speed optimization including both algorithmic and hardware acceleration techniques, we can mitigate this issue and let StelLA be trained with only 15% additional training time compared to LoRA. Please kindly refer to the appendix Sec. D and our response to Reviewer Qyxg for the detailed and thorough discussion of computational cost of StelLA.

The mentioned quotient structure allows for decreasing the ambiguity of the choice for U and V matrices, but according to the text such an approach "is not designed for low-rank adaptation". Could you clarify the details for the optimization method that you use in this case?

Thanks for the question, we will clarify this detail in the main text of the revised version. We use Alg. 1 to optimize the quotient structure, same as in StelLA. The difference between the quotient structure and StelLA is that the quotient structure is equipped with the Riemannian metric proposed in "R3MC: A Riemannian three-factor algorithm for low-rank matrix completion" [36]. This Riemannian metric is designed for the low-rank matrix completion problem, which is different from the low-rank adaptation (LoRA) problem. We provided the details of R3MC's Riemannian metric in the appendix Sec. B. Specifically, as noted in the appendix L34-39, the Riemannian metric is induced by the block approximation of the Hessian of the loss function in the low-rank matrix completion problem: $||USV^T-W||_F^2$. Hence, we say that it is not designed for low-rank adaptation, because the Hessian of the loss function in low-rank adaptation is different from that in low-rank matrix completion.

Nevertheless, the quotient structure can be optimized using Alg. 1, as long as the three equations for

• converting Euclidean gradients to Riemannian gradients (see [36, Eq. 9]),
• projecting a vector onto the tangent space (see [36, Eq. 5]), and
• retracting a point onto the manifold (see [36, Eq. 7]),

are computed accordingly for R3MC's Riemannian metric.

Did you try to use basic Riemannian optimization algorithms for $\mathrm{St}\, (r, m) \times \mathbb{R}^{r \times r} \times \mathrm{St}\, (r, n)$ and $(\mathrm{St}\, (r, m) \times \mathbb{R}^{r \times r} \times \mathrm{St}\, (r, n)) /(\mathcal{O}(r) \times \mathcal{O}(r))$ without mixing them with the Euclidean optimizers?

Yes, we have tried using the Riemannian Adam optimizer [6] on StelLA in the ablation studies. The results are shown in the last row of Table 5, where the parallel transport is enabled. The setting of StelLA + parallel transport is equivalent to using the basic Riemannian Adam optimizer on the product manifold $\mathrm{St}\, (r, m) \times \mathbb{R}^{r \times r} \times \mathrm{St}\, (r, n)$. We will make this information explicit in the discussion of Table 5 in the revised manuscript.

Thanks for the insightful review and constructive comments. We have addressed all the comments and suggestions in the revised version of the manuscript. Below we clarify the questions and concerns in detail.

The paper lacks a comparison of computational efficiency. While the proposed method introduces additional computation for manifold operations (e.g., projection and retraction), which may impact training speed. A runtime or efficiency analysis would strengthen the empirical evaluation.

Thanks for the suggestion. A detailed discussion of computational cost of Alg. 1 is provided in the appendix Sec. D. We will move this information from the appendix to the main text for better visibility.

Here we summarize the key information:

• The main computational overhead of Alg. 1 comes from the SVD operations in the polar retraction step, which can be reduced by using efficient CUDA-accelerated SVD implementation.
• Training the LLaMA3-8B model with StelLA using Alg. 1 takes 6.5h for the commonsense benchmark, which is approximately 1.4x the time of LoRA (4.5h) on one H100 GPU. It can be further improved, e.g., by increasing parallelism with batch processing (details in our response to Reviewer Qyxg), which reduces the training time from 6.5h to 5.2h, about only 15% slower than vanilla LoRA.

Please kindly refer to our response to Reviewer Qyxg for more details on the computational cost, training time and how to improve them.

The method is related to LoRA-XS [3]. However, the paper does not include experimental comparisons against LoRA-XS, weakening the completeness of the evaluation.

Thanks for the suggestion. Below we list the results on commonsense reasoning benchmark using LLaMA3-8B model from LoRA-XS's paper (they followed the same experiment setting as ours, so the results are directly comparable.) The result shows that StelLA outperforms LoRA-XS by a large margin. This is understandable because the number of learnable parameters in LoRA-XS is much smaller (only the $S$ matrix is learnable) than that in StelLA.

Table: Comparison of StelLA and LoRA-XS on the commonsense reasoning benchmark using LLaMA3-8B model.

It would be beneficial to compare the orthogonality of the proposed method to other methods like TriLoRA and LoRA-XS, as this would strengthen the effectiveness of the proposed Stiefel manifold constraints.

Thanks for the suggestion. We have briefly discussed the difference between StelLA and TriLoRA/LoRA-XS in the related work section (L94-97). Here we reiterate the key points and will add more details in the revised version:

• TriLoRA does not initialize the low-rank matrices with orthogonal matrices (they initialize U with Gaussian, and V with zeros), and it does not maintain orthogonality during training.
• LoRA-XS initializes the low-rank matrices with the leading $r$ singular vectors and the singular values of the weight matrix (Suppose $M=U_r S_r V_r^T$, then $U:=U_r$ and $V=S_r V_r^T$), in which only one of the low-rank matrices $U$ is orthogonal. Both $U$ and $V$ are frozen during training, which limits the expressiveness of the low-rank adaptation.

We summarize the usage of orthogonality in StelLA, TriLoRA, and LoRA-XS in the table below.

Table: Comparison of orthogonality in StelLA, TriLoRA, and LoRA-XS

The performance improvements on image classification benchmarks are limited. This raises questions about the robustness and generalizability of the proposed method

Thanks for raising the concern. The relatively limited performance improvements on image classification benchmarks are due to the fact that the performance might have being saturated on the benchmark: the accuracy of LoRA is above 90% on 6/8 of the datasets. Yet, the performance improvements of StelLA over LoRA are significant on the two less saturated datasets: DTD (79.2% -> 79.89%) and Sun397 (75.62% -> 76.28%). Despite facing the challenge in further improving performance, our method still outperforms the two SOTA baselines, DoRA and PiSSA.

Thanks for the detailed review and constructive comments. We have addressed all the comments and suggestions in the revised version of the manuscript. Below we clarify the questions and concerns in detail.

W1: Lack of comparison of computational overhead or training time.

and

W2: Will there be a major drawback to computational complexity? Will the training cost be significantly increased due to the involvement of complex operations?

and

Q1: Theoretically and practically, how much would the proposed method slow down the model training speed?

and

Q3: Is there potential for algorithmic or hardware-level optimizations to speed up Stella?

Thanks for the comments and questions regarding the computational overhead and training time, which is indeed an important aspect of our proposed method. We would like to point out that a detailed discussion of computational cost of Alg. 1 has already been provided in the appendix Sec. D. We will move this information from the appendix to the main text for better visibility.

We summarize the key information in the appendix here:

• The main computational overhead of Alg. 1 comes from the SVD operations in the polar retraction step, which can be reduced by using efficient CUDA-accelerated SVD implementation [e].
• Training the LLaMA3-8B model with StelLA using Alg. 1 takes 6.5h for the commonsense benchmark, which is approximately 1.4x the time of LoRA (4.5h) on one H100 GPU. It can be further improved, e.g., by increasing parallelism with batch processing (details below), which reduces the training time from 6.5h to 5.2h, about only 15% slower than vanilla LoRA.

Detailed discussion:

Theoretically, the computational complexity of the polar retraction step is $O(mn^2)$ for $m\times n$ shaped tall matrices ($m > n$) due to the SVD computation. But practically, there are efficient SVD algorithms for tall matrices that can reduce the computational cost significantly. Specifically, we used the gesvda solver [e] in StelLA, which is a CUDA-accelerated SVD implementation that can handle tall matrices efficiently. Below is a micro-benchmark of different SVD solvers for the matrix shapes used in StelLA for the commonsense benchmark on LLaMA3-8B. The results show that the gesvda solver is significantly faster than the default GPU solver in PyTorch.

Table: Runtime of different SVD solvers in PyTorch on H100 GPU

The speed could be further improved by using batch processing to increase parallelism. Specifically, we can stack all the low-rank matrices $U$ and $V$ with the same shape into a batch, and then perform the polar retraction step on the batch of matrices. This reduces the training time from 6.5h to 5.2h, about 15% slower than vanilla LoRA. Below is a micro-benchmark of the SVD using batch processing on the matrix shapes used in StelLA for the commonsense benchmark on LLaMA3-8B (there are 192, 64, and 64 matrices with shapes 4096x32, 1024x32, and 14336x32, respectively). The results show that the batched SVD is significantly faster than the sequential one, and the polar retraction step is no longer the bottleneck of the training time.

Table: Runtime of SVD with/without batch processing

W3: References should be in numerical form according to NeurIPS 2025 Paper Formatting Instructions.

Thanks for pointing this out; we have fixed the reference format in the revised manuscript.

W4: There is no discussion of limitations.

We have discussed the limitations in appendix Sec. G, which we will move to the main text for better visibility.

W5: Why maintain the orthogonality of subspaces? This approach lacks motivation.

and

Q4: I expect to get some intuitive explanations: Why does applying SVD decomposition with orthogonal constraints during training improve model performance? Could this operation be merely playing a similar role to simple regularization?

The columns of the orthogonal matrix form an orthonormal basis for the subspace, which represents the "direction" of the subspace and is decoupled with the "magnitude" of the representation. Maintaining orthogonality is to ensure that the learning focuses on optimizing the "direction" of the subspace rather than their "magnitudes". Specifically in StelLA, $V$ represents the "direction" of the input subspace, and $U$ represents the "direction" of the output subspace, and $S$ represents the "magnitude" of the adapter.

The separation of direction and magnitude is a useful regularization in many areas of machine learning, such as in weight normalization [a] (separates direction and magnitude of network weights), CosFace [b] (normalize both features and class representative vector to focus on learning directions), and Grassmann Class Representation [c] (represent class prototype by subspaces, which ignores the magnitude). Research in these areas shows that decoupling direction and magnitude can improve the performance of the model. Recently, DoRA [d] also shows that this separation also benefits low-rank adaptation. As such, we investigated whether explicitly optimizing the direction of subspaces during low-rank training can further improve LoRA's performance. Our empirical results confirm that preserving orthogonality throughout training can indeed lead to performance gains.

As reviewer CC7z pointed out, maintaining U and V on the Stiefel manifold is also beneficial for certain downstream analyses, e.g. adversarial robustness.

We will enhance the discussion of the motivation for maintaining orthogonality in the revised version of the manuscript.

Q2: By adding strict constraints during training, could the convergence speed become slow or even fail to converge? Convergence speed is already a pain point for LoRA, and modification in this direction will further reduce the convergence speed.

We didn't observe any convergence issues in our experiments. We used the same number of training epochs for both StelLA and all baselines. At the end of training, StelLA achieves better performance than other methods, which indicates that the added constraints do not hinder convergence but instead contribute to more effective optimization.

Generally speaking, literature on adding orthogonal constraints [c, f, g, h, i] in neural networks shows that the convergence speed is improved with the orthogonality constraints.

We will add comparison of convergence plots to the revised manuscript.

Q5: Why are the rank settings different in each scenario? The choices lack justification.

For each learning task, we carefully followed the same setting as in previous work for fair comparison. For example, for the commonsense reasoning benchmark, we followed the setting in DoRA (L201), and for the math and code generation benchmarks, we followed the setting in PiSSA (L224).

Difficult to determine whether the formulas in the article are accurate.

Thanks for raising this concern. We provided an illustrative example in the appendix Sec. A, where a two-dimensional case is used to explain the concepts related to Stiefel manifold and geometric optimization. The equations are simplified in two-dimensional space, hoping to help sanity checking the workflow of Alg. 1.

• [a] Weight normalization: A simple reparameterization to accelerate training of deep neural networks, NIPS 2016.
• [b] Cosface: Large margin cosine loss for deep face recognition, CVPR 2018.
• [c] Get the Best of Both Worlds: Improving Accuracy and Transferability by Grassmann Class Representation, ICCV 2023.
• [d] DoRA: Weight-Decomposed Low-Rank Adaptation, ICML 2024.
• [e] https://docs.pytorch.org/docs/stable/generated/torch.linalg.svd.html
• [f] Unitary Evolution Recurrent Neural Networks, ICML 2016.
• [g] All you need is beyond a good init: Exploring better solution for training extremely deep convolutional neural networks with orthonormality and modulation. CVPR 2017
• [h] Can we gain more from orthogonality regularizations in training deep networks? NeurIPS 2018.
• [i] Orthogonal Convolutional Neural Networks, CVPR 2020.