Efficient optimization with orthogonality constraint: a randomized Riemannian submanifold method

We sincerely appreciate the reviewer for acknowledging the strengths of our work. We also thank the reviewer for all the thoughtful comments and constructive feedback, which we address in details below. All revisions are highlighted in red in the revised manuscript.

1. (W1 + Q5) Parallelization of RSDM compared to existing methods.

Thank you for raising this point. We believe the use of term "parallel" may have caused some confusion. What we really mean is the following: Although the total complexity is the same as existing coordinate-based methods [3,4], RSDM converges in fewer iterations despite its higher per-iteration complexity. This characteristic makes RSDM more suited to modern hardware, such as GPUs, which are optimized for parallel processing. To this end, we have modified the wording in Section 1.1 by removing the "parallel" discussion.

To be more precise, let $n=p$ and coordinate based methods [3,4] requires $O(n^2\epsilon^{-2})$ iterations for convergence to $\epsilon$-stationary points and each iteration costs $O(n)$, which gives total complexity of $O(n^3\epsilon^{-2})$. In contrast, proposed RSDM (with permutation sampling) requires $O(n^2r^{-2} \epsilon^{-2})$ iterations and each iteration costs $O(nr^2)$, yielding the same total complexity. Therefore, existing methods require more iterations (sequentially), which limits their efficiency on modern hardware designed for parallel processing.

This claim can be validated from Figure 2(a) and Figure 6(a), where the runtime complexity of RCD [3,4] is significantly higher despite a lower per-iteration cost. On the contrary, the proposed RSDM is more favorable for large-scale settings when implemented on hardware such as GPUs.

2. (W2 + Q6) On the assumption of iterates staying in the local region.

Thank you for the question. We can get rid of such an assumption when $X^*$ is an isolated local minima in the neighborhood $U$. This is because RSDM converges to a stationary point almost surely (shown in Theorem 2). Therefore, once $X_{k_0} \in U$. Then we can show $X_k \in U$ for all $k \geq k_0$ by a quadratic growth argument as shown in [5]. We have modified Theorem 1 and 2 accordingly, along with the proof.

Nevertheless, in the stochastic setting where we fix the stepsize, the almost sure convergence to $X^*$ is lost. For such cases, we rewrite Theorem 3 by showing linear convergence only for points $X_{k_0}, ..., X_{k_1} \in U$. It is worth noticing that many prior works on stochastic Riemannian optimization directly assume all iterates are contained in $U$, e.g., Theorem 3 in [6], Theorem 4.6 in [7].

3. (Q1) Clarify the statements on per-iteration cost and non-standard linear algebra.

Thank you for the comment. Indeed, the main aim of the paper is to reduce the complexity of non-standard operations associated with retraction, which is the main bottleneck for optimization with orthogonality constraints (see e.g., [8]).

Apart from the retraction, the per-iteration cost also includes gradient and iterate computation via matrix multiplication. Although these costs require higher worst-case complexities, they are less of concern on modern hardware such as GPUs, where they can be easily accelerated.

4. (Q2) Not suitable to use "denote...as...". The meaning of $X^\star$ in line 155.

Thank you for the suggestion. We have now corrected the wording involving "denote...as...". We have also clarified the meaning of $X^*$ as any local minimizer.

5. (Q3) Write grad F_k(I_n) with Euclidean gradient.

Thank you for the suggestion. We have now added the expression explicitly in terms of Euclidean gradient.

6. (Q4) Why gradient computation in line 222 require $O(nr^2)$?

This is because the cost of truncated permutation, required for forming $A = P_k(r) \nabla F(X_k)$ and $B = P_k(r) X_k$ is smaller than the matrix multiplication of $AB^T$, which costs $O(nr^2)$.

7. (Q7) Line 484 should be Figure 4 instead of Figure 2.

Thank you for noticing the typo. We have now fixed it.

8. (Q7) Comparison to RSSM [9].

Although we had briefly compared proposed RSDM to RSSM [9] in Section 1.1, we have now included a dedicated section in Appendix E of the revised manuscript to provide a detailed and comprehensive comparison. We also include a numerical comparison despite that RSSM [9] does not report experimental results or provide the code.

Specifically, we show that RSDM requires a lower retraction cost of $O(r^3)$ while RSSM [9] requires $O(nr^2)$. In addition, RSDM allows easy extension to quotient manifolds, such as Grassmann manifold, while RSSM [9] appears challenging to adapt to such manifolds.

Our experimental results clearly demonstrate that the proposed RSDM converges significantly faster than RSSM in terms of runtime, further validating the advantages of RSDM over RSSM.

We are immensely grateful for the reviewer’s positive feedback on our contributions. We also deeply appreciate the insightful comments and questions, which we address in detail below. All revisions are highlighted in red in the revised manuscript.

---

1. (W1 + Q3) Contribution limited to $p = \Omega(n)$ while many applications involve $p \ll n$. Experiments across various $p$.

Thank you for the comments and questions on $p$. First, the $O(n^3 \epsilon^{-2})$ complexity of the proposed RSDM represents the worst-case complexity, which may not reflect its empirical runtime performance. While the per-iteration cost of RSDM is $O(nr^2)$, it is dominated by matrix multiplication (as shown in Section 4), which can benefit significantly from modern hardware (such as GPUs). In contrast, the per-iteration complexity of RGD includes $O(np^2)$ for retraction, which involves non-standard linear algebra operations that cannot be easily accelerated. As shown in Section 7, RSDM achieves faster convergence in runtime than RGD across various settings, despite the complexity analysis. This suggests the theoretical complexity could be tightened with a refined analysis, which we leave as future work.

Additionally, we have conducted further experiments for PCA problem under various settings of $p = 100, 300, 500, 800, 1000, 1200, 1500, 1800$, with $n = 2000$ fixed. The results are given in Appendix D of the revised manuscript. We observe that RSDM shows superior convergence than RGD across all settings except for $p = 100$. This empirically validates the advantages of RSDM over RGD over a range of $p$ (even when $p$ is small relative to $n$). However, when $p$ is significantly smaller, the cost of retraction becomes negligible, which can lead to RGD outperforming RSDM in such cases.

In our future work, we plan to enhance RSDM for cases where $p$ is significantly small, e.g., by developing an adaptive submanifold selection scheme that allows RSDM to behave similarly to RGD when $p$ approaches zero, while maintaining its superiority for larger $p$.

---

2. (Q1) Comparison to RSSM [1].

Thank you for the question. Although we have briefly compared RSDM to RSSM in Lines 89-92, we have now included a dedicated section (Appendix E of the revised manuscript) where we detail the updates for RSSM and compare it to RSDM. For your convenience, we create the following table summarizing the per-iteration complexity of RSDM and RSSM [1].

From the table, it is clear that RSDM (with permutation sampling) requires significantly lower per-iteration complexity than RSSM, especially in terms of the cost of retraction, which requires non-standard linear algebra operations and cannot be easily accelerated compared to matrix multiplication.

Although both RSDM and RSSM [1] share the same goal of reducing the retraction cost, they employ largely different strategies. One of the major differences is that RSDM operates on the rows while RSSM [1] operates on the columns. This distinction provides RSDM with additional advantages, including easy generalization to quotient manifolds, such as the Grassmann manifold (as discussed in Section 6). In contrast, this generalization appears challenging for RSSM.

Additionally, while RSSM [1] did not include any numerical experiment nor provide the code, we have implemented and compared RSDM to RSSM numerically in Appendix E of the revised manuscript. The results show that RSDM achieves significantly faster convergence compared to RSSM [1], thus further validating the benefits of proposed RSDM.

---

3. (Q2) On the extension of Theorem 3 to mini-batch gradient descent.

Thank you for the comment. We have now extended Theorem 3 to the mini-batch gradient descent setting, in Appendix F of the revised manuscript.

In summary, we consider $F(X) = \frac{1}{N} \sum_{i=1}^N f_i(X)$ where each component function $f_i$ has bounded variance $\sigma^2$ and the mini-batch of size $B$ is sampled uniformly with replacement. We prove that using a mini-batch reduces the gradient variance from $\sigma^2$ (as in the stochastic case) to $\sigma^2/B$. Consequently, the convergence rate remains the same as in the stochastic case, with the variance term $\sigma^2$ replaced by $\sigma^2/B$ for mini-batch setting.

We sincerely thank the reviewer for acknowledging the strengths of our work and for providing constructive comments and questions. Below, we provide detailed responses to each of the points raised. All revisions are highlighted in red in the revised manuscript.

1. (W1 + Q3) On the possibility of exact convergence.

Yes. We have now (in Appendix C of the revised version) derived an exact convergence of our method by adding a condition on the sampling. In particular, we design a two-loop procedure for RSDM where we construct a set of $P_k^s(r), s=0,...,S-1$, for each outer iteration $k$. Then as long as the sampling of $P_k^s(r)$ satisfy some non-degeneracy condition regarding the gradient norm (eq. (14)). Then we can show the exact convergence rate of $O(1/k)$.

We also give one such example of permutation sampling that satisfies this condition, i.e., by ensuring a non-overlapping over the set of (unique) truncated permutation matrix (details in Appendix C).

These results confirm the possibility of exact convergence of our proposed RSDM.

2. (W2) Total complexity shows no improvement compared to Riemannian gradient descent.

We highlight that all subspace based methods (including coordinate descent, subspace gradient) cannot achieve improved total complexity over gradient descent for general nonconvex functions, either in the Euclidean space [1] or on Riemannian manifolds [3,4].

However, this does not de-merit subspace based methods because when $n,p$ becomes extremely large, our proposed RSDM (with permutation sampling) can make progress once $O(nr^2)$ operations, whereas Riemannian gradient descent (RGD) can only make progress after $O(np^2)$ operations.

It is also important to note that the complexities presented are worst-case complexities and may not reflect the actual runtime behavior. In fact, the per-iteration cost of RSDM, i.e., $O(nr^2)$ is dominated by matrix multiplication whose cost could be much lower when exploiting modern software, such as GPU. However, the cost of RGD, i.e., $O(np^2)$ is dominated by non-standard linear algebra operations (by retraction), which is not as easily accelerated. This suggests that in practice, the total complexity of RSDM could be much lower than RGD.

In our experiments, we have indeed demonstrated our proposed RSDM consistently requires lower runtime compared to Riemannian gradient descent on various problem instances.

3. (W2) "The algorithm only converges to a neighbourhood rather than a first-order stationary point"

This is in fact incorrect. As we have shown in Theorem 2 that RSDM under both permutation and orthogonal sampling converges to a first-order stationary point almost surely.

4. (W3 + Q4) Experiments on different $r$ and small $p$.

Regarding the choice of submanifold sizes $r$, we have already shown in Figure 3(a) on the performance of proposed RSDM with different $r$. We observe there exists a range of $r$ values where RSDM performs better than Riemannian gradient descent (RGD) method.

For settings with small $p$, we have now conducted additional experiments on PCA problem under a range of $p$ values: $p =100, 300, 500, 800, 1000, 1200, 1500, 1800$ with $n = 2000$ fixed. In addition, we have evaluated the performance of RSDM with different $r$ values for each $p$. The experiment results are presented in Appendix D of the revised manuscript.

The results demonstrate that indeed the performance gap of proposed RSDM over RGD decreases when $p$ becomes smaller, as predicted by our theoretical analysis. However, RSDM is still able to show faster convergence than RGD across all settings, except for the case when $p = 100$. This verifies the benefits of RSDM over RGD across a range of $p$. Nevertheless, when $p$ is significantly smaller, the cost of retraction becomes negligible, which can lead to RGD outperforming RSDM in such cases.

In our future work, we believe RSDM can be enhanced for scenarios where $p$ is significantly small, e.g., by designing an adaptive submanifold selection scheme such that RSDM behaves similarly to RGD when $p$ is small while retaining its advantages for larger $p$.

5. (Q1) How to ensure $X_k$ remains inside $U$ under $X_{k_0} \in U$?

Thank you for your thoughtful question. In Theorem 1, we can ensure $X_k$ remains in $U$ almost surely when $X^*$ is an isolated local minima in $U$. This is because as we have shown in Theorem 2, RSDM converges almost surely in gradient norm. Thus, we can show $F(X_k)$ converges to $F(X^*)$ (once $X_k \in U$) and by quadratic growth condition (equivalent to PL condition for $C^2$ functions) in [2], we have $||X_k - X^* || \rightarrow 0$ almost surely. As a result, we can guarantee all $X_k \in U$ almost surely for $k \geq k_0$.

We have now modified the Theorem 1 and 2 in this regard. We have also added the proof accordingly in the Appendix.

We thank the reviewer for the thoughtful feedback and general appreciation of our work. Here are our detailed responses to the comments. All revisions are highlighted in red in the revised manuscript.

1. "Only the first $r$-columns of the matrix $P_k$ that define the subspace for the rotation."

We appreciate the reviewer for pointing this out. Indeed the random submanifold is only determined by the first $r$ rows of $P_k$. To address this, we have explicitly clarified this by specifying "the first $r$ rows of $P_k$" in line 173 of the revised manuscript.

2. "Enough to generate only a random orthogonal frame from St(n,r) that defines the subspace?"

This is correct. We have already addressed this in lines 209–211 for the uniform orthogonal sampling. Additionally, we have now clarified this in lines 218–219 for the permutation sampling, where only the first $r$ indices are relevant for defining the subspace.

3. On the local approximation of $\tilde F(Y)$.

As discussed in Section 3, we wish to minimize $\tilde F(Y) = F(U(Y) X_k)$ locally around $Y = I$ such that the updated $X_{k+1} = U(Y) X_k$ is not far from $X_k = U(I) X_k$. Due to the proximity of $Y$ to $I$, minimizing the local approximation of $\tilde F(Y)$ is natural, which gives rise to the Riemannian gradient update of $Y$.

Based on Riemannian gradient update, we have rigorously derived convergence in expectation (Theorem 1), in high probability and almost surely (Theorem 2). These results justify that minimizing the local approximation is both reasonable and effective.

4. Mistake/typo after the first equality in Eq. 4

After carefully reviewing the equation, we have not identified any mistakes or typos to the best of our knowledge. We would be greatly appreciated if the reviewer could be more specific on the potential mistakes/typos so that we can address and provide clarification.

5. On the order of the first $r$ indices of permutation.

The order of the indices does not matter and one can sample $r$ indices without replacement. This is exactly what we have implemented in the experiments. We have now added this remark in line 220 of the revised version.

6. Comparison to random subspace sketching [1].

Thank you for pointing out the connection to random subspace sketching method. However, the reference [1] is not relevant to our setting to the best of our knowledge.

In their paper (Page 3, second paragraph), the setting is optimization over generalized Stiefel manifold with constraint $X^T B X = I$, where $B = Z^T Z + \lambda I$ with $X \in R^{d \times p}$, $Z \in R^{n \times d}$ with $n >> p$. The aim of [1] is to develop a sketching method to reduce complexity dependency on $n$ and improve the conditioning of the optimization problem. However, we focus on when $B = I$ and to reduce the retraction complexity of $O(dp^2)$. Hence, their method is not directly applicable to our setting.

We have now included a discussion to this work in the Relate work section (Section 1.1).

7. Experiments on the choice of $r$. Missing information of $r$ for procrustes problem.

We have already provided an experiment comparing the different choices of $r$ for the PCA problem. See Figure 3(a), where we compare various choices of $r$ for both orthogonal and permutation sampling. We have noticed that there exists a range of $r$ that can lead to improvement over the Riemannian gradient descent method. In addition, we have performed additional experiments on varying both $r$ and $p$ for the PCA problem. The results are included in Appendix D of the revised version.

Regarding the choice of $r$ for Procrustes problem, we set $r = 150$ for the small-size problem and $r = 900$ for the large-size problem. We have now added this information in the revised version.

8. Discussion on the computational cost on Haar uniform sampling on $O(n)$.

As we have discussed in Section 4, it is not necessary to sample the entire orthogonal matrix from $\mathcal{O}(n)$. We only need to sample the first $r$ rows of the orthogonal matrix, which can be efficiently achieved using QR decomposition (via the Gram-Schmidt process) on a randomly generated Gaussian matrix. The computational cost of this sampling process is $O(nr^2)$.

In our experiments, we observe that despite the higher computational cost of orthogonal sampling compared to permutation sampling, RSDM with orthogonal sampling can still achieve better performance than Riemannian gradient descent.

9. "Use fast random subsampling methods, such as random fast JL transform?"

Thank you for the comment. We believe random fast JL transform is not directly applicable in our setting due to the orthogonality constraint we address in this paper. Nevertheless, we believe this is an interesting direction to explore in the future.

---

[1] Shustin, B., & Avron, H. (2024). Faster Randomized Methods for Orthogonality Constrained Problems. JMLR, 25(257), 1-59.

Thank you for clarifying this. I really like this paper's idea and also authors' response to reviewers' comments. I decided the increase the score of my review.

I also fully agree with the authors reply to the reviewer xyJ5 that the case of $n=p$ is an important one and that the argument that the retractions are fast enough when $p\ll n$ is subjective since ``fewer columns'' is a subjective statement. One could take this argument to an extreme by considering $p=1$, when the retraction corresponds to just a norm normalisation of a vector, but, even though it is very fast to compute, this is not where the useful structure of the orthogonality is employed.