Retraction-free optimization over the Stiefel manifold with application to the LoRA fine-tuning

Thank you for appreciating our work and providing detailed comments. The following are our responses.

Weaknesses:

All the presented theory is developed on a function defined on $\text{St}(d, r)$, while LoRA fine-tuning gives rise to an objective function $f(B, A) = L(BA)$, where $B \in \text{St}(d, r)$ (or $\text{Ob}(d, r)$), and $A \in \mathbb{R}^{r \times m}$. This objective function has to be minimized over $\text{St}(d, r) \times \mathbb{R}^{r \times m}$, and the advantages of optimizing on a compact manifold are thus lost.

Unfortunately, this makes the presented theoretical results not directly useful for the practical case under consideration. To give a more precise statement, for example in Lemma 3, the constant $\hat{L}$ would depend on $A$. By the mean value theorem, we would get a bound of the kind$\lVert \text{grad}\_B f(A, B_1) - \text{grad}\_B f(A, B\_2) \rVert \leq C(A) \lVert B\_1 - B\_2 \rVert$ (as noted in equation after line 183 for the Euclidean gradient).

Since the space in which $A$ resides is not compact, one would need at least a uniform control on $\lVert A_k \rVert$ over the iterations to make the theory interesting for LoRA fine-tuning. It is interesting to note, however, that the authors observe exact numerical convergence to the constraint in all cases.

Reply: The variability of Lipschitz constants over an unbounded domain is not unique to our method. Consider a simple two-layer neural network with weight matrices $U$ and $V$. For a loss function $\mathcal{L}(U, V)$, it is natural that the Lipschitz constant of $\nabla_V \mathcal{L}$ depends on the norm of $U$, where the domain of $U$ spans the entire Euclidean space, thus being unbounded.

This phenomenon is common in various optimization scenarios and can be managed by introducing regularity assumptions on the loss function, such as cocoercivity and differentiability. Specifically, if the iterates of an algorithm adhere to a certain descent property on their loss functions, they will remain within a bounded level set from the cocoercivity of the loss function. We also note that there is a line of research that analyzes the convergence of algorithms by employing local Lipschitz smoothness instead of global Lipschitz smoothness.

Questions:

1. Concerning the comment made in the "weaknesses" section: In all the proofs, the Lipschitz constants $L_f$ and $D_f$ depend on $A$, making the convergence analysis not directly applicable in the numerical case under investigation. I would appreciate the authors' comment on this, maybe I am missing something here?

Reply: Gradient Lipschitz continuity is a widely used assumption in the analysis of optimization algorithms. Though this condition may seem restrictive for unbounded domains, particularly in unconstrained optimization, it can be relaxed by incorporating cocoercivity (which implies bounded level sets) and ensuring differentiability of the objective function. The key is how to ensure the iterates lie in a level set of $L$. This is easy for the deterministic setting, where the line search is an effective tool. However, this would be difficult in the stochastic setting where line search is not applicable. It should be noted that such issue is not specific to our algorithm, but all stochastic algorithms.

2. I believe however that the observed convergence behavior is not a coincidence, even if the setting is different. The issue can be addressed by assuming $L(X) \to +\infty$ as $\lVert X\rVert \to +\infty$. With enough regularity assumed, this ensures that the sublevel sets of $L$ are compact. Given an initial condition $X_0 = B_0 A_0$ in the compact set $L^{-1}((-\infty, M])$, the flow $\dot{A} = -\nabla_A L(BA)$ decreases $L$, thereby remaining within the initial compact set $L^{-1}((-\infty, M])$. This would provide an effective bound for the Lipschitz constant, without compactness of the original domain.

Reply: In our response to the weakness, we acknowledge that the variability of Lipschitz constants over an unbounded domain is a common issue in machine learning tasks, not just specific to our approach. To address this, one effective strategy is to assume that the loss function is proper, coercive, and smooth. See also our responses to the weakness and Q1.

3. I don't see however a way to control $\lVert A\rVert$ uniformly in time for "classification-like" problems, where the global interpolating minima may be off at infinity.

Reply: As in our response to previous questions, we may need to control the step size and the estimation accuracy of the gradient to ensure the boundedness of $\lVert A\rVert$ by maintaining iterations within a bounded level set. Could you please provide more details on your concerns regarding the potential blow-up of $\lVert A\rVert$?

4. Line 93: the definition of $f$ has a typo. Also, "differentiable" is more common to use.

Reply: Revised.

5. Line 131: Even if it's clear what you mean, $\bar{U}_{\text{St}(d, r)}\left(\frac{1}{8}\right)$ was defined nowhere in the main manuscript.

Reply: It is defined as $\bar{U}_{\rm St}(1/8) = \\{ Y \in \mathbb{R}^{n \times p} \mid \lVert Y - X \rVert_F \leq \frac{1}{8} \\}$. We have added this into the notation part.

6. The choice of notation is a bit unfortunate in some points. I would personally substitute the loss function $L$ with $\mathcal{L}$, just to avoid confusion with the Lipschitz constants.

Reply: Thank you for the suggestion. We will revise our text accordingly to enhance readability.

Thank you for reviewing our paper. The following are our responses.

Weaknesses:

1. Though the value of mu and an upper bound of $|x_0 - \bar{x}_0|$ are given concretely, the choice of step size is unknown. Theoretically, the step size needs to be sufficiently small (See Theorem 1). Any theoretical suggestion for the choice of the step size?

Reply: It is not specific to our algorithm for unknown step size, including retraction-based algorithms, where the step size relies on unknown Lipschitz constant of the Riemannian gradient for problem (1).

Theorem 1 demonstrates that a suitable constant step size results in convergence. In contrast, the landing algorithm in [1] converges only to a neighborhood whose size is dependent on the step size. According to Theorem 1, an exact step size can be calculated if $L$, $\hat{L}$, and $\hat{D}_f$ are known. If these parameters are unknown, it is necessary to estimate their upper bounds or to utilize numerical strategies, such as a grid search. The role of this theory is to ensure the existence of a constant step size that allows the algorithm to converge. In our numerical experiments, we employ the grid search.

2. Numerical experiments report results of the comparisons. However, the definition of result is not given. Is the result computational time or classification accuracy or a notion of correctness or something else?

Reply: We report the overall (matched and mismatched) accuracy for MNLI, Matthew’s correlation for CoLA, Pearson correlation for STS-B, and accuracy for other tasks. Higher is better for all metrics. Additionally, since the additional computational cost per step induced by our method is negligible, the number of epochs needed can be roughly equated to the time required.

3. Problem (12) does not remove all the ambiguity. Note that if $B \in St(d, r),$ then $B A = B O O^T A = \tilde{B}\tilde{A}$ where $O$ is an orthonormal matrix and $\tilde{B} = BO$ is still in $St(d, r)$. Likewise for $B \in Ob(d, r)$. Is it possible to completely remove the ambiguity by considering the quotient manifold?

Reply: The distinction between the Stiefel and Grassmann manifolds is not significant, as observed in applications from low-rank optimization (e.g., matrix completion and matrix decomposition) and principal component analysis. The primary reason for choosing the Stiefel manifold over the Grassmann manifold is the ease of analytical treatment. Numerically, the differences between these two approaches are negligible, especially since the projection operators on their respective tangent spaces are identical, a common scenario in low-rank optimization and principal component analysis.

4. Why is the numerical performance of Manifold-LoRA for using the Stiefel manifold and the oblique manifold in (12) different? The optimization problem is equivalent in the sense that the local minimizer/stationary point does not change.

Reply: Setting the constraint set of $B$ to either the Stiefel manifold or the oblique manifold in equation (12) leads to two distinct optimization problems, characterized by unique optimality conditions and stationary points. For the Stiefel manifold, the stationary points are defined by:

$$grad f(X) - X \text{sym}(X^T\nabla f(X)) + X(X^TX - I) = 0.$$

For the Oblique manifold, the condition is:

$$grad f(X) - X \text{diag}(\text{diag}(X^T\nabla f(X))) + X \text{diag}(X^TX - I) = 0.$$

Therefore, the sets of stationary points for the two manifolds are not equivalent. The differing geometrical constraints imposed by each manifold are expected to influence the numerical performance of the algorithm.

Questions:

The questions are given in the weaknesses section. More questions are given below.

1. L117, $X \in {\rm St}(d, r) ~ or ~ X \in \mathbb{R}^{d \times r}$

Reply: It should be $X \in \mathbb{R}^{d \times r}$. Revised it.

2. What is the percentage of the computational cost of the retraction in the overall computations?

Reply: Considering a LoRA layer represented by:

$$H = (W + BA)S,$$

where $H \in \mathbb{R}^{n \times k}$, $S \in \mathbb{R}^{n \times k}$, $W \in \mathbb{R}^{n \times n}$, $B \in \mathbb{R}^{n \times p}$, and $A \in \mathbb{R}^{p \times n}$, we define the gradient with respect to $H$ as $\mathcal{D}_H$. The gradients with respect to $B$ and $A$ are computed as $\mathcal{D}_H S^T A^T$ and $B^T \mathcal{D}_H S^T$ respectively. Thus, the computational cost for these gradients amounts to $\mathcal{O}(n^2k + n^2p) + \mathcal{O}(2pnk)$. Our method is retraction-free, adding only the computational tasks of projection and calculating the penalty gradient, both of which are $\mathcal{O}(np^2)$. This additional cost is negligible when $p$ is small.

Thank you for carefully reading our manuscript and appreciating our work. The following are our responses.

Novelty:

1. Citation [1] considers optimization with constraints on the orthogonal group (i.e., St(n,n)). It seems that the core idea on how to implement retraction free optimization already appears there. The authors mention this in "related work", and say that [1] does not discuss the $r<n$ case. Inspection of [1] reveals that this is not the entire story. In Sec 3.5 of [1] the case of $r<n$ is discussed briefly, and it is said the results can be extended for that case.

   Reply: Our method differs with [1] from the following aspects:

   • The construction of our landing algorithm is more straightforward. It involves the summation of the Riemannian gradient of the loss and the gradient of the penalty function of the Stiefel manifold. While the landing field in [1] shares a similar structure, its first term (as seen in Eq. (4) of [1]) is less interpretable. Our method may provide a clearer approach to design landing algorithms for general manifolds.
   • We explore the strong convexity-like property, specifically the restricted secant inequality, of the penalty problem and provide an explicit value for the penalty parameter, which is not given in [1]. Consequently, our landing algorithm only requires tuning the step size, whereas [1] necessitates careful selection of both the step size and the penalty parameter.
   • Our theorem demonstrates the exact convergence when using a constant step size. In contrast, the landing algorithm in [1] only converges to a neighborhood whose size depends on the step size. This issue is acknowledged in their paper, particularly in the paragraph following Proposition 10. Developing a theory that ensures exact convergence with a constant step size is crucial for making these methods competitive with retraction-based approaches.

   Although the landing algorithm in [1] can be extended to the case where $r<n$ in Section 3.5, our analysis introduces a novel penalty parameter-free approach. Additionally, our theoretical results on exact convergence with a constant step size are both new and significantly different from those presented in [1]. We will revise our manuscript to make the comparison clearly.

2. However, the authors of [1] are skeptical of the value of this, as they mention that there are fast retraction methods (i.e., Cayley) for the Stiefel manifold.

   Reply: Both [1] and [MP 2015] indicate that the computational cost of the Cayley transformation is $4nr^2 + \frac{40}{3}r^3$, which is more than twice the cost of our method at $2nr^2$ for any $r$. Additionally, performing a retraction on the Stiefel manifold involves a specific orthogonalization procedure that is challenging to scale and parallelize, such as the matrix inversion required in the Cayley transformation. In contrast, our landing algorithm can be executed efficiently using BLAS3 operations.

   [MP 2015] Bo Jiang, and Yu-Hong Dai. A framework of constraint preserving update schemes for optimization on Stiefel manifold. Mathematical Programming 153, no. 2 (2015): 535-575.

3. Setting the penalty parameter: the authors advocate that they give an explicit value for the penalty parameter. And indeed Theorem 1 sets the parameter $\mu$ to 1/3. However, I do not think the situation is so simple. The theorems have the additional assumption that the iterates starts close to the manifold (1/8). This is, of course, easy to achieve - just start on the manifold itself. However, for the proof to work shouldn't all iterates stay inside this bound? This necessitates for the other parameter (step size) to be small enough. And indeed, the theorem requires that the step size be small enough, and does not specify how small.

   Reply: Our results show that setting $\mu = \frac{1}{3}$and $0 < \alpha \leq \frac{1}{2c_1}$ with $c_1$ specified in Line 418 will yield convergence. To ensure that iterates remain within the designed neighborhood, the step size must not exceed $\frac{1}{24 \hat{D}_f}$, based on the derivations from Equation (9) and the norm of $\lVert\nabla f(X)\rVert$ over $\bar{U}\_{{\rm St}}(1/8)$. That is, if $\alpha \le \frac{1}{24 \hat{D}\_f}$, all iterates $X_k$ will stay within $\bar{U}\_{{\rm St}}(1/8)$ for all $k$. This step size bound is implied by the condition $\alpha \le \frac{1}{2c_1}$. We will revise our manuscript to clarify it.

For responses to the remaining questions, please refer to the general response.

Thank you for reviewing our paper. The following are our responses.

1. W1: Some experimental results are unclear and not well defined.

   Reply: We have revised the descriptions of numerical experiments accordingly. Specifically, we report the overall (matched and mismatched) accuracy for MNLI, Matthews correlation for CoLA, Pearson correlation for STS-B, and accuracy for other tasks. And for Table 2, we report Exact Match(EM) and F1 scores. Higher is better for all metrics.

2. W2: Lack of the discussion about limitations of your method.

   Reply: In our conclusion, we address the theoretical limitations of our paper. While the constraints of the Stiefel and Oblique manifolds contribute to advancements in fine-tuning large language models, this method is difficult to extend to other areas such as pre-training and alignment. Therefore, it is essential to generalize our approach to more comprehensive manifolds, such as the Grassmann manifold. In addition, in the standard setting, LoRA uses a fixed rank for all layers, which has been shown to yield sub-optimal performance, as demonstrated in paper [46]. They address this issue by modifying the model architecture. In our conclusion, we also acknowledge that we did not use the coefficient matrix $A$ to dynamically select the rank $r$, which may have impacted our fine-tuning performance.

3. W3: Some findings of the experiments are hard to understand.

   Reply: As shown in Table 1, our method outperforms the other baselines on average scores. Specifically, Sphere (r=16) achieves better performance than other baselines on 7 out of 9 tasks, with the exceptions of MNLI-MM and QQP. In Tables 2 and 3, our method performs consistently better than the baselines. We will make the statement about the experiments more precise.

Questions:

1. Q1: In Table 1, could you please explain the meaning of m/mm? As I know, it's not an estimation of the performance of NLP models. You also didn't define Acc and Mcc. Also, why not use Acc to evaluate the performance of all the datasets?

   Reply: MNLI-M and MNLI-MM represent two similar datasets for which the evaluation metric is accuracy. The MNLI dataset is a collection of sentence pairs annotated for textual entailment. MNLI-M is the matched split of the MNLI dataset, while MNLI-MM is the mismatched split.

   Mcc stands for Matthews Correlation Coefficient, used for evaluating CoLA. The Pearson correlation is used for STS-B, and accuracy is the metric for other tasks. These metrics are widely used and are considered standard in the field, such as papers [22], [15], and [25].

2. Q2: The finding "It can be seen that our method 211 is consistently superior to other baselines." doesn't match the results in Table 1. What makes you conclude this?

   Reply: The best results are in bold in Table 1, our method outperforms other methods in most datasets except in MNLI-MM and QQP compared to full fine-tuning. Also in the setting of Sphere(r=16), it achieves the best average scores among all methods, followed by Stiefel(r=8).

3. Q3: Some findings are not clearly represented. For example, "We conclude that the proposed 228 Manifold-LoRA method achieves a 2x speed-up in training epochs compared to AdamW, while 229 simultaneously improving model performance".

   Reply: In Figure 2(a), it is easy to see that our method requires only half an epoch to reach a training loss of 1, whereas LoRA requires 2 epochs to achieve the same result. Additionally, since the additional computational cost per step induced by our method is negligible, the number of epochs needed can be roughly equated to the time required. In Figure 2(b) and Figure 2(c), it is obvious that our method achieves a higher EM (Exact Match) and F1 score.

4. Q4: The two datasets in the QA task are wildly used, but also weak. Have you tried other datasets? For example, the HotpotQA dataset.

   Reply: Please refer to the following table for the results of HotpotQA dataset. Our method, Sphere $({r=16})$, performs the best.

   Methods	Params	HotpotQA (EM/F1)
   Full FT	184M	63.3 / 76.7
   Adapter${r=16}$	0.61M	60.2 / 74.2
   Bitfit	0.07M	57.2 / 71.6
   LoRA${r=8}$	1.33M	61.3 / 75.4
   LoRA${r=16}$	2.65M	61.5 / 75.4
   Sphere${r=8}$	1.33M	62.4 / 76.3
   Sphere${r=16}$	2.65M	63.4 / 76.9
   Stiefel${r=8}$	1.33M	61.6 / 75.4
   Stiefel${r=16}$	2.65M	61.4 / 75.4

   Table: Results with DeBERTaV3-base on HotpotQA. We report EM/F1. The best results in each setting are shown in bold.

Thank you for your valuable feedback.

For the Table 1, we will revise it to make it more clear.

We use MCC for CoLA for two main reasons:

1. MCC is a widely recognized metric in the context of LoRA-fine-tuning. It is considered a standard in LoRA fine-tuning due to the fact that the original LoRA paper[22](in their section 5.1 Table 2) reported MCC scores for the CoLA experiments in the GLUE benchmark. This is the primary reason.

2. MCC is advantageous in handling imbalanced datasets because it evaluates all elements of the confusion matrix—true positives, true negatives, false positives, and false negatives—providing an unbiased assessment that accuracy (ACC) may not offer in skewed data situations. Moreover, MCC's scale from -1 to +1 provides detailed information about model performance, distinguishing between perfect predictions, random guesses, and complete inaccuracies, unlike ACC's more limited 0 to 1 range.

We are, of course, willing to accommodate your suggestion by including the ACC scores for the CoLA dataset as well.

We will also refine the statement "Our method is consistently superior to other baselines" to be more rigorous based on the results as you suggested.