Efficient Optimization with Orthogonality Constraint: a Randomized Riemannian Submanifold Method

We sincerely thank the reviewer for the positive evaluation of our work. We are glad that the reviewer found our work to be interesting and original.

1. Whether the proposed idea requires a compatible Riemannian metric to effectively reduce the computational cost.

We thank the reviewer for the question. In our work, we choose the Euclidean metric as the Riemannian metric for the Stiefel manifold. Under such a metric, we have derived the expression for the Riemannian gradient on the submanifold, which can be efficiently computed. We agree that using alternative metric, such as canonical metric may be interesting, which we leave for future exploration.

We sincerely appreciate the reviewer's thoughtful feedback and overall positive assessment on our work. We would like to take the chance to address the comments in detail.

1. On the convergence of MNIST/CIFAR experiment.

R1: In our MNIST/CIFAR experiments, our aim was to highlight the efficiency of the proposed method in terms of convergence speed, which aligns with our goal of improving optimization efficiency under orthogonality constraints. Faster convergence is particularly relevant in scenarios where reaching a certain level of accuracy is sufficient or where computational resources are limited.

To ensure a fair comparison, we used a fixed learning rate schedule and tuned both RSDM and RGD for their best performance. This setup is consistent with our theoretical developments, and the experiments serve to demonstrate the practical potential of RSDM in neural network training. We agree that exploring different learning rate schedules is valuable and leave this for future work.

To further evaluate the performance of RSDM in neural network training, we have now added experiment on training a Vision Transformer, where we observe that our method consistently achieves higher test accuracy throughout the entire training process (rather than only in early iterations). Please see R2 for more details.

2. Lack of large-scale experiments.

R2: Thank you for the suggestion. We have added an additional experiment on training (orthogonal) Vision Transformer. Following [1], we impose orthogonality constraint on the query, key and value projection matrices. We train a 6-layer, 4-head transformer with embedding dimension 1024 and 64 patches, on CIFAR10. This scales the number of parameters from 5.4M in Section 6.4 to 28.5M. We set $r = 300$ for RSDM and tune the stepsize for both RSDM and RGD. The experiment results are included on https://anonymous.4open.science/r/ICML2025-additional-experiments-08E5. We observe that RSDM converges faster than RGD in test accuracy with a non-negligible gap, throughout training process. This demonstrates the potential of RSDM to improve large-scale model training.

As our main focus is on the theoretical development of the framework, testing more diverse applications is beyond the scope of this paper. We plan to explore this direction in the future.

3. On the terminology of "non-standard linear algebra operations".

R3: Thank you for the question. In this paper, "non-standard linear algebra operations" refer to operations such as matrix decomposition, matrix inverse and matrix exponential. These operations generally have significantly higher computational complexity compared to standard operations like matrix multiplication or elementwise operations. We will clarify the definition in our revised version.

4. On the orthonormal layers in Sec 6.4.

R4: This is a typo. All the first five layers are orthonormal.

5. The number of inner iterations used in the experiments.

R5: In the experiments, we implement RSDM according to Algorithm 1, which does not have inner iterations.

6. Reasonable range for $r$.

R6: The suitable range of $r$ depends on the problems. As shown in Figure 3, RSDM demonstrates improved convergence over RGD across a wide range of $r$ values. Determining the optimal $r$ is indeed an important and valuable direction for future work.

7. On the RSDM with momentum.

R7: Thank you for the suggestion. The idea of combining RSDM with momentum is indeed worth exploring. There are several feasible solutions. One strategy is that we fix $P_k$ for several iterations where we apply momentum. This is equivalent to taking multiple gradient descent (with momentum) for minimizing $\widetilde F_k(Y)$, initialized from $I_r$. To test the viability of the proposed strategy, we evaluate RSDM with momentum and compare against RGD with momentum on the PCA problem. The result is included on https://anonymous.4open.science/r/ICML2025-additional-experiments-08E5. We see that the RSDM with momentum improves the convergence of RSDM, which supports the potential of incorporating momentum into our framework.

We hope our responses have addressed your comments. If there are any further questions or suggestions, we would be happy to address them.

We sincerely thank the reviewer for the insightful comments and feedback. We would like to address your comments as follows.

1. Discussions on relevant references.

R1: We thank the reviewer for highlighting references [1,2]. Reference [1] re-parameterizes variables in Euclidean space via the Lie exponential map. However, this approach requires differentiating through the exponential map, which can be computationally expensive. Moreover, the re-parameterization may alter the loss landscape, making convergence analysis more difficult. In contrast, our method comes with convergence guarantees, which [1] does not provide.

Reference [2] formulates the problem of training a Siamese network as an optimization problem on the Stiefel manifold and employs Riemannian (stochastic) gradient descent with retraction for parameter updates. Our proposed algorithm is task-agnostic, so it can be applied to improve the optimization in [2] as well.

We will include the above discussions in the revised manuscript.

[1] Cheap Orthogonal Constraints in Neural Networks: A Simple Parametrization of the Orthogonal and Unitary Group

[2] Siamese Networks: The Tale of Two Manifolds.

2. More practical and diverse datasets and tasks.

R2: We thank the reviewer for the comment. As noted by the reviewer, our primary focus is on theoretical developments, and we have followed prior works in using standard benchmarks. As suggested by the reviewer, we have now conducted additional experiments on (orthogonal) vision transformer. We followed [1] by imposing orthogonality constraint on the query, key and value projection matrices. We trained a 6-layer, 4-head transformer (embedding dimension 1024, 64 patches) on CIFAR10. We set $r = 300$ for RSDM and tune the stepsize for both RSDM and RGD. The experiment results are included on https://anonymous.4open.science/r/ICML2025-additional-experiments-08E5. We observe our proposed RSDM converges faster than RGD in both test accuracy and training loss. This validates the potential of RSDM in more practical settings. We plan to extend the evaluation to additional real-world tasks and datasets in future work.

[1] Fei et al. O-vit: Orthogonal vision transformer. arXiv:2201.12133.

We hope we have addressed your concerns. We would be happy to respond to any remaining questions you may have. Once again, we appreciate the reviewer's time and valuable comments.

We are grateful that the reviewer found our work interesting and well-written. Below are our responses to your comments.

1. On the novelty of this work compared to related works.

R1: Thank you for the comment. We would like to clarify the novelty of our work. We have carefully compared with related works both theoretically and numerically in the main paper and in Appendix I, J. Our method is well-motivated from Riemannian geometry, which gives a principled foundation for the design and convergence analysis of the algorithms within the framework of Riemannian optimization. This geometric perspective also allows our method to generalize naturally to more complex manifolds, including quotient manifolds, as shown in Section 5.4 and Appendix C. We believe this provides both theoretical and practical value. We will expand the discussion on the novelty and contributions of our work in the revised manuscript to make this clearer.

2. Add label to x-axes of Figure 3, 4, 5.

R2: Thank you for the question. We have already included the x-axes label for Figure 3,4,5, which is (wall-clock) Time.

We hope we have addressed all your comments and concerns. If you have further questions, we would be happy to address them.