Efficient Optimization with Orthogonality Constraint: a Randomized Riemannian Submanifold Method
Authors: Andi Han
Pierre-Louis Poirion
Akiko Takeda
Venue: ICML 2025 poster (2025)
Keywords: Riemannian manifold
Stiefel manifold
Efficient optimization
randomized submanifold
Primary area: optimization->nonconvex

Abstract: Optimization with orthogonality constraints frequently arises in various fields such as machine learning. Riemannian optimization offers a powerful framework for solving these problems by equipping the constraint set with a Riemannian manifold structure and performing optimization intrinsically on the manifold. This approach typically involves computing a search direction in the tangent space and updating variables via a retraction operation. However, as the size of the variables increases, the computational cost of the retraction can become prohibitively high, limiting the applicability of Riemannian optimization to large-scale problems.  To address this challenge and enhance scalability, we propose a novel approach that restricts each update on a random submanifold, thereby significantly reducing the per-iteration complexity. We introduce two sampling strategies for selecting the random submanifolds and theoretically analyze the convergence of the proposed methods. We provide convergence results for general nonconvex functions and functions that satisfy Riemannian Polyak–Łojasiewicz condition as well as for stochastic optimization settings. Additionally, we demonstrate how our approach can be generalized to quotient manifolds derived from the orthogonal manifold.  Extensive experiments verify the benefits of the proposed method, across a wide variety of problems.

Efficient Optimization with Orthogonality Constraint: a Randomized Riemannian Submanifold Method
Andi Han 12
Pierre-Louis Poirion Akiko Takeda 13
Abstract
Figure 1: Proposed random sub- manifold method on 2-sphere. Each iteration restricts the update I-dimensional randomly se- lected submanifold, i.e., circle.
Optimization with orthogonality constraints fre- quently arises in various fields such as machine learning. Riemannian optimization offers a pow- erful framework for solving these problems by equipping the constraint set with a Riemannian manifold structure and performing optimization intrinsically on the manifold. This approach typ- ically involves computing search direction in the tangent space and updating variables via a re- traction operation. However, as the size of the variables increases, the computational cost of the retraction can become prohibitively high, limiting the applicability Riemannian optimization to large-scale problems. To address this challenge and enhance scalability, propose novel ap- proach that restricts each update on random submanifold, thereby significantly reducing the per-iteration complexity. We introduce two sam- pling strategies for selecting the random submani- folds and theoretically analyze the convergence of the proposed methods. We provide convergence results for general nonconvex functions and func- tions that satisfy Riemannian Dolyak-Lojasiewal condition as well as for stochastic optimization settings. Additionally, we demonstrate how our approach can be generalized to quotient manifolds derived from the orthogonal manifold. Extensive experiments verify the benefits of the proposed method, across wide variety of problems.
where the matrix variable X E Rxx with is column orthonormal and F Ruxp - IR. Optimization with orthog- onality constraint arises naturally in various domains of ap- plications because it is crucial for achieving certain desired properties, such as linear independence, numerical stability and geometry preserving. Examples of applications include principal component analysis (Hotelling, 1933), indepen- dent component analysis (Theis al., 2009), multi-view clustering (Khan & Maji, 2021: Liu et al., 2021; Chen et al., 2022), low-rank matrix completion (Vandereycken, 2013; Mishra et al., 2014), robust optimal transport (Lin et al., 2020; Huang et al., 2021). training of deep neural networks (Helfrich et al., 2018; Li et al., 2019; Wang et al., 2020), continual learning (Chaudhry al., 2020) and fine-tuning large foundation models (Qiu al., 2023; Liu et al., 2024), among many others.
Riemannian optimization (Absil et al., 2008; Boumal, 2023; Han et al., 2024b) provides a powerful framework for solv- ing (1) by leveraging the geometry of the orthogonality constraint. Indeed, the set of orthogonal constraint forms a smooth manifold known as the Stiefel manifold, denoted by (t. = (X E RRxP X X Tp)- By equipping the manifold with a suitable Riemannian metric, optimization can be performed intrinsically on the manifold. A crucial step in this process is the retraction operation, which ensures that iterates remain on the manifold after each update. Vari- ous retractions have been proposed for the Stiefel manifold, such as those based on QR factorization (Absil et al., 2008). polar decomposition (Absil & Malick, 2012), the Cayley transform (Wen & Yin, 2013), and the matrix exponential (Edelman et al., 1998). All these retractions require non- standard linear algebra operations with complexity of at least (npp²) (see Section 2 for details). As result, the retraction step becomes the primary bottleneck for Rieman- nian optimization solvers as and p increase.
1. Introduction
In this paper, consider optimization problems with thogonality constraint, i.c.,
min XEReXP:XT X=I,
F(X)
(1)
'RIKEN AIP UUniversity of Sydney University of Tokyo, Correspondence to: Andi Han handi.hanee wwwnn@yyyddy......o
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Efficient Optimization with Orthogonality Constraint
In this work, we propose novel approach that updates the variable a random submanifold. In particular, our contributions are summarized as follows.
operations that cost O(nr²), where our method scales with ( As we elucidate the differences to these works in Section 4, our proposed submanifold update includes the coordinate descent as a special case, yet being more efficient in runtime empirically. Another line of research, including (Gao et al., 2019; Xiao et al., 2022; Ablin & Peyré, 2022; Ablin et al., 2023) develop infeasible methods for solving (I), where the updates do not necessarily satisfy the orthog- onality constraint. A recent work (Shustin & Avron, 2024) proposes randomized sketching method on the generalized Stiefel manifold with constraint X BX = Tp. However, they assume B = Z Z, with Z E gdx= with d > n. The aim is to reduce complexity in constructing B and im- prove the conditioning of optimization, which is different to our setting where B I, and the aim is reduce the complexity related to retraction.
We propose a novel parameterization for the update via the action of orthogonal group. This allows to update the current iterate in a random submanifold of orthogo- nal group. This reduces the complexity of non-standard linear algebra operations (such matrix decomposition, matrix inverse, and matrix exponential) from (np²) to (rr)), where r is the dimension of the submanifold se- lected.
We introduce two strategies for the parameterization, through permutation and orthogonal transformation. We derive the convergence results both in expectation and in high probability. We show the trade-off between the two in terms of efficiency and convergence guarantees. We show the other standard computations are reduced from (mp²) (nr²) or (nprr) under permutation and orthogonal sampling respectively.
2. Preliminaries
Stiefel manifold St(n,p) = {X E REXP: : = 1p} is the set of column orthonormal matrices. When n = p. St p) = (n)) which is called orthogonal manifold, also forming a group. We use St,,(n,)) p) to denote the tangent space X St(n,, p) and consider the Euclidean metric as the Riemannian metric, i.e., for any U,V V E St p), (U, V)x = (U, V)), we use (1.-) to represent the Euclidean inner product. For a smooth function F: Su(,, - IR, the Riemannian gradient gX) E StSt(n, p) is derived as gXadf(X) = V/(X) - X X F is the Euclidean gradient and {A]s Retraction Retrx St. - St(n, is a smooth map that allows to update iterate following tangent vector direction. Many retractions are proposed on Stiefel manifold, including (I) QR-based retraction: Retrx (U) = qf(X + U). where of extracts the Q- factor from the QR decomposition; (2) Polar retraction: Retrx (U) (X + Ulp + UT)--1/2 (3) Cayley
We establish convergence guarantees for a range of set- tings, including general nonconvex optimization prob- lems, nonconvex functions that satisfy the Riemannian Polyak-Lojasiewice (PL) condition, and stochastic set- tings under both general nonconvex and PL conditions. We show how our developments can be extended to quo- tient manifolds derived from the orthogonal manifold, including Grassmann and flag manifolds.
where VF())ss, vVX)vs(X)
(A+
We validate the effectiveness of the proposed method through extensive experiments, showcasing its fast con- vergence across a variety of problems, spanning both synthetic and real-world applications.
1.1. Related Works
Lezcano-Casado & Murtinez-Rubio (2019) re-parameterizes the orthogonality constraint in Euclidean space via the Lie exponential map. However, this approach requires differen- tiating through the exponential map, which can be computa- tionally expensive. Recently, Shalit & Chechik (2014); Gut- man & Ho-Nguyen (2023); Yuan (2023); Han et al. (2024a); Cheung et al. (2024) extend the idea of coordinate descent to Stiefel manifold by only updating few rows/columns while adhering to the orthogonality constraint. Despite the promise in cheap per-iteration update, they either suffer from poor runtime on modern hardware, such as GPUs by requiring a significant number of iterations to converge (Shalit & Chechik, 2014; Gutman & Ho-Nguyen, 2023; Han et al., 2024a) or involve subproblem that may become difficult to solve in general (Yuan, 2023). It is worth high- lighting that Cheung et al. (2024) lift the coordinate updates to the ambient space and then project back to the manifold. Their algorithms still require non- ota-dtad linear algebra
retraction: Retrx(U) WX for some skew-symmetric W E Raxn; (4) Exponential retraction: Retrx (U) -UTU) Texpmn-X U)]
(In -W)- + W)X
where U
[ U
X U] expmm
Ip
0
where expm(-) denotes matrix exponential. We highlight that all retractions require linear algebra operations other than matrix multiplications that costs at least (mp²)).
One classic Riemannian solver is the Riemannian gradi- ent descent (Udriste, 2013) that updates the variable as X = Retrx, ( F (Xx)), where 7e > 0 is the stepsize. Other more advanced solvers include Riemannian accelerated gradient methods (Ahn & Sra, 2020; Alimisis et al., 2021; Han et al., 2023b), variance reduced gradient methods (Bonnabel, 2013: Zhang et al., 2016; Kasai et al., 2018; Han & Gao, 2020; 2021a;b; Utpala et al., 2023), quasi-
2

---

Efficient Optimization with Orthogonality Constraint
Newton methods (Qi et al., 2010; Huang et al., 2015; 2018), and second-order methods (Absil et al., 2007; Agarwal et al., 2021), just to name few. Some more recent works develop algorithms for solving non non-smooth optimization (Ferreira & Oliveira, 1998; Chen et al., 2020; Li et al., 2024), min- max optimization (Zhang et al., 2023; Han et al., 2023a;c), bi-level optimization problems (Han et al., 2024c; Li Ma, 2025) on Riemannian manifolds.
Algorithm 1 RSDM
1: Initialize Xo E St(n..).
2: for k: = 0, .... K do
3: Sample Ps E (n) and let FY) F(Ux(Y)
where Ux (Y) is defined in (2).
4: Compute Riemannian gradient grad
Update Y2 = (-ngrd.-ngga
Set Xx+1 = (Yx))
Despite the recent progress in Riemannian optimization, All the aforementioned methods utilize the retraction operation and thus it becomes critical to reduce its complexity before scaling to large problems.
7: end for
We remark that our approach can be viewed as a gener- alization of the random subspace gradient descent in the Euclidean space (Kozak et al., 2021) to the Stiefel manifold. Specifically, the Euclidean random subspace updates the variable as J = 2x + ua (v) where uv:(v) = -P for some random matrix Ps that spans the subspace.
Notations. We use O(-) and ((.)) to denote the big-O and big- Big-Omega notation and (n) to represent the orthogonal manifold of size n X n.. P(n) C ((n) represents the set of permutation matrices and 8 {u E R" : :vv v = 1} is the unit sphere. We use = to represent diffeomor- phism between two manifolds. We use (+,-),ll || || to denote the Euclidean inner product and Euclidean norm, and use ( x)x, II- |Ix to denote Riemannian inner product and norm on Tx St(n,,) p). Because we only consider Euclidean metric as the Riemannian metric in this work, we use || and -l ||x interchangeably. We use P(r) to denote the first rows of matrix P.
To compute the Riemannian gradient Fr. let (a(()) E R. denote the first r rows of PR. Using the expression of Riemannian gradient, we can derive
gradFx(1,) = 2 (VFA( -VFFCC))VVAAATT))
= 2 ( X.VF(X)))) PA(O))
= (((P)graaFF(())P(())))
(3)
3. Riemannian Random Submanifold Descent
To ensure the updates adequately explore the full space with high probability, we re-sample the orthogonal matrix P& E (n)) each iteration. The distribution from which the orthogonal matrix is sampled will be discussed in detail in the subsequent section. The full algorithm outlined in Al- gorithm I, where we call the proposed method Riemannian random submanifold descent (RSDM).
This section introduces the proposed method that reduces the complexity of retraction by restricting the update random submanifold. In particular, each iteration k we parameterize the next iterate X XR for some UR E (n) and thus converts the problem to optimization over Us E ((n) on the orthogonal manifold. Such a param- eterization can be justified by the fact that the action of the orthogonal group O(n) over St (n, p) is transitive. Hence, at any point Xg, there exists a matrix Ug E (n) such that X" U. Xx, where X' is any local minimizer. Further, we parameterize the orthogonal matrix U3 by random orthog- onal matrix PR E (nn) and a low-dimensional orthogonal matrix Y E (r) where r is the lower dimension that we choose, and we define
4. Sampling Strategies and Complexities
In this section, we introduce two sampling strategies and re- spectively analyze the per-iteration complexity of Algorithm 1. We propose to sample Ps from two distributions, (1) a uniform distribution over the set of orthogonal matrices and (2) a uniform distribution over the set of permutation matri- ces. The second strategy of sampling from a permutation matrix is considered due to its sampling and computational efficiency compared to the orthogonal sampling.
Y
0 [n_r]
1
Un(Y)= = P.
Pk.
(2)
0
By minimizing Fx(Y) = F(UR(Y)) over instead of minimizing FU) = Fx) over U. we update the iterates on a random submanifold detined via (the first r rows of) the random orthogonal matrix Pt. Rather than min- imizing F3 to global optimality, we minimize the first-order approximation of the function Fà around 4, such that the update remains close to X₃. This suggests we can compute Y by taking a Riemannian gradient update from I,, i.c., y Retry,(- (-1k grad gg..FF(T.)) (T.)) for some stepsize mi > O.
The per-iteration complexity of Algorithm is attributed to four parts, i.e., sampling, gradient computation, gradient descent update and iterate update. For both sampling strate- gies, the gradient descent update (Step 5) shares the same complexity. In particular, the gradient update involves the retraction on O(r), which is on the order of O(r"). Next we respectively discuss the sampling and computational cost of each sampling strategy. As we show later, the per- iteration cost of orthogonal sampling strategy O(ppr)
3

---

Efficient Optimization with Orthogonality Constraint
while the per-iteration cost of permutation sampling strategy is O(rr2).
5. Theoretical Guarantees
In this section, we analyze the convergence guarantees for the proposed RSDM under both orthogonal sampling and permutation sampling. The proof' of all results are included in Appendix sections. We make use of the following nota- tions throughout the section. Recall we have defined in Sec- tion 3 that (U) F(UX+) and ( = (YUx())) at iteration k. We also introduce generic notations that Fx (U) = FUUX) and Fx(Y)= = XX)( for some sampled P E O(n).
Uniform Orthogonal. We first consider sampling Ps from the unique translation invariant probability measure (i.e. the Haar measure) on (n). Because only r rows of Pk is required, the sampling can be performed using QR decom- position (with Bram-Schhidd method) on a randomly sam- pled Gaussian matrix PE E _Æ"*" (Meckes, 2019). Hence, the cost of sampling is O(nr9). Furthermore, from (3), the computation of gradFs & (1,) requires a complexity of (nprr) by first computing Ps (r) and Xx before multiplication. For the iterate update, the computa- tion of U& only depends on the first rows of Pk as U = Xx + Pa(r)) (Y- - (r) which requires (nprr). This suggests the total per- verriteration com- plexity for uniform orthogonal sampling is O(ppr).
Assumption 5.1. F has bounded gradient and Hessian in the ambient space, i..., [[ I Co [[ <
Ci||U/] for any X E St(n, U E R"*x.
Assumption 5.1 is naturally satisfied given X E St(n,p), which is a compact submanifold of the ambient Euclidean space Rxxp, The next lemma verifies the (Riemannian) smoothness of (Y)),, which is due to Assumption 5.1. Lemma 5.2. Under Assumption 5. for any X E pt(n,,). (Y) is (Co + Ci)-smooth on O(r).
Uniform Permutation. Sampling from uniform distri- bution on permutation matrices corresponds to sampling permutation π : [x]] + [ra]. and thus the sampling complex- ity is negligible compared to other operations. In practice, sampling r indices without replacement from [rx] is suffi- cient. In terms of gradient computation, because Pk is a truncated permutation matrix, P'e (F()) and (x) corresponds to permuting the rows of (X())) and Xg. This largely reduces the cost compared to matrix multiplica- tion. Thus the gradient computation only requires Onr²). Lastly, for the iterate update, we highlight matrix multi- plication involving both (r) and (r)))) T corresponds to rearranging the rows and thus the cost can be reduced to O(r²)). ). Thus the total cost is 0(nr²).
Further, we show the following lemma that relates the gra- dient of Fx at identity to gradient of F(X).
Lemma 5.3. For any X E St(n,p), we can show (4,)) Z (X)))
Apart from general non-convex functions, we also analyze convergence of RSDM under (local) Riemannian Polyak- Lojasiewicz (PI.) condition.
Definition 5.4 (Riemannian Polyak--.ojasiewicr). For a
su C Sttnn,,) a smooth function F u4 -> IR satisfies the Riemannian Polyak-Lojasiewicz (PL) condition on U if there exists μ such that VX e U, we have F(X)
Remark 4.1 (Riemannian coordinate descent is special case). We show that with the permutation sampling and r = 2, RSDM is equivalent to the Riemannian coordinate descent on Stiefel manifold (Shalit& & Chechik, 2014; Han et al., 2024a; Yuan, 2023). To see this, we first recall that a Givens rotation (o) E On)) represents a sparse or- thogonal matrix such that its non-zero entries satisfy (1)
minxeu F(X) ك 1 (llraddF(X))
2μ
5.1. Main Results
This section derives theoretical guarantees for the proposed method. A summary of the main results are presented in Table 1. We first give the following proposition that relates the gradient of F to gradient of F at identity.
(0),( 1 for all # is and + 1; (2) (0),, =
0 for all i = I; (3) [G,,(0) -iin 0. Further we know that when r = 2. any Y E (()) can be parameterized by an angular parameter 0 and is
(G), =
Proposition 5.5. Assume that P is uniformly sampled from P(n) or uniformly sampled from O(n). Then for any X E St(n,p), we have Ellgrad (T,)) = (lgralllradFxx(
cos sin
either a rotation matrix R(Ø)=
cos or a
sin 0
reflection matrix F(0)- = coss easy to verify that (o) = Pk.1
Thus it is sin 0
where the expectation is with re- spect to the randomness in P.
sin 0
cos 0
R())
0
PT In-2
Remark 5.6 (Proof techniques of Proposition 5.5). We ob- tain the same rate for both the permutation and orthogonal sampling strategies. Nevertheless, the proof techniques are largely different. In the permutation case, the proof boils down to counting the number of permutations that satisfy some criterion and in the orthogonal case, we have to com- pute, for all set of indices, E. Pic, Pzt, Pak, Pjtz] for # j.
where Ps.i corresponds to the permuitation such that (1) = k n (2) = 1. This suggests that the update of RSDM reduces to X = (0), which is how coordinate descent is implemented in (Shalit& & Chechik, 2014; Han et al., 2024a; Yuan, 2023). In (Yuan, 2023), F(0) is further considered as an alternative to the rotation.
4

---

Efficient Optimization with Orthogonality Constraint
Table I: Summary of the main results under deterministic and stochastic settings, with both orthogonal (ortho.) and permutation (permu.) sampling. The global and local rates refer to the convergence under general nonconvex and PL conditions respectively. Size of Stiefel manifold is n x p and is the iteration number. Function constants I., 1 are ignored.
EXPECTATION
HIGH PROBABILITY
GLOBAL
LOCAL (PL)
GLOBAL
LOCAL (PL)
Deterministic
Ortho.
0 n
O(exp(= k))
0
=))
(Theorem 5.9,
Permu.
0 0(exp(- 1/4 (-)) -k))
Stochastic
Ortho.
n² (PVk)
O(expl (exp((=t))) + n20²)
(Theorem 5.13)
Permu.
This is achieved by leveraging the rotational invariance of the distribution of P.
This suggests RSDM matches the O(nppe- complexity of Riemannian gradient descent (RGD) in the regime when p 2 Cn for some constant > 0. This corresponds to the challenging regime where the retraction cost dominates other matrix operations. In contrast, when p n, the cost of retraction becomes relatively negligible.
Proposition 5.5 shows that the submanifold gradient is on the order of 2)) of the full-space gradient. In con- trast, the Euclidean subspace gradient method (Kozak et al., 2021) achieves a scaling of O(rn-1). This is because our proposed submanifold approach requires applying the pro- jection matrix Py twice, whereas the Euclidean subspace method requires only single Pe.
5.1.2. CONVERGENCE IN HIGH PROBABILITY
Theorem 5.7 suggests both permutation and orthogonal sam- pling guarantee the same convergence rate in expectation. However, we show in the following theorem that orthogonal sampling achieves much tighter convergence bound in high probability compared to the permutation sampling.
5.1.1. CONVERGENCE IN EXPECTATION
Theorem 5.7. When Ps is sampled uniformly from P(n)
or (n), under Assumption 5.1 and select 1) = 3 with
Theorem 5.9. Under Assumption 5.1 and nj = i with obtain
L= Co + C1, we obtain that for all k:
== Co + C1, if we use orthogonal sampling,
min ElllgadFF(X)))
ALAo n(n 1) k 1)
and for all k 2 1, with probability at least 1 - exp -11-7h (111-(n,1)) r)))).
where we denote Do F(Xx)- - F. Suppose further Xx converges to neighborhood LA that contains an (isolated) local minimizer X*. Further, F' satisfies Riemannian PL
16LA.
n(n - 1) (1 T ((n, (n2,)))r 1)
min ==0....k-I |grad
k
condition on U. Let ku be that Xko E U. Then we have Xko+R E U, Vk 2 and ([F(E(F(XX (x [F(()) S exp (- r(r-1) 1) k) 2L p(n
where T(u,r) = ex
(2( 112
(n If we use per-
mutation sampling, with probability at least 1 - exp
//4 (*) (k))
Theorem 5.7 shows that the convergence rate for general non-convex functions maintains the same sublinear con- vergence, compared with the Riemannian gradient descent (RGD) (Boumal, 2023), albeit with an additional -22 factors. Such a factor can be compensated by the lower per- iteration complexities of RSDM, which leads to a matching total complexity compared to RGD.
= 16LAo n(n n2 min gradF((X))] i=.....k
r(r-1)
Hence, in both cases, we have that almost surely, lim inf [lgradF (X+)) O. k-to0
Under the same setting in Theorem 5.7, suppose X E u. Then for orthogonal sampling, with probability at least 1 exp(- F(X*) S exp SL x(-1) 11 ((n,r))k)
Remark 5.8 (Total complexity of RSDM to RGD). The- orem 5.7, we show the convergence is at most On2²r~2/k) for both sampling strategies. This implies that in or- der reach an e-stationary point in expectation with min;=o...,A-1 F[l]ggrad((x))) eration complexity of 0 where per-iteration complexity is either (rspr) ) for orthogonal sampling or Onr²2) for permutation sampling strategy for permutation sampling, as analyzed in Section 4. This gives a total com- plexity of at least -2))
T(n,r))k), we have F(Xk++k)
F(X*)). For permutation sampling, with probability at least exxppp (- -1)), we have ((X++1) exp ( #
< e², we require an it-
4T. n(n (*) (F(X))))
Theorem 5.9 derives a high-probabbliin bound for for both orthogonal and permutation sampling strategies. For gen-
5

---

Efficient Optimization with Orthogonality Constraint
Algorithm 2 RSDM with partial deterministic sampling
that the following non-degenerate condition (4) holds:
1: Initialize Xu E St(n.p).
S-1
2: for k == 0,... K - do
Σ [PF(r)graaFF(())))))))))) CpllgradFF.
3: Sample {
S-1 such that condition (4) holds.
4: Set X Xk-
(4)
do 6: Let (Y) = XY) where UY) is de-
5: for s = -....
Such a non-degenerate condition over the selection of ran- dom matrix P ensures the projected gradient does not vanish. We show in Appendix F that there exist certain sampling schemes that satisfy condition (4).
fined in (2) with random matrix PE.
7:
Compute Riemannian gradient gradF (T.).
8:
Update E Retry, (-n (J.)))
We now derive the exact convergence guarantees as follows,
9:
Set x =
Xz
Theorem 5.11. Under Assumption 5.1, suppose RSDM is implemented as in Algorithm 2. Then let n = i. with L= Co + Ct. We can obtain for all k 2 1,
10: end for
11: Set XA+1 = X
12:
end for
1 K C C₁ M2222)
2LD.
min (=0 K-1 [ggrad((X))]
where S the inner iteration number, M is a constant depending on the retraction, Cp depends on the sampling.
eral nonconvex functions, it can be seen that the high- probability bound for permutation sampling can be much worse than for the orthogonal sampling due to the addi- tional binomial factor. In addition, compared orthog- onal sampling, permutation sampling requires the num- ber of iteration to be significantly larger in order for the bound to hold with arbitrary probability. To see this, we first can bound T(n,r) r) E (0.0.9995) due to r I n. 0.0005 < - T(n, r) S and thus 1 - T(n,r) = @(1). In order to require the high probability bound hold with probability 1 8' (for arbitrary 8' E (0, we require 2 4000 log require & 2 2(7)) (og(118')0" significantly large when n >> r.
5.3. Stochastic and Finite-sum Optimization
This section adapts Algorithm for stochastic optimization, defined as
min XER XTxXXPXXX=I,, =
where we obtain noisy estimates of the gradients by query- ing &. In Algorithm 1, we replace the Riemannian gradi- ent grad (T,) with the stochastic gradient grad (Tr: where we denote f(Y;e) = XYix for some randomly sampled & iteration k.
= f2(1) for the orthogonal case but
((())), which can be
(The trade-off between efficiency and con- vergence). The worse convergence guarantee of permuta- tion sampling relative orthogonal sampling in high prob- ability indicates trade-off between efficiency and conver- gence. Specifically for general nonconvex functions, permu- tation sampling requires only ((rr²) complexity per itera- tion while suffering from a convergence of (m²²r=(())/ /k) in high probability. In contrast, orthogonal sampling re- quires (mprr complexity per iteration but converges with rate of ( 2/kk) with high probability. Similar ar-
Remark 5.10
For convergence analysis, apart from Assumption 5.1, we also require the assumption of stochastic gradients being unbiased and having bounded variance, which is standard in analyzing stochastic algorithms (Ghadimi & Lan, 2013).
Assumption 5.12.
The stochastic gradient is unbiased, i.c.. [v(x(((;;) = ( and has bounded variance, i.e., (x - VF(X) S o², for all X E (n.
Theorem 5.13. Under Assumption 5.1 and 5.12, suppose
we choose 1) = min(L-} . / Lo -oo-- '-----lp 1/23, where we denote Do = F(Xo) F*. Then we can show
guments also hold for local linear convergence under PL condition.
E BlggraF((Xx)1 An(n 1) T.d + 22 AoLK
min
s r(r -
5.2. Exact Convergence of RSDM
1=
K
Suppose there exist Xko..., Xk, E U for some k > ko. where u is defined in Theorem 5.7. Then we have [(F(XH,) F(X*) S exp ( 2L m (k1 -
In the previous section, we have derived the convergence rate of RSDM in both expectation and with high probabil- ity. Here we adapt RSDM to achieve exact convergence. To this end, we adapt RSDM as described in Algorithm 2. Specifically, we implement RSDM with a double-loop procedure, where for each outer iteration k we sample pro- jection matrices (Pg) for the inner iterations, such
r(r-1)
n(n-1)
r(r-1)
Theorem 5.13 derives convergence guarantees for stochastic optimization and is comparable to Fuclidean analysis under
6

---

Efficient Optimization with Orthogonality Constraint
101
10"
RGD Landing
10
RGD
Optimailty Gap
10-1
PCAL
Landing
Optimality Ga
10
Landing
PCAI
RSOM-P
Optimallty lit Gap
20
10-1
RSDN P RSDM-0
Optimallty iity Gap
20-1
R5DM P
Landing
PCAL
1071
R P RSDM.Q
RSDM-0 -0
10
40
RSDM-O
10-4
10-1
25
75
100
125
150
175
0
10
20
30
42
50
10
15
20
25
20
20
40
60
20
100
Time
Time
Time
Time
(a) Procrustes (200, 200)
(b) Procrustes (2000, 2000)
(c) PCA (2000, 1500)
(d) PCA (3000, 2500)
Figure 2: Experiments on Procrustes problem and PCA problem under various settings. The numbers in brackets represent the size of n, p. For the Procrustes problem, we see RSDM converges competitively against the best baselines due to the simplicity the problem. For PCA problem, we see RSDM converges the fastest.
general nonconvex functions (Ghadimi & Lan, 2013) and under PL. conditions (Garrigos & Gower, 2023). The intro- duction of additional factor of O(n2r=3 2) is consistent with the deterministic setting, analyzed in Section 5.1. Lastly, we remark that although we focus on stochastic gradient in the main paper, the analysis can be easily extended to mini-batch gradient descent, as we show in Appendix
case, X2 a column orthonormal matrix whose columns span the subspace. Therefore, Algorithm 1 can be directly applied to the quotient manifolds. Because F& and F& are only defined on the orthogonal manifold, all our results de- rived for the Stiefel manifold still hold for general quotient manifolds. Apart from the orthogonal group, our develop- ments can also be similarly generalized to other compact matrix groups, such as 5O(n).
5.4. Generalization to Quotient Manifolds
6. Experiments
We now extend the developments of RSDM to general quo- tient manifolds. In fact, Stiefel manifold can be equiva- lently viewed as the quotient space of orthogonal mani- fold (Edelman al., 1998). More precisely, we can write St(n, p) = (n p). i..., point in the Stiefel man- ifold corresponds to the equivalence class
This section conducts experiments to verify the efficacy of the proposed method. We benchmark our methods with several baseline: (1) Riemannian gradient descent (RGD) on Stiefel manifold (Absil et al., 2008; Boumal, 2023); (2) Coordinate descent type of algorithms on Stiefel manifold, namely RCD (Han et al., 2024a) and TSD (Gutman & Ho- Nguyen, 2023); (3) Infeasible and retraction free methods, including PCAL (Gao et al., 2019) and Landing (Ablin & Peyré, 2022; Ablin et al., 2023).
[2) f0= :Ue : On om
In other words, each point in St(n,p) is the set of all thogonal matrices with the same first p columns. Such viewpoint allows to generalize the previous developments and analysis to more general quotient manifolds of the form
For all experiments, we tune the learning rate in the range of [0.01, 0.05, 0.1, 0.5, 1.0, 1.5, 2.0). For the infeasible meth- ods, we tune the regularization parameter in the range of [0.1, 0.5, 1.0, 1.5, 2.0). For the proposed method (RSDM), consider both permutation sampling and orthogonal sam- pling for Pk, which we denote as RSDM-P and RSDM-O respectively. We set the submanifold dimension accord- ingly based on the problem dimension and fix for both sampling strategies. By defaults, we use QR-based re- traction for RGD and proposed RSDM. All experiments are implemented in Pytorch and GPU. The code is available on https: / /gttttbtbb////////////// com/ andy im3 RSDM.
M = Onn)/K( -K.
U U U {()))
(5)
where K. is closed subgroup of On)). An element of M the equivalence class [Q] (K Q K E K.) Quo- tient manifold of the form (5) includes the famous Grass- mann manifold (Edelman et al., 1998), i.e., Gr(n,p) R (m)/(O(p) x O(n - p)) as well as the flag manifold (Zhu & Shen, 2024), i.e., (rl.(m.....
single RTX4060
Pad n) (On))(O(nn) x Om2 - n1) x ... x O(n - na-1) x O(n - nd)). Since the action of (nn) over M transitive, we can follow the same approach for St(n, p). and introduce function F&
6.1. Procrustes Problem
(m) + and F5=: : O(r) + R. where (U) x()),) and (Y) = (Y)), where (Y) is defined as in (2). We highlight that X₃ E . a representation of the equivalence class. For example, in the Grassmann manifold
We first consider solving the Procrustes problem as to find an orthogonal matrix that aligns two matrices, i.e., minxest(m.p)) s(X) |XA- B|² for some matrix A E

---

Efficient Optimization with Orthogonality Constraint
10
10"
RGD
10'
RGD (arl
:0"
RGD
Optimality Gap
55.20 P 1909
153N-6 0 5500
13-1
R5DM4-P R5DM-0
Optimality Gap
RGD (polar)
RGD (eep)
Optimality Gap
PCAL
Landing
152N-0 0 520
RGD (cayley)
RSDM P
10-1
553M P1 1901
0 (730)
10-1
RSDM (or)
RSDM (exp)
33-1
R5DM-0
10-
30-
RSDM (pdla))
20-
553N 01
0 2916
19
20
RSDM (cayley)
32-1
0
20
80 Time
200
120
140
14
20
15
Time
Time
Time
(a) Varying
(b) Varying randomness
(c) Varying retraction
(a) QUAD
Figure 3: Experiment results on PCA (n = 2000, p 1500) by (a) varying low-dimension
Figure 4: Experiment results on quadratic assignment prob- lem (QUAD), n = = 1000.
r and (b) random seed with r- = 700. The results suggest the outperformance of proposed RSDM over RGD is robust to changes in r as well as random seed.
RPxp, B RRxp, The optimal solution of this problem is x UVT where UEV BAT is the thin-SVD of the matrix BAT. Hence the optimal solution is computed as s(X*) = ||A + [B] 2tr(E).
matrix with a condition number of 1000 and exponentially decaying eigenvalues. Due the existence of an analytic solution, we measure the optimality gap in the same way as in Section 6. In Figure 2(c) and (d), we consider the setting of n 2000, p = 1500 and = 3000, p 2500, which represent large-scale scenarios. For the two settings, we set r = 700, 1000 respectively. We see RSDM achieves the fastest convergence among the baselines. Especially around optimality, we see RSDM switches from the sub- linear convergence to linear in contrast other baselines that maintains the sublinear convergence throughout. This behavior may be attributed to the random projection, which potentially provides a more favorable optimization land- scape close to optimum (Fuji et al., 2022). A formal theoretical verification of this claim is left for future work.
Setting and Results. We explore small-size as well as large-size problems by considering (1) n = p = 200 and (2) n = p = 2000. For the two settings, we set r = 150 and r 900 for RSDM respectively. We generate A, B where each entry follows random Gaussian distribution. For this problem, X" by taking SVD of BA and measure the optimality gap in terms of opt = (X(X) f(X* (LXxx)) We notice that for feasible methods, like RGD, RCD and proposed RSDM, the problem can be reduced to a linear function as maxxest(n.p)) (X, BAT) while for infeasible methods, the problem remains quadratic. We highlight that because p, RCD and TSD are equivalent.
compute the closed-form solution
To further validate the robustness of RSDM, we conduct additional experiments by varying the low dimension r, al- tering the random seed and utilizing different retractions. The results, presented in Figure 3, demonstrate that the per- formance of RSDM is largely insensitive to the choice of r (within a reasonable range) and the randomness throughout the iterations. RSDM also consistently outperforms RGD across all available retractions.
In Figure 2(a) under the setting performs notably worse compared to other benchmarks in runtime. This is because, although requiring fewer floating point operations (as shown in (Han et al., 2024a)), RCD requires more iterations, which is not GPU-friendly, On the other hand, we see the proposed RSDM performs com- petitively compared to RGD. When increasing the dimen- sionality to = p 2000, we notice RCD requires overly long runtime to progress and thus we remove from the plots. From Figure 2(b), we verify the superiority of RSDM over RGD.
200, we see RCD
6.3. Quadratic Assignment Problem
The quadratic assignment problem (Burkard et al., 1997: Wen & Yin, 2013) aims minimize a quadratic function over permutation matrix. In (Wen & Yin, 2013), the problem is re-formulated as a problem over the Stiefel manifold: (sn.)))( F(X) (r(A (X C X)B(X O X)T).
6.2. PCA Problem
Next, we consider quadratic problem, originating from principal component analysis (PCA), as to find the largest eigen-directions of a covariance matrix. This can be for- mulated as minxeSt(n.p) F(X) -xx AX). where A E Raxn_ This problem also has analytic solution given by the top-y eigenvectors of A.
Setting and Results. We consider the setting of n = 1000 and generate A, B as random normal matrices. Since no closed-form solution exists for this problem, we first run RGD for sufficient number of iterations, using the result- ing variable as the optimal solution. As shown in Figure 4(a), RSDM converges the fastest among the baselines, espe- cially near the optimal solution. Moreover, in this example,
Setting and Results. We create A to be a positive definite
8

---

Efficient Optimization with Orthogonality Constraint
RGD RSDM-
ACC (96)
ACC (%)
RSOM-P
ACC (%)
RSOM
23
0b
70
100
200
550
Time
Time
(a) O-FEN on MNIST
(b) O-FFN on CIFAR1O
(c) V-T on CIFARIO
Figure 5: Test accuracy for training orthogonal neural network (O-FNN) and orthogonal vision transformer (O-ViT) on MNIST dataset and CIFARIO dataset in five runs.
orthogonal sampling outperforms the permutation sampling.
non-negligible performance gap. This suggests that RSDM may offer greater advantages for training larger and more complex architectures such as ViTs.
6.4. Orthogonal Neural Networks
We consider optimizing a neural network with orthogo- nal constraints. We first consider a feedforward neural network (FFN) with ReLU activation for image classifi- L(n(()) (X), where
7. Conclusion
In this paper, we have introduced a novel randomized sub- manifold approach for optimization problems with orthog- onality constraints in order to reduce the high complexity associated with the retraction. We have derived convergence guarantees of the proposed method variety of func- tion classes and empirically demonstrated its superiority in number of problem instances. We also discuss two sampling strategies based on orthogonal and permutation matrices, and discuss the trade-off terms computational efficiency versus convergence guarantees.
cation task, i.e., min nn(X)= at... -- or (XW be) .... )W by denotes L- layer feedforward neural network with bias terms and L(--)) denotes the cross-entrope loss. Apart from the feedforward neural network, we also consider optimizing an (orthogonal) vision transformer (ViT) (Fei et al., 2022). We follow (Fei al., 2022) to impose orthogonality constraint on the query, key and value projection matrices.
Setting and Results. We optimize neural networks (orthog- onal FFN and orthogonal ViT) to classify MNIST (LeCun et al., 1998) and CIFARIO (Krizhevsky et al., 2009) im- ages. For preprocessing. the MNIST images are resized into 32 32 for MNIST and CIFAR 10 images to 20 x. The images are then normalized into [-1, 1] and vectorized as input for the neural network with a size of 1024 for MNIST and 1200 for CIFAR1O. We train 6-layer FFN, where we constrain the weight of the first 5 layers to be column orthonormal with a hidden size of 1024. The output layer weight, with a size of 1024 X 10, remains unconstrained. We also train a 6-layer, 4-head VIT with embedding dimen- sion 1024 and 64 patches, and constrain the query, key and value matrices of all attention layers to be orthogonal. For optimization, we employ RGD and RSDM with a batch size of 16. We set learning rate for unconstrained parameters to be 0.1 and only tune the learning rate for the orthogonal parameters. We plot the test accuracy in Figure 5 where we compare RGD with RSDM-P with five independent runs.
We believe our developments represent a significant ad- vancement in scalable Riemannian optimization by offering simple, yet effective solution for large-scale problems with orthogonality constraints. In the paper, we only dis- cuss the application of randomized submanifold strategy to Riemannian gradient descent, Nonetheless, we believe such strategy can be combined with more advanced opti- mization techniques, such as line-search (Boumal Cartis, 2019), momentum (Li et al., 2020; Kong et al., 2023), pre- conditioning (Kasai et al., 2019) and higher-order methods (Huang et al., 2015; Absil et al., 2007), to further enhance convergence efficiency and robustness with orthogonality constraints and beyond.
Limitation. We remark that our random submanifold method has matching complexity as RGD when p 1(2). Extending development to the case where p re- mains an important future direction.
When training an orthogonal FFN, we observe that RSDM achieves faster convergence during the early iterations in terms of runtime, indicating its greater efficiency in quickly reaching high accuracy. On the other hand, when training an orthogonal ViT, RSDM consistently outperforms RGD in test accuracy throughout the training process, with a
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

---

Efficient Optimization with Orthogonality Constraint
References
Chen, M.-S., Wang, C.-D., Huang, D., Lai, J.-H., and Yu, P. S. Efficient orthogonal multi-view subspace clustering. In ACM SIGKDD conference on knowledge discovery and data mining, pp. 127-135, 2022.
Ablin, P. and Peyré, G. Fast and accurate optimization on the orthogonal manifold without retraction. In International Conference on Artificial Intelligence and Statistics, pp. 5636-5657. PMLR, 2022.
Chen, S., Ma, S., Man-Cho So, A., and Zhang, T. Proximal gradient method for nonsmooth optimization over the Stiefel manifold. SIAM Journal on Optimization, 30(1): 210- 130-239,,00 2020.
Ablin, P.. Vary, S., Gao, B., and Absil, P.-.. Infeasible deterministic, stochastic, and variance-reduction algo- rithms for optimization under orthogonality constraints. 0rXiv:2303.165/1 2023.
Cheung, A. Y.-M,, Wang, J., Yue, M.-C., and So, A. M.- C. Randomized submanifold subgradient method for optimization over Stiefel manifolds. arXiv:2409.01770, 2024.
Absil, P.-A. and Malick, J. Projection-like retractions matrix manifolds. SIAM Journal on Optimization, 22(1): 135-158, 2012.
Darmwal, Y. and Rajawat, K. Low-complexity subspace- descent over symmetric positive definite manifold. arXiv:2305.0204 020041 2023.
Absil, P.-A., Baker, G., and Gallivan, K. A. Trust-region methods on Riemannian manifolds. Foundations of Com- putational Mathematics, 7:303-330, 2007.
Edelman, A., Arias, T. A., and Smith, S. T. The geometry of algorithms with orthogonality constraints. SIAM jour- nal Matrix Analysis and Applications, 20(2):303-353, 1998.
Absil, P.-A., Mahony, R., and Sepulchre, R. Optimization algorithms matrix manifolds. Princeton University Press, 2008
Agarwal, N., Boumal, N., Bullins, B., and Cartis, C. Adap- tive regularization with cubics manifolds. Mathemati- cal Programming. 188:85-134, 2021.
Fei, Y., Liu, Y., Wei, X., and Chen, M. O-vit: Orthogonal vision transformer. arXiv:2201.1113,, 2022.
Ferreira, O. and Oliveira, Subgradient algorithm on Rie- mannian manifolds. Journal of Optimization Theory and Applications, 97:93-104, 1998.
Ahn, K. and Sra, From Nesterov's estimate sequence to Riemannian acceleration. In Conference Learning Theory, pp. 84-118. PMLR, 2020.
Fuji, T., Poirion, P.-L., and Takeda, A. Convexification with bounded gap for randomly projected quadratic opti- mization. SIAM Journal on Optimization, 32(2):874-899, 2022. doi: 10. 0,.133/21MM133673
Alimisis, F., Orvieto, A., Becigneul, G., and Lucchi, Mo- mentum improves optimization on Riemannian manifolds. In International Conference Artificial Intelligence and Statistics, pp. 1351-1359. PMLR, 2021.
Gao, B., Liu, X., and Yuan, Y-x.. Parallelizable algorithms for optimization problems with orthogonality constraints. SIAM Journal on Scientific Computing, 41(3):A1949- A1983, 2019.
Bonnabel, S. Stochastic gradient descent on Riemannian manifolds. IEEE Transactions on Automatic Control, 58 (9):2217-2229,, 2013.
Garrigos, G. and Gower, R. M. Handbook of con- vergence theorems for (stochastic) gradient methods. arXiv:2301. 2023.
Boumal, N. An introduction to optimization on smooth manifolds. Cambridge University Press, 2023.
Boumal, N. and Cartis, C. Global rates of convergence for nonconvex optimization on manifolds. IMA Journal of Numerical Analysis, 391::1133 2019.
Ghadimi, S. and Lan, G. Stochastic first-and zeroth-ordd methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4):2341-2368, 2013.
Burkard, R. E., Karisch, S. E., and Rendl, F. QAPLIB-2 quadratic assignment problem library. Journal of Global oplimization, 10:391-409 -033, 1997.
Götze, and Sambale, H. Higher order concentration on stiefel and grassmann manifolds. Electronic Journal of Probability, 28:1-30, 2023.
Chaudhry, A., Khan, N., Dokania, P., and Torr, P. Continual learning in low-rank orthogonal subspaces. Advances in Neural Information Processing Systems, 9991,, 2020.
Gutman, D. H. and Ho-Nguyen, N. Coordinate descent without coordinates: Tangent subspace descent on Rie- mannian manifolds. Mathematics of Operations Research, 48(1):127-159, 2023.
10

---

Efficient Optimization with Orthogonality Constraint
Han, A. and Gao, J. Escape saddle points faster on manifolds via perturbed Riemannian stochastic recursive gradient. arXiv:2010. 12191, 2020.
Huang. W., Gallivan, K. A., and Absil, P.-A. A Broyden class of quasi-Newton methods for Riemannian optimiza- tion. SIAM Journal on Optimization, 25(3):1600-1685, 2015.
Han, A. and Gao, J. Improved variance reduction methods for Riemannian non-convex optimization. IEEE Transac- tions on Pattern Analysis and Machine Intelligence, 44 (11):7610-7622) 2021a.
Huang, W., Absil, P.-.,, and Gallivan, K. A. A Rieman- nian bigs method without differentiated retraction for nonconvex optimization problems. SIAM Journal on Op- timization, 28(1):470-443 2018.
Han, A. and Gao, J. Riemannian stochastic recur- sive momentum method for non-convex optimization. In Proceedings of the Thirtieth International Joint Conference on Artificial Intelligence, pp. 2505-2511, 2021b. URL https: / http//oi..co or1/000.......... /249///3/.223////// jjcai. 2021/345.
Kasai, H., Sato, H., and Mishra, B. Riemannian stochastic recursive gradient algorithm. In International Conference on Machine Learning, pp.
PMLR, 2018.
Kasai, H., Jawanpuria, P., and Mishra, B. Riemannian adap- tive stochastic gradient algorithms on matrix manifolds. International Conference on Machine Learning, pp. 3262 3271. PMLR, 2019.
Han, A., Mishra, B., Jawanpuria. P., and Gao, J. Nonconvex- nonconcave min-max optimization on Riemannian man- ifolds. Transactions Machine Leurning Research, 2023a.
Khan, A. and Maji, Multi-manifold optimization for multi- view subspace clustering. IEEE Transactions on Neural Networks and Learning Systems, 33(8):3895-3907, 2021.
Han, A., Mishra, B., Jawanpuria, P., and Gao, J. Rieman- nian accelerated gradient methods via extrapolation. In International Conference on Artificial Intelligence and Statistics, pp. 1554-1585. PMLR, 2023b.
Kong, L., Wang, Y., and Tao, M. Momentum stiefel op- timizer, with applications to suitably-orthogonal atten- tion, and optimal transport. In International Conference on Learning Representations, 2023. URL https: openreview net / (wworrmid=vcJ-----
Han, A., Mishra, B., Jawanpuria, P., Kumar, P., and Gao. J. Riemannian Hamiltonian methods for min-max optimiza- tion on manifolds. SIAM Journal on Optimization, 33(3): 1797-1827, 2023c.
Kozak, D., Becker, S., Doostan, A., and Tenorio, L. A stochastic subspace approach gradient-free optimiza- tion in high dimensions. Computational Optimization and Applications, 79(2)3339-368, 2021.
Han, A., Jawanpuria, P., and Mishra, Riemannian coor- dinate descent algorithms on matrix manifolds. Interna- tional Conference on Machine Learning, 2024a.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers features from tiny images. 2009.
Han, A., Jawanpuria, P., and Mishra, B. Riemannian Op- timization. In Encyclopedia of Optimization, pp. 1-6. Springer, 2024b.
LeCun, Y., Bottou, L.,, Bengio, Y., and Haffner, Gradient- based learning applied document recognition. Proceed- ings of the IEEE, 86(1 1:2278-2322. 1998.
Han, A., Mishra, B., Jawanpuria, P. K., and Takeda, A. A framework for bilevel optimization on Riemannian manifolds. Advances in Neural Information Processing 2024c.
Lezcano-Casado, M. and Martinez-Rubio, D. Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group. In International Conference on Machine Learning, pp. 3794 3803. PMLR, 2019.
Systems, 37:103829-103388
Helfrich, K., Willmott, D., and Ye, Q. Orthogonal recurrent neural networks with scaled cayley transform. In Interna- tional Conference on Machine Learning, pp. 1969-1978. PMLR, 2018
Li, J. and Ma. Riemannian bilevel optimization. Journal of Machine Learning Research, 26(18): 1-444 2025.
Hotelling, H. Analysis of complex of statistical variables into principal components. Journal of Educational Psy- chology, 24(6):417, 1933.
Li, J., Fuxin, L.,, and Todorovic, S. Efficient riemannian opti- mization the stiefel manifold via the cayley transform. urXiv:2002.01113 2020.
Huang, M., Ma. S., and Lai, L. A Riemannian block co- ordinate descent method for computing the projection robust Wasserstein distance. In International Conference on Machine Learning, pp. 4446 4446-4455. PMLR, 2021.
Li, J., Ma, S., and Srivastava, T. A Riemannian alternating direction method multipliers. Mathematics of Opera- tions Research, 2024.

---

Efficient Optimization with Orthogonality Constraint
Li, Jia, K., Wen, Y., Liu, T., and Tao, D. Orthogonal deep neural networks. IEEE Transactions Pattern Analysis and Machine Intelligence, 43(44)1352-1368, 2019.
Theis, F. J., Cason, T. P., and Absil, P. A. Soft dimension reduction for ICA by joint diagonalization on the Stiefel manifold. In Independent Component Analysis and Sig- nal Separation: 8th International Conference, ICA 2009, Paraty, Brazil, March 15-18, 2009. Proceedings 8, pp. 354-361. Springer, 2009.
Lin, T., Fan, C.,, Ho, N., Cuturi, M., and Jordan, M. Pro- jection robust Wasserstein distance and Riemannian op- timization. Advances in Neural Information Processing Systems, 33:9383-9397, 2020.
Udriste, C. Convex functions and optimization methods on Riemannian manifolds, volume 297. Springer Science Business Media, 2013.
Liu, W., Qiu, Z., Feng, Y., Xiu, Y., Xue, Y., Yu, L., Feng, H., Liu, Z., Heo, J., Peng. S., et al. Parameter-efffcient orthog- onal finetuning via butterfly factorization. In The Twelfth International Conference on Learning Representations, 2024.
Utpala, S., Han, A., Jawanpuria, P., and Mishra, B. Im- proved Differentially Private Riemannian Optimization: Fast Sampling and Variance Reduction. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
Liu, X., Liu, L., Liao, Q., Wang, S., Zhang, Y., Tu, W., Tang, C., Liu, J., and Zhu, E. One pass late fusion multi- view clustering. In International conference on machine learning, pp. 6850-6859, PMLR, 2021.
Vandereycken, B. Low-rank matrix completion by Rieman- nian optimization. SIAM Journal on Optimization, 23(2):
1236, 2013.
Vary, S., Ablin, P., Gao, B., and Absil, P.-A. Optimization without retraction on the random generalized Stiefel man- ifold. In International Conference on Machine Learning. PMLR, 2024.
Meckes, E. S. The Random Matrix Theory of the Classical Compact Groups. Cambridge Tracts in Mathematics. Cambridge University Press, 2019.
Vershynin, R. High-dimensional probability: An introduc- tion with applications in data science, volume 47. Cam- bridge university press, 2018.
Mishra, B., Meyer, G., Bonnabel, S., and Sepulchre, R. Fixed-rank matrix factorizations and Riemannian low- rank optimization. Computational Statistics, 29:591-621, 2014.
Wang. J., Chen, Y., Chakraborty, R., and Yu, X. Or- thogonal convolutional neural networks. In Conference Computer Vision and Pattern Recognition, pp. 11505- 11515, 2020.
Peng, L. and Vidal, R. Block coordinate descent on smooth manifolds: Convergence theory and twenty-one examples. arXiv:2305.14744, 14744, 2023.
Wen, Z. and Yin, W. A feasible method for optimization with orthogonality constraints. Mathematical Program- ming, (42(1):397-434 2013.
Qi, C., Gallivan, K. A., and Absil, P.-.. Riemannian BFGS algorithm with applications. In Recent Advances in Op- timization and its Applications Engineering, pp. 183- 192. Springer, 2010.
Xiao, N., Liu, X., and Yuan, Y.xx. A class of smooth exact penalty function methods for optimization problems with orthogonality constraints. Optimization Methods and Software, 37(4): 1e011120512200 1000-121, 2022.
Qiu, Z., Liu, W., Feng, H., Xue, Y., Feng, Y., Liu, Z Zhang, D., Weller, A., and Schölkopf. B. Controlling text- to-image diffusion by orthogonal finetuning. Advances in Neural Information Processing Systems, 36:79320- 79362, 2023.
Yuan, G. A block coordinate descent method for nonsmooth composite optimization under orthogonality constraints. arXiv:2304. 0364044,,.61.... 2023.
Zhang, H., J Reddi, S., and Sra, S. Riemannian SVRG: Fast stochastic optimization on Riemannian manifolds. Advances in Neural Information Processing Systems, 29, 2016.
Rebjock, Q and Boumal, N. Fast convergence to non- isolated minima: four equivalent conditions for c2 func- tions. Mathematical Programming, pp. 1-49, 2024.
Shalit, U. and Chechik, G. Coosdinate-descent for learning orthogonal matrices through Givens rotations. In Inter- national Conference on Machine Learning, pp. 548-556. PMLR, 2014.
Zhang, P., Zhang. J., and Sra, S. Sion's minimax theorem in geodesic metric spaces and a Riemannian extragradient algorithm. SIAM Journal on Optimization, 33(4):2885- 2908, 2023.
Shustin, B. and Avron, H. Faster randomized methods for orthogonality constrained problems. Journal of Machine Learning Research, 25(257):1-59, 2024.
Zhu, X. and Shen, C. Practical gradient and conjugate gradient methods on flag manifolds. Computational Opti- mization and Applications, 88(2):49155) 2024.
12

---

Efficient Optimization with Orthogonality Constraint
A. Other related works
Apart from orthogonality constraints, Han et al. (2024a) derive efficient coordinate updates for other matrix manifolds, such as Grassmann and positive definite manifolds. Vary et al. (2024) extend the idea of infeasible update for generalized Stiefel manifold and Darmwal & Rajawat (2023) propose efficient subspace descent algorithms for positive definite manifold with intine-ivvarinnce metric. Several other studies (Huang al., 2021; Peng & Vidal, 2023) investigate (block) coordinate descent on product of manifolds, where each update targets an individual component manifold. This however is less relevant to our setting, where we exploit submanifolds on a single manifold.
B. Preliminaries on Stiefel manifold and Riemannian optimization
In this section, we provide more detailed introduction to the Stiefel manifold and Riemannian optimization, Stiefel manifold (n,n, p) {X Raxp= : X X / is the set of column orthonormal matrices. When n = p. (t, p) (()) which is called orthogonal manifold, also forming group. The tangent space of Stiefel manifold is St,,t(n,)) p) = {U RAxp: X + U X 0}. One can choose the Euclidean metric (restricted to the tangent space) as a Riemannian metric for St(n,p), p). i.e., for any X E St(n,,), p). and U use ( to represent the Euclidean inner product. Other metrics such as canonical metric (Edelman et al., 1998) can also be considered. Orthogonal projection of any W E _ to TxSt(nn p) with respect to the Euclidean metric is derived as Px(WW W x W's, where we denote {A}s == (A + //2. For a smooth function F St p) + E Riemannian gradient of F X Su(n,)), p). denoted as g(adF(X), ). is tangent vector that satisfies for any U Tt, p), ()) = (VF (X)),) U), where VF denotes the classic Euclidean gradient. The Riemannian gradient on Stiefel manifold can be computed as gradF(X) Px(VF(()) VF(X) X VF(X)}s.
V E Sx,,(n,))) p). Riemannian metric (U,V) = {U, V) where we
Riemannian optimization works by iteratively updating the variable on the manifold following some descent direction. Throughout the process, a retraction is required to ensure that the iterates stay on the manifold. Specifically, a retraction, denoted as Retrx= : TxSt p) + St,,,p) is a map from tangent space to the manifold that satisfies (etrx(0) = X and DRetrx (0) = V for any V E StS,(n, p). where
the differential operator. There exist various retractions on Stiefel manifold, including (1) QR-based retraction: Retrx (U) qf(x + U). where of extracts the Q-factor from the QR decomposition; (2) Polar retraction: Retrx (U) = (X + U) UT)-1/2
(3) Cayley retraction: Retrx(U) (In (I, + W)X where U = W for some skew-symmetric W E R (4) Exponential retraction: Retrx(U))
T T
XTU -UTUD
expm((X UU] where expm(-) denotes matrix exponential. We highlight that all
[X U] expmm
0
retractions require linear algebra operations other than matrix multiplications that costs at least O(np²).
C. Additional experiment results
In the main text, we present the convergence results only in terms of runtime. Here we also plot the convergence with respect iteration number in Figure 6. We see for Procrustes problem, of the simplest optimization problems on Stiefel manifold, both RGD and Landing algorithm yields fastest convergence in iteration number. We also notice in small-sized problem, RCD converges quickly. Nonetheless, each iteration of RCD requires to loop through all the n² indices, resulting in poor parallelizability. This is reflected in the runtime comparisons presented in the main text. For other problem instances, including PCA and quadratic assignment, RSDM attains the fastest convergence not only in runtime (as shown in the main text) but also in terms of iteration count (Figure 6).
Finally, we plot the convergence in iteration for training orthogonal neural networks on MNIST and CIFARIO. We see that RSDM is not able to beat the RGD in terms of convergence in iteration, due to the difficulty of the optimization problems.
D. Proofs
D.1. Proof' of Lemma 5.2
We first recall the Hessian of a function G: ::(St, p) + R along a any tangent vector V is
= PX - vLx V[XTVG((X))
where (A)s (A+ + AT)/2 and Px = -
13

---

Efficient Optimization with Orthogonality Constraint
Optimailty Gap
32-1
Landing
PCAL
107
RSDM-
PCAL
Landing
Optimallt Gap
20:1
10-1
10"
RSDM P
RGD Landing
Optimallty iity Gap
20-1
10
RSDM P
Landing
RGD
40-1
RSDM-P
Optimala
R5DM 0
10
RSDM-D
RSDM-0
10-4
230
400
sco a00
1300 1400
230 300 750 1000129100007000000
250
405
600
800
1000
200
403
650
802
1030
Iteration
Iteration
iteration
Iteration
(a) Procrustes (200, 200)
(b) Procrustes (2000, 2000)
(c) PCA (2000, 1500)
(d) PCA (3000, (3000,2500
RGD
Landing
Optimaling ity Gap
RSDM
(96
20m-
R5DM-0
RED
10.2
3002
2000
3002
4000
20
50
03
70
10
60
Iteration
Iteration
Iteration
(e) QAssign (1000, 1000)
(f) MNIST
CIFARIO
Figure 6: Convergence in terms of iteration Procrustes problem and PCA problem and quadratic assignment problem under various settings. We observe that except for the Procrustes problem and training of orthogonal neural network, RSDM also converges the fastest in terms iteration number.
2.2
1.5
RSD
3.0
RSDMP
RSDM-
0.2
20
50
120
20
Time
Time
MNIST
(b) CIFARIO
Figure 7: Convergence in loss plot for on image classification.
Proof of Lemma 5.2. Recall the (Euclidean) gradient and Hessian of Fx (Y) is derived
PWW.UI/MPPPYYFFUUUXXXXXXXX P(r)
for any V E (r)(r). This leads to the following Riemannian Hessian
Hess H = Py -YY
We wish to bound (X) in terms of [Vly.. First we notice that ||Hess (V)(V]llx & (
14

---

Efficient Optimization with Orthogonality Constraint
V(Y ((YT)Xx(Y))) because for any ç.
(P = -1115- 1 (2l1e1I² 2(5,
1 (211E11 2voo(5) T (Y @ Y)voo(E
S 4 (211511 2115119)
||
where we use the fact that |I @ Y'ul
||a||| for any and Y O(r).
Then we bound
- {G]U((V) +
-Co (Co+G))lVV)
where triangle inequality in the first inequality. The second incquality uses [P()) < 1, ||X|], ||]| s 1. The third inequality is by assumption on F()( (XX).. The last equality is by definition of (V)..
D.2. Proof of Lemma 5.3
Proof of Lemma 5.3. From the definition that FX( F(UX), we let
grad F x (In)]12 -1 -X
= Iggrd = [VF(X) - VX(XVx(X))s
We first notice that
A (I -
4
1
(I (XVVVE(XXXXVXXX)))
2
Similarly,
= | X(X
X V
where we let C [VVZ)))==MMXXXXX])]XXX(X133
((VF((X)) XX VF(X)). Then have
C +IXVXX))XXXXXXXXXXXXXX xIP 1 x XVF(x)
XX IXVF((X
- xx VF(X)) ( VXCXXV(())V x)+ +0
- =t((V((X) XX [ VF(X))
Therefore,
B= - - TXX
xx V 1
XVF(X)
2
15

---

Efficient Optimization with Orthogonality Constraint
and,
A - (
We will now prove that To this end, we consider the quantity [IF(X X F( O. By expanding obtain
A B A.
[IF(X) x X - F(X)]| 2xx(F(X) XF(X) x)
= =2(1=((x)'X]'=-11) (xxX)
(xxx(VE()"xxxVV> (T(XF(X)T O.
This shows
tx(XVF(X) (XXXXX)XXXXX))))))) (u((P(()) xX Xx XxVFF(X))
which concludes A B 2 A and thus > B/2.
D.3. Proof of Proposition 5.5
Proof of Proposition 5.5. We We separately prove the results for the strategies of permutation sampling and orthogonal sampling.
Permutation sampling. Recall that the gradient grad Fx (In) on ((n) is given by
gr)dFx((n) 1 (VF(()X - XVF(X)))
Hence by (3) and using the fact that P is permutation matrix, we have
grad (I+) P(r)grrdFFx(())PP))
where P(r) R consist of the first r rows of P. us define by S% the set of all truncated permutation matrix P(r) E Rexx, To each element of Sh we can associate a unique "truncated" permutation i defined by
Vi S r: tii) is such that PT E; = (x(()
Notice that क is defined only for the first r integers as the matrix P E S has only rows. Using the fact that grad (In) Skew-symmetric, we have
[P =2 Σ
(gn)
Σ
Isicisr
1stJsr
We can check the number of element in S is given by
(n r)!
as n! the number of permutations on {1.... n) and (n - r)! is the number way to complete the truncated permutation
into a permutation on {l...., .n.. Therefore, we deduce that
E Elgrad = - rr)!
Σ
Σ
(gradFx
(Ln)
nd
TEST i<j
Let us now fix < k n. We will now count how many times does the term (grrdFF((4)))).. appears in the sum
ΣΣ (grad
REST
16

---

Efficient Optimization with Orthogonality Constraint
More formally, let denote by ng, E N, the number of time (gradFx(In))))))) appears in the above sum. Then we have that
B [IgradFx(4))]] 2 (n-r)) -r
Σ
(n.(graddFx(x)))
n!
IERIEN
Let us now compute Rg.J- For that is equivalent to count how the number of elements т E S such that (k; E ((T(1, +()) i < j}. This equivalent to choosing an ordered pair (i,j) inside {1,-.. r and then count the number of elements v Sy such that (k: I) (m))), m()). Hence, we deduce that
r(r (n 2)!
TR
(n - r']
We deduce therefore that as there is exactly (n 2) permutation such that for fixed (i.j), (k, = (z)
E (igradFx(4.)1) = (n - r)! r(r - 1) (n 2)
(
n!
2 (n r)!
Iskkss
r(r - 1)
n(n 1) |gradFx(4.)11
which completes the proof.
Orthogonal sampling. Recall
grad (I,) = ((r)gradF((.))P()
where P(r) E IR"*" consist of the first r rows of P. Let us denote by M = gradFxx (In). We have
x
Σ
Pa.Ma
1
Σ
Piky Mant Pata Pike Musta Pata
Σ Par Paty Pag ManM
Σ
where the last inequality holds as PMPT is skew symmetric as M is, hence (PMP )a = 0 for all i, Hence, to prove the theorem, we must compute, for all set of indices, E[Pa, Puk, for F
First, let us consider the case where all indices kj 1, 12 12 are different. We will prove that in such case Indeed, notice that since all the four indices are different, Paky Pst. Pikgs Psta belongs to four different columns of P. Hence, by multiplying P on the left by an identity matrix where the 1 are k1 position on the diagonal has been replaced by 1,, we can change Pak, to - Pas, and Put, Pjt, Pity Pitz to Pity Pjt, Pat, Pzta- Since the distribution of P is invariant with such operation, we deduce, by symmetry, that IE[Pika PJt. Pike p
E[Pa,, Pj Pial
More generally, let us now consider the case where k1. 11, k2, l2 take at least 3 different values. Then by a similar reasoning, since once column of P must contain least single index among ki, k2, l l2 then we can show (by multiplying this column by 1) that Pit, Pjs Piks Psta has the same distribution Hence, we need to consider two cases: kn l2, h l (notice that kn l2 is the same case as k1 41).
Pak, PJt, Pakz Pyta. proving again that E Pak, Pjt, Pak. Pital 0.
First, let us assume that k1 l2. Then, by the previous point, we must also have that / 12. Now, let consider the case ly = 41 Again, we must have by = 1z, otherwise we would have at least 3 distinct columns.
In summary, we have proved, by considering the three cases: k = 12 l = lz; k1 = 11, ho = la andk; = 12, A:2 = 4, that: (lgradF((T.)11
n
T Σ 1,1=1
Σ
Σ 1=
B
P2 M
Σ
PPPPPPPPPPP
#1
1,k=1
17

---

Efficient Optimization with Orthogonality Constraint
Which implies, by anti-symmetry of M (My-
Ma and MRs=
Mu = =0)
(graFF((I))1
Σ
Σ
[E[P
E[PiPPP&PPPp))) Ma.
(6)
1zj=1 1
Let us now compute EIP Pj). for k + L. Notice for all i, Ek=1
P2 1. Hence, by multiplying two this equality (for i
and j) and taking the expectation, we get that for all i + J.
B[P
R
30 1,
th
1==1
which implies that
R
n
E.P p2 =1
However, notice that for all i * 1, the joint law of Pak, Pjt is the same. Indeed the law of P does not change by permuting the columns of P, which implies that for all k, the law of Pik, Pst is the same the joint law of Pai Pjz. Hence we have that from the previous equation that
1n²-n)BBl n)
nE[P,
(7)
Furthermore using that
Σ 1=
Psk = =0,
leading to
Σ
PaPPsPaPPe +
Σ
ik 3k
p2 0.
k
Since, by permuting the rows of P. PxPPxPPPP has the same law as P.I PiPPPPPPPpp we found that
Rea
P} = 0.
(8)
Notice that (6) leads to
E
Igrad Fx.))1
Σ
[E PPPPPpPPPP3)) Ma.
1/j=1 1+k=
Hence
[ggradFxx (4.)1
r) (EIP (x(2)
(
21 IggradFx (4,))12.
(9)
Using (7), we found that
BIPă 1 1- P3
n n
and using (8), we found that
nE[P?
E P3PPPPPPPppl
n 2--N n
Hence,
,2
E
[Igradxx (1.)
ExaaFF((!.))
n2
Thus the proof is now complete.
18

---

Efficient Optimization with Orthogonality Constraint
D.4. Proof of Theorem 5.9
Here we first prove Theorem 5.9, which is the high probability convergence guarantees. Towards this end, we require the following proposition that deduces concentration inequality on IlgradFx (I.) Proposition D.I. When P is sampled uniformly from (1i), we have that
P (IeradFx (4,11 2 r(r -1) (3radF((4.)13 2 1 eexp
exp (- ,²(r - 1)2
2048n2 (-2048n-(1 - 1)2
2n(n 1)
When P is sampled uniformly from P(n) we have that
2 r(r r(r-1) 1) n(n - 1) (eradF(((.))))) 2 (7)-
P (leradFx(1x)1
Proof of Proposition D.I. First, let us consider the case where P is sampled uniformly from O(n). When P is sampled uniformly from (n), we can see P(r) is sampled uniformly from St(n,r), r). Then by Proposition 5.5,
E ( r(r - 1) n(n - 1) Igrad/(x
In order to derive a high-probability result, we define the following function h St(n, r) + R that 1(X) = I| M where M E R is any Skew-symmetric matrix. And it can be verified that when M = grad (In). we can show h(P(r)T = x
Let us now compute a Lipschitz constant La for the function h. For that, we compute the Riemannian gradient gradh(X). Le us first compute the Euclidean gradient Vh(X). We have, by anti-symmetry of M:
Vh(X) A
We therefore deduce the Riemannian gradient:
gradh(X) 4MXX MX + +XXXX" MXX MXX'XX)s -MMX MX + 4XX MXX MX
This implies that in order to find the Lipschitz constant La, we need to bound || M and |XX MXX MX]] Using that for any matrix A, B, we have that ||AB|| I ||/ and that || S for any X E St(n,r), we can bound the two term above by || [l.. Hence, we deduce that we can take Ln ||| the Lipschitz constant for h. From (Götze& & Sambale, 2023), we deduce that for any >> O, we have
P h(X)= < r(r r(r-1) 1)
(n 1)t²
n(n |M -t S Sexp
512
Hence, by taking t r(r-1) || 112, we deduce that m(
-(r- 1)
(r 1)2
P ((()) S
2n((n - 11 |1
exp
2048n2 1)²
This ends the proof for the case where P is sampled uniformly from On)). Notice that the permutation case is obvious as each element sampled with probability (") and at least one element P should induce a value h(P(r)T) larger than [)
-1
Proof of Theorem 5.9. By Lemma 5.2, we see (Y) is L-smooth with L = Co + C1. Then we have for any Y and W& E Ty rr,,),
(x(eery((W)) Fx( + (ggaaddxxx) L wwi2
19

---

Efficient Optimization with Orthogonality Constraint
Applying this inequality to the update and recalling Fx(() F(U(Y)x&)) we have
F(X++1)= =
S FT,,)) 11
- 2
1 (gradF2 (1,)1
=({+)
(10)
2L
where we note that (I,) F-Xx)-
Hence, we deduce from Proposition D.1 and Lemma 5.3, that
1 r(r - 1)
[grrd < r(r - 1) 8L n( |grad
F(X++1) - F(X+) S 4L n(n - 1)
()--1)² 20-ss : 1- r.. Let us denote, for all k. by Yv [0,1] the
holds with probability at least 1 exp random variable equal to one if and only if the above inequality holds. We have that B|Y]] 2 1 (n,r). furthermore since F(X++1) I F(xx). we have that for all k;
S 8L n(n 1)) (F((2) - F(X++1)).
[gxadF(Xx)12y%
r(r (
Hence
min
((x.)F(X))])
1 Y
Σ
IgradF(Xx)) <
8L n(n - 1)
(F(Xo) - F*).
K
k r((r 1)
We have by a Chernoff bound (see (Vershynin, 2018)), that for all 8 E (0,
k-1
82
Σ
(1-
2 1 - cxp
(-2 (1 - (1.(
(11)
5² (-5(1- (1_r(nn)))k)) Hence, we deduce that with probability at least 1 - exp (-
exxp
8T.
n(n - 1)
min 0= R-1 [lgrad((x.)) S h (1 5)(1 - Tnn () (F(Xo) -F)).
Finally, the proof for the permutation case is exactly similar and thus we omit here. lim inf (Xx)) = 0 almost surely.
This suggests
For the analysis under the Riemannian PL. condition, once there exists kn such that Xāo E U, then by the convergence almost surely in gradient norm and the fact that X is an isolated local minima, we conclude that there exists a subsequence Xky. for Ko < h1 < li2 < converging to X" For such a subsequence, is clear that F(X,,) converges to F(X*).
Furthermore, because FX))) converges as in (10), then F(X;)) must converge to F(X*) almost surely. By (Rebjock Boumal, 2024), we also know that quadratic growth holds (due to PL. condition), i.e., F(X²) F(X*) 2 I X 1|2 (by X" is isolated). Then we have ||_X X 112 + 0 almost surely. Thus, X₃ E U for all ko-
Next, we derive the convergence rate. If we orthogonal sampling, show the following results by induction:
< (1 (
2.or X. 4L n(n - r(r-1)
l r(r 1) )
(F(Xko) F(X*))
where Y, E {0, 1} is the same random variable defined above. It is clear at k 1, by (10) and by the same argument as above, we have
F(X)) = F(X F(X..) + F(Xxo) - F(X))
r(r - 1)
S 8L n(n 11)
(X (F()))))))))))))))))()))
-
s <(1 11-
4L n(n 1) -
" r(r
Yo (O(XEa--F(X"))
20

---

Efficient Optimization with Orthogonality Constraint
Now there exists iteration k 2 1 such that for all ls k', we have
E V2
F(Xko+k) - F(X) - (1- ((1
P r(r (0-1] 1) 4L n (1)n
((Xx))-F(XX))
Then we verify for k = k',
F(Xko+#) F(X))) F(Xko+t))
+ F FX -1) F((*)
r(r 1)
(X + ('(Kko+t-)) F((x))
8L n(nn -)
1 1r r(r 11-11 11
YK
4L n(n (11
- (F((
1)
1 P r(r 11-11
YH-1
(1-
14 r(r - En'o²%
(F(X))-F(X))
4L nn 1)
4L n(n 1
P r(r-1) E/n'' Y; 4L n(n-1)
((Xxt)-F(X))
where the second last inequality is by induction. This completes the induction. Then using similar argument, we have
Et=d Y
1e r(r - 1)
F(X F(X))
1
-P(X)))F((X3)
4T. n(n - 1)
(1-6)(1-7(rr)14
P r(r 1)
1
(P(XKo)-F(X*)
4T. n(n 1)
r(r 1) (1 - 8) (161)11(1000 ((n,r))) (F(X
exp
4L n(n
with probability least 1 exp (- 82 (1 - r(n,r))k by (11). The proof for the permutation case is the same and thus omitted.
For simplicity, we fix 8 1/2 such that the results hold with probability at least 1 1-exp (-(1 - 7(n, k))k/8) for orthogonal sampling and results hold with probability at least exp(ppppp -1 k/8) for permutation sampling.
D.5. Proof of Theorem 5.7
Proof of Theorem 5.7. By Lemma 5.2, we see Fx(Y) is L-smooth with L = Co + Cj. Then we have for any Y and W E TrO(r.))
L (Retry (W)) s FX + (gradFx W) + 2 |Wi..
Applying this inequality to the update and recalling (x(() = F(U(Y))))))) we have
1121.
F(X++1) = FR ((,))) < F(,,))
(1.))
17 2
= =({))
Fg (T.)l12
2T.
where we note that Fa, = F(Xx). Taking expectation on both sides with respect to the randomness in the current iteration, we have
1
E&F((++1) - F(Xx)
21. E& E (4.)
1 r(r 1)
=((x))
2L n(n 1) [ggradFa
21

---

Efficient Optimization with Orthogonality Constraint
where recall (u(()= R.)
Hence by Lemma 5.3, we have (lgradFii 2 1 (Xx)r and taking full expectation,
1 r(r - 1)
E[F(X++1) 1 F
4L n(n - 1) [
(12)
Hence telescoping the inequality from i = O..... k - gives
4T. n(n 1) (F(Xo) F*). r(r-1)
E ((x)
k
Then notice that min.=o.....k --1 E Elllgrad (XX)]] < E ElllgraF((x.) finishes the proof for the non-convex case.
To show the second result on convergence under PI. condition, we notice that by Definition 5.4, and that min F(X) F(X*). Then once XRo E Ut we can follow the same proof that X E U for k 2 ko. Then we have using, (12) and Definition 5.4 that
((XXE+))F((XX))) =( F(X+) + B(F(XA) F(X*))
r(r 1)
S 4L n(n 1)
B ([lgrraFF(x)) B(F(X)) - F(X))
2p r(r 1)
S 4L n(n - 1)
B[F(Xx) F(X*)] + F(X*)]
P r(r 1)
B[F(XL) F(X))..
2L n(n 1)
Let Ko be sufficiently large iteration such that Xko E u. Then, we have B[ F(X*))
exp(-
a 2L min- (X))))) F(X")], where we use (1 a)* S exp(-ake for & > O.
D.6. Proof of Theorem 5.13
Proof of Theorem 5.13. From Lemma 5.2, we know that F is I-smooth, where I. = Co C1. Then F(X++1) (Retrr, (-
T2L (I::))
(T,)) ((gradF
Taking expectation with respect to εk. we obtain
Eg. (F(x-)) S Es. [F(x))] - (T.))a + v2 R. j
where we notice (L,) F(Xx) and use the unbiasedness assumption. addition, we can bound
Es, Ilgrad (Tri (Tri grad
(lAFPPP)))) (rad)
< [ - VF(X)) + (XX) X (f(x:EE))) ? < (V VF(X.)) 0²
where we use the definition of Ja. Jw, Fe, Fx and orthogonality of Pe(r) and Xạ. The last inequality is by bounded variance assumption.
Then we further expand
Eg, Igrad (T(: Es, Ilgrad Fg (T,)) + Rg, (gradd((TT::)) Fr
< Flgrad/F (1,)) +
22

---

Efficient Optimization with Orthogonality Constraint
where we use the unbiasedness in the first equality. This gives
ER, [**(++1)]] - Ex, [P(x)] (1 (lgrad/(.. (1, + Lf²
L2L
2
Now taking expectation with respect to the randomness in Pk, we can show from Proposition 5.5 that
[x([(X&+1)] S F(Xx) - -(1- 1/2 L r(r 1)
Igrad (Tn)) n²
2 n(n 1)
2
S F(X - -11 11² L r(r 1)
+ of Lo²
2
2n(n - 1)
2
where we denote Eg to represent the expectation over randomness in iteration k and the second inequality is by Lemma 5.3. Taking full expectation, we obtain
E| [[F(X++1) - F(X:)] < - -(7 -17- 22L - r(r 1)
Lo²
EE FlggradF(Xx)] +
(13)
2 2n(n 1)
2
Rearranging the terms and summing over [K] gives
1/L r(r -1-1 1) K-1
(xa) F(Xo) - F + Kof
(1
2 2n(n 1)
2
k=
Choosing 1 = min 100" -1 K K--//// 1/23 for some constant C > Then v) - 112 2 1/2 and thus we can show
K
E FXx)) $ 4m(n - 1)
4n(n - 1)
11) OoxxLL,VK//)AA
r(r nLOo2
An(n - 1) LAo o
Loe
r(r - 1)
CVK
VK
4n(n - 1) (
T.A
+
-(r-1) 1) K
VK
where we let Ao = F(Xo) - F and we choose c = /ooLL to minimize the upper bound. Finally, noticing
min; K-1 Fx S * Σk- (Xx)) completes the proof under nonconvex loss.
To show the second result on convergence under PL condition, suppose X₁ E U. For such k. we have by (13)
( =
k+1) F|
- F(X")]
s -1 r2L L r(r -1) 1)
E ElgrrdFF(X))l + L0 +E(F(XX) F(X*)]
2n(n 1)
S (n 1 2 L r(r -1) 1)
L22
n(n - 1) ( - F(X*)]
2 -E F(X")
(1 - ul 1/ T.22) r(r 1) E(F(Xx) F(X*))
12
n(n - 1)
2
Given 77 - L-1, we obtain B[ - F[X)]] < (1 # - ) F(X*)] z..
Then for ko, ..., ki such
that X koo -..; Xa. U1 we have
[((((x.)) F(X*]
11 r(r 1)
11- E(F(X F(X*)+ +
σ² ki hi-
Σ
(1
1 r(r 100111 2L n(n - 1)
2L nnn - 1)
2L
P r(r - 1)
0² n(n - 1)
B[F(Xx0) F(X*)]
2T. n( 11)
11 r(r - 1)
11 r(r
0² n 1) 11 r(r - 1)
< exp
(k1 ko) ((F((K))) F(X"]]
22 n(n
1)
23

---

Efficient Optimization with Orthogonality Constraint
where the second inequality is by Ei a)' < 1.
E. Extension to finite-sum settings
In this section, we generalize the analysis of Section 5.3 to mini-batch settings, namely
1 Σ |(X) {F(X) =
min X (St(p))
(14)
where fi E [N] represents N component functions.
We accordingly modify Algorithm by sampling Bx C. [V] uniformly from [N] with replacement. Without loss of generality, we set [Bx] = B for all k. Then we replace the Riemannian mini-batch gradient as
gradf (1+) (Ir) = B
Σ BE
grad (4)),
where we use (Y) = f((Ux(Y)Xx)) similarly as before. Then we update the iterate with the mini-batch gradient gradfa, (I)). Notice that by the definition, we have Fx) = N E== ok (Y). Because B₁ is sampled uniformly with replacement, we can easily see that ER, Igrad (1-)] - (I.).
Before analyze the convergence, require the following assumption on bounded variance. Assumption E.1. The gradient of each component function has bounded variance, i.e., V( -VFXX)) S 0² for any E [N] and X E St(n.p)-
Theorem E.2. Consider mini-batch setling as (14) and solve with mini-batch gradient descent with batch size B. Under Assumption 5.1 and E.1, suppose choose 1 min / LBooLLo* K -1/2,, where denote Do = F(Xo) F Then we can show
min [llgradF(Xx)ll < 4n(n 1) k=
LAo
20
AoL
r(r 1)
K
VKB
Suppose there exist Xko: ..., Xx, E U for some k > kig where is defined Theorem 5.7. Then we have ((F(X o² m(n-1)
F[X*]] S exp
2L m(n-1) (ki (()( F(X*]]
1 r(r-1]
Bu r(r-1)
Proof of Theorem E.2. From Lemma 5.2, know that F is L-smooth, where L = Cu C1. Then F(Xx-1)= = Fx (Retrs, (-ngradf (T,)))
n2 Igradf (I,)112.
< Fx,) - nigrad grad (I₂))
2
Taking expectation with respect to Bg, we obtain
1 En Iggrad n, (7,)
ER, [F(X++1) < ER, [F(x,)] ollgrad
2
where we notice (T,) F(Xx) and by the unbiasedness. In addition, we can bound
En, Igraddn (T+) - (T.)]
= PA (r) gradda (4.) gradF
EB. (V.BA (X+) VF(Xx) + +X(VXX)
(X:)))
< EA, VV (Xx) - VF(X+112
B(V(()))) VF(XII - /² B
B
IEB.
where the second last inequality is by independence of samples in Bg and the last inequality is by bounded variance assumption (Assumption E.I).
The subsequence analysis follows exactly the same as in Theorem 5.13, where we replace o² with a² B.
24

---

Efficient Optimization with Orthogonality Constraint
F. On the exact convergence of RSDM
In this section, we examine the exact convergence RSDM. We require the projection matrices for the inner iterations satisfy the condition (4), which we recall below:
[PP(r)graddFAA( 112 Fa
for some constant C O. This is non-degenerate condition the selection of random matrix P over certain iterations such that projected gradient does not vanish.
Let us define by S, the set of all truncated permutation matrix P(r) E grx". To each element P of S we can associate a unique truncated permutation ı defined by
vi S r: (i) is such that PT Cx(()-
Notice that π is defined only for the first r integers as the matrix P E S has only r rows.
n! Proposition F.1. Let 5 (n-rj!' and assume that the matrices (Pr) randomly sampled, without replacement, from St (at each iteration s, we pick randomly a matrix from S that has not already been chosen) then condition (4) holds with Cp ()1((r-))
Proof of Proposition F.I. The proof follows the same idea with the proof Proposition 5.5 in the permutation case. Indeed let us fix a matrix (r) Sh with associated permutation æ. Using the fact that gradF& is skew symmetric, we have
Σ Igicjsr
(gradF
Σ 1SUJ
(L.)
Hence, we can write that
Σ Σ
(xrad/((4)))
RESE ISHJST
Notice that we obtain the same expression as the proof of Proposition 5.5 but without the factor (n-r):- Indeed, the summation on term the above equation corresponds to what we denoted by (1])radFF*4..) in the proof of proposition. Hence following the same argument, we have that
n! r(r - 1) (n r)! n(n 1)
Igrad d F5 (T.)))
that is
- ())))l(r- - 1) (n (n-r)!
(P
|graad Fi (1.)
Next, we analyze the exact convergence we sample according to condition (4).
Before we derive the results, we recall the notation that F2(Y) = F(Ug(Y)Xg) and FE(U)= F(UXg). We also require the following lemma from (Chen et al., 2020) that bounds the retraction on Stiefel manifold with the Euclidean retraction, ie.., addition.
Lemma E.2 (Chen et al. (2020)). For all X E St(n,, p) and U e St p). there exists a constant M such that |etrx (U) - X|| s M|.
25

---

Efficient Optimization with Orthogonality Constraint
Proof of Theorem 5.11. Following the analysis of Theorem 5.7, have for any k;
s s F(X))
2L (T.))
(15)
Next, recall that (T..) -F())) Then we show for any O,..., S - 1 and any k.
[ggradF;([..) gradFF. & (F()) VF(Xx)] Cil Xxll
6-1
Σ
| Xill
= C Σ
[ -
Σ
1 4.ll
s SiiA gradFF(())
(16)
where the first and third inequalities are by triangle inequality and the second inequality is by Assumption 5.1 that (Euclidean) Hessian is upper bounded. The fourth inequality is by the orthogonality of P (r) and X The last inequality is by Lemma F.2.
Then we can bound for any s,s (1
20 - gra)
FT.)1
where the second inequality is by (16). Summing over the above inequality yields
$
(P T12
22+2
Fg
(17)
Finally, we sum over the inequality (15) for s = O,... S 1, which obtains
S-1
Igad F (I,)
2L
==0
I F(X2) (L-1 3 2
< <F(X)) C,(L C²L M 3 [Iggadd
< F M
where the second inequality by (17) and 1/L. The second inequality by (4). The last inequality by Lemma 5.3. Noticing tht X Xx41 and x X2 and telescoping the result, have
2L
(x.)F(X C1 + C (F(X)) - F*),
which shows the desired result.
26

---

Efficient Optimization with Orthogonality Constraint
10t
RGD
10"
12"
ND
Optimaltyy Gap
10
RSDM-C
Optimal Gap
10
Optimality Gap
20
10
10-1
1S0M (00)
#SpM-P
20-1
RSOM: 0 (250)
RSDM- (250
10-1
8536F
153M: (400)
10-0
K50M6 (001)
150M 150M-0180 11000
ISo 115-1
RSDK i (500)
15.M F (200)
13UMP (103)
050440 23001
RSDN-O (300)
R53N-0 (200)
Optimal
R53M-0
10"1
10
20
12
24
2.5
2.5
5.0
7.5
33.0
12.5
15.
375
10
25
20
25
20
30
Time
Time
Time
Time
(a) PCA (2000, 100)
1G0 (b) PCA (2000, 300)
(c) PCA (2000, 500)
(d) PCA (2000, 800)
120
101
RGD
10"
RGD
1785)
(5365
Optimality Gai
10-1
10-1
13-2
Optimallty Gap
10-1
1892)
20-2
10
33
70
10-2
RSDMM
O (793)
Optimallty
10
RS F 1900)
152K-0 - (200)
Optimality
RID (902)
RSDM 1992)
20
RSDA-9 (800)
RSDIN
10
10
10-1
10-1
15
25
30
40
20
40
20
20
20
40
60
15
15
20
26
56
1b
Time
Time
Time
Time
(c) PCA (2000, 1000)
(f) PCA (2000, 1200)
(g) PCA (2000, 1500)
(h) PCA (2000, 1800)
Figure& 8: Experiments on PCA problem with various settings of The numbers in brackets correspond to n. respectively.
G. Additional numerical experiments on different p
In this section, we investigate the performance of proposed RSDM across various settings of on the PCA problem. Following Section 6.2, we employ the same procedures in generating the data and fix dimension n = 2000. We sweep across a range of values of p. i.e., p = 100, 300, 800, 1000, 1200, 1500, 1800. For each problem instance, we also run RSDM with different submanifold size
The convergence of RSDM in comparison to RGD is given in Figure 8. We observe that indeed when p becomes small, the performance gap between RSDM and RGD is decreasing, which is in accordance with derived total complexities. Nevertheless, we still observe that RSDM is able to outperform RGD across all the settings, except for the case when 100. This validates the benefits of RSDM even when p becomes smaller. Nonetheless, we admit that when becomes significantly small relative to n (as in the case when p = 100), RGD may perform better because the cost of retraction is much less pronounced. However, this motivates hybrid design of RSDM such that when p is relatively small, effectively behaves similarly to RGD. This is left for future exploration.
An interesting observation is that when p 300 and 500, selecting convergence especially near optimality. We conjecture this is due to the randomized submanifold descent leads to better- conditioned loss landscape around optimality, and thus performs well particularly for ill-conditioned problems. The theoretical analysis of such claim is left for future works.
can still yield significantly improved
H. Comparison to RSSM
In this section, compare our method with RSSM (Cheung et al., 2024), which can be viewed as projected Euclidean coordinate descent on the Stiefel manifold. In particular, Algorithm 1 of (Cheung et al., 2024) translates into the following update steps for smooth optimization. For each iteration k,
1. Pick index set C' C [p] with no repetition in C.
2. Compute partial Euclidean gradient and project to tangent space:
Frad,F(''' =
where Xc E [C is the columns of X corresponding to the index in C' and skew(A) = (A-AT) AT2/2 denotes the skew-symmetric operation. VxF the partial Euclidean gradient with respect to the columns of X in C.
27

---

Efficient Optimization with Orthogonality Constraint
RGD
10"
RGD
Optimaltty Gap
10
RSSM
RSDM-P
Optimality Gap
20:1
RS3M
RSDM-P
10-1
R5DM O
10-1
RSDM 0
10:
0
30
25
30
40
55
50
220
405
400
000
1000
Time
Iteration
(a) PCA (2000, 1500)
(b) PCA (2000, 1500)
Figure 9: Comparison of proposed RSDM with RSSM (Cheung et al., 2024) on the PCA problem. We observe RSDM converges significantly faster than RSSM.
3. Update the columns of X in C' by projected gradient descent while keeping other columns the same:
x = Projss(s..cc)) (X& ngrad,
where Projss(r..)) (.) denotes the projection from Euclidean space to Stiefel manifold via SVD.
Comparison to proposed RSDM. RSSM can be viewed as Euclidean coordinate descent projected to Stiefel manifold. Let r = [C] < p be the number of columns sampled. The gradient computation requires ((ppr) complexity and iterate update requires O(nr²) for non-standard linear algebra operation. Thus the per-iteration complexity costly than our proposed RSDM (with permutation sampling), which requires O(nr²) for gradient complexity and O(r3) for non-standard linear algebra operation. Thus, see RSDM-P requires much per-iteration complexity compared RSSM. In addition, apart from the advantages in per-iteration complexity, RSDM also allows easy generalization to quotient manifolds, such as Grassmann manifold, while this appears challenging for RSSM due to its operations along the columns.
Numerical comparisons. To further validate the benefits of RSDM to RSSM (Cheung et al., 2024), we have implemented RSSM with a fixed stepsize and choose C from [p] uniformly without repetition. We have tuned both r = |C] and stepsize ? for RSSM to the best performance. We compare the performance on the PCA problem where we tune r 700 and 1 0.1 for RSSM.
The results are included in Figure 9. We notice that RSDM (either with orthogonal or permutation sampling) achieves significantly faster convergence compared to RSSM. This verifies the numerical benefits of RSDM over RSSM.
I. Comparison to OBCD
This section compares proposed RSDM to OBCD (Yuan, 2023). We first remark that (Yuan, 2023) is primarily designed for nonsmooth optimization and thus optimality conditions and convergence analysis are largely different. Here we adapt the algorithm of OBCD to the smooth case.
In this case, because they parameterize X4+1 = X₂ + UB I,JUT Xr where UB E Raxr is random truncated permutation matrix, and V E R"*r, they minimize quadratic upper bound for F(X, - UB(V - J X4). Suppose F is L, smooth, then the subproblem translates into
LF VIV
min VESI(ry))
-V , +
11
(18)
for some constant Le that depends on the smoothness of F. And thus there exists a global solution to (18), i.e., V* the top eigenvectors of LF UF Un.
This is related but different particular, we update Y by
update of Y (according to our notation) when we use permutation sampling strategy. In
Y Retri,
= -eetr,, (- (XCCCXX)XXXXXXXXXXX)XX))))
It is worth mentioning that (Cheung et al., 2024) did not include any numerical experiments nor provide the code.
28

---

Efficient Optimization with Orthogonality Constraint
RGD
10"
RGD
Optimality Gap
10
0BCD
RSDM-P
Optimality Gap
2011
0BCD
10-1
R3DM O
10-1
RSDM 0
10
10
20
30
40
50
00
220
405
600
000
1000
Time
Iteration
(a) PCA (2000, 1500)
(b) PCA (2000, 1500)
Figure 10: Comparison of proposed RSDM with OBCD (Yuan, 2023) on the PCA problem. We observe RSDM converges
significantly faster than OBCD.
The key difference that we have skew-symmetric operation for (F())) which renders the update direction properly defined as Riemannian gradient (rr). In contrast, OBCD leverages the Euclidean gradient for the update.
This difference leads large deviation in the proof strategy, making the analysis of (Yuan, 2023) less aligned with common analysis on Riemannian manifolds. This makes their developments more difficult to generalize to other manifolds of interest and incorporate additional optimization techniques on manifolds, such as adaptive gradients, acceleration, Newton based methods, etc.
Apart from this main difference, we also summarize other differences of their developments compared this work:
They show convergence to block-k stationary points, which seems to be weaker than our established convergence to stationary points (as shown in Theorem 5.5 of (Yuan, 2023).
Their convergence rate in Theorem 6.3 depends on a large binomial coefficient C while our convergence has a coefficient -2
They only consider Up to be (truncated) permutation while we consider both permutation and general orthogonal matrix.
We have shown convergence in stochastic settings and shown extension to other quotient manifolds, which is not the case for (Yuan, 2023).
They only show convergence in expectation while we show convergence both in expectation and with high probability and almost surely.
Finally, we compare the proposed RSDM to OBCD numerically on the PCA problem. We have solved V from (18) analytically with SVD. We choose r = 700, which is the same as RSDM for comparability and tune stepsize accordingly. The convergence plots are given in Figure 10 where we observe that OBCD (Yuan, 2023) converges significantly slower compared to RSDM. This suggests the critical difference the update directions (Riemannian gradient for proposed RSDM and Euclidean gradient for OBCD) has led to significant convergence disparities, thus verifying superiority of the framework of Riemannian optimization employed by RSDM in this paper.
J. RSDM with momentum
In this section, we explore the potential of RSDM when coupled with momentum. We adopt the strategy of fixing Pt, for several iterations, where we apply momentum. This is equivalent to taking multiple gradient descent (with momentum for minimizing (Y) initialized from I,. The procedures are included in Algorithm 3.
29

---

Efficient Optimization with Orthogonality Constraint
Algorithm 3 RRDDM-momentum
Initialize Xu St(n.p).
2: for k == 0, K do
3: Sample Ps E (z) and let FY) (X)x( where U+(Y) is defined in (2).
4: Set v = Ir.
5: for s = 0, ..... S do
6:
Compute Riemannian gradient gradFF (Yx))-
7: Update v Retry; (-n F (Vg) + yY - 1)-
8: end for
9: Set XK+1 = Un (Y
10:
end for
is worth mentioning that we now require to compute the gradient gradFa(Y) for any Y E O(r), while previously we
only need to compute grad (T, Specifically, we compute
gradF((Y) ; -
-
where 2x = Ux(Y)X, = + Py - (.).
To test the feasibility of the proposed algorithm, we evaluate RSDM with mo- mentum (RSDM-M) with permutation sampling and compare against RGD (RGD-M) with momentum on the PCA problem. We consider the setting of 2000, 1500 and 700. We set the momentum parameter to be 0.5 for
RGD-M
Opsimality
both RSDM and RGD. We tune and set the learning rate of 0.1 for RGD-M and 1.0 for RSDM-M.
From Figure 11, we see that the RSDM with momentum improves the convergence of RSDM, which demonstrate the potential incorporating momentum into our framework.
Time
Figure 11: Comparison of RSDM
with momentum with RGD with mo- mentum on the PCA problem.
30


Reviews:
[{"review_id": "vkYyOK9ECh", "reviewer": "Reviewer_gbXw", "summary": "This paper develops a new approach to performing optimization with orthogonality constraints, with an emphasis on keeping computational complexity low. The authors are able to provide convergence bounds in expectation for nonconvex losses. The authors then run a number of experiments on well-known baselines to confirm the utility of their method.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful that the reviewer found our work interesting and well-written. Below are our responses to your comments. \n\n\n**1. On the novelty of this work compared to related works.**\n\n**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.\n\n\n**2. Add label to x-axes of Figure 3, 4, 5.**\n\n**R2**: Thank you for the question. We have already included the x-axes label for Figure 3,4,5, which is (wall-clock) Time. \n\nWe hope we have addressed all your comments and concerns. If you have further questions, we would be happy to address them.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "I did not go through the main proofs in detail, but the convergence results in 5.1.1 and 5.1.2 make sense, and are well-explained by the authors.", "experimental_designs_or_analyses": "N/A.", "relation_to_broader_scientific_literature": "This would be interesting to anyone who requires optimizing with orthogonality constraints in their research.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "The paper is nicely written, in a way that I (a non-expert in this area) could follow along and understand what was being done. I worry that this work might be too incremental, after having looked through some of the references provided in the introduction. But as I state above, I am not particularly familiar with this literature so I could be wrong about the novelty of this work.", "other_comments_or_suggestions": "On the x-axes of Figures 3,4, and 5, please add a label. In this wall-clock time? Number of optimization steps? This information is helpful to know, as most deep learning researchers have a sense of how long it takes to train e.g., MNIST, and it is important to understand the slowdown (if there is one) your algorithm induces.", "questions_for_authors": "No.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aSi8IK78tE", "reviewer": "Reviewer_ns8p", "summary": "The paper presents an efficient Riemannian optimizer for Stiefel manifolds, introducing two parameterization strategies that reduce the computational complexity of optimization steps while ensuring rigorous convergence analysis.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the insightful comments and feedback. We would like to address your comments as follows.\n\n\n\n**1. Discussions on relevant references.**\n\n**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.\n\nReference [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.\n\nWe will include the above discussions in the revised manuscript.\n\n\n[1] Cheap Orthogonal Constraints in Neural Networks: A Simple Parametrization of the Orthogonal and Unitary Group\n\n[2] Siamese Networks: The Tale of Two Manifolds.\n\n\n\n**2. More practical and diverse datasets and tasks.**\n\n**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.\n\n[1] Fei et al. O-vit: Orthogonal vision transformer. arXiv:2201.12133.\n\n\nWe 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.", "claims_and_evidence": "The paper introduces two parameterization strategies—orthogonal and permutation sampling—applied in both deterministic and stochastic settings. It establishes that Riemannian Stochastic Differential Manifolds (RSDM) achieve the same convergence rate as Riemannian Gradient Descent (RGD) when $p \\geq Cn$ and demonstrate significantly higher efficiency than RGD when $p \\leq n$. Additionally, the paper provides a detailed analysis of the trade-offs between convergence and efficiency across different parameterization methods. These claims are substantiated by rigorous theoretical results and empirical experiments.", "methods_and_evaluation_criteria": "The evaluation is well-aligned with the paper's contributions and application. Notably, the theoretical results are thoroughly developed and rigorously justified. While the experiments are conducted on a limited set of datasets and architectures, this is a minor limitation given the paper’s primary focus on theoretical advancements.", "theoretical_claims": "I have carefully reviewed all the theoretical proofs and found no major issues.", "experimental_designs_or_analyses": "I have checked the experimental designs and found no significant issues.", "relation_to_broader_scientific_literature": "While the paper introduces a Riemannian optimizer for the Stiefel manifold, its underlying ideas have the potential to extend to general Riemannian manifolds or, at the very least, a broader class of quotient manifolds.", "essential_references_not_discussed": "The paper overlooks several relevant orthogonal and Riemannian optimization methods, including:  \n\n1. **Cheap Orthogonal Constraints in Neural Networks: A Simple Parametrization of the Orthogonal and Unitary Group**  \n2. **Siamese Networks: The Tale of Two Manifolds**  \n\nIncorporating a discussion of these works could provide a more comprehensive comparison and contextualize the proposed approach within the broader landscape of orthogonal optimization techniques.", "other_strengths_and_weaknesses": "Several important methods for optimization with orthogonal constraints are not discussed or compared with the proposed approach. Additionally, the experiments are conducted on a very limited set of datasets and architectures. To strengthen the evaluation, it is recommended to test the proposed method on more recent architectures, such as Vision Transformers, and more challenging datasets, such as the Fine-Grained Visual Categorization (FGVC) datasets [3]. Furthermore, the practical benefits of the proposed parameterization are not clearly demonstrated in real-world applications that require orthogonal constraints, limiting its applicability beyond theoretical analysis.  \n\n**References:**  \n\n[3] Zhang, Y., Tang, H., & Jia, K. (2018). Fine-Grained Visual Categorization using Meta-Learning Optimization with Sample Selection of Auxiliary Data. *arXiv preprint arXiv:1807.10916.*", "other_comments_or_suggestions": "The paper would be significantly strengthened by incorporating more diverse experiments on challenging tasks. However, given its primary focus on theoretical contributions, this limitation appears to be a minor concern.", "questions_for_authors": "I have no questions for the authors.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BvR7Ya7JyA", "reviewer": "Reviewer_awSm", "summary": "This paper improves the efficiency of the retraction operation in the geometric optimization algorithms over the Stiefel manifold by constraining the optimization into a randomly sampled smaller manifold. Specifically, for an element X in the Stiefel manifold, it learns an orthonormal matrix U that acts on X, where only a random subspace in U is learnable. The computation complexity is reduced if the dimension of the random subspace is sufficiently small. The convergence of the proposed algorithm is analyzed, and experiments on four settings show some improvement in the convergence speed.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n**1. On the convergence of MNIST/CIFAR experiment.**\n\n**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.\n\nTo 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.\n\nTo 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.\n\n**2. Lack of large-scale experiments.**\n\n**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.\n\nAs 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.\n\n**3. On the terminology of \"non-standard linear algebra operations\".**\n\n**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.\n\n**4. On the orthonormal layers in Sec 6.4.**\n\n**R4**: This is a typo. All the first five layers are orthonormal. \n\n**5. The number of inner iterations used in the experiments.**\n\n**R5**: In the experiments, we implement RSDM according to Algorithm 1, which does not have inner iterations. \n\n\n**6. Reasonable range for $r$.**\n\n**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.\n\n**7. On the RSDM with momentum.**\n\n**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. \n\nWe hope our responses have addressed your comments. If there are any further questions or suggestions, we would be happy to address them.", "claims_and_evidence": "The claims are supported by evidence.", "methods_and_evaluation_criteria": "The method and evaluation make sense.", "theoretical_claims": "I did not check the proofs of the theorems.", "experimental_designs_or_analyses": "The experiment design makes sense, except for the MNIST/CIFAR experiment. The final test accuracy matters more than the \"faster convergence in terms of runtime\". The \"convergence\" in the context of neural networks can be easily influenced by learning rate schedules, and the test accuracy at the beginning of the run does not tell much about the final performance.", "relation_to_broader_scientific_literature": "The paper relates to the geometric optimization algorithms on matrix manifolds.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "The paper is well-written, and the presentation is clear. The idea is straightforward, and the results look competitive compared to baselines. The weakness is the lack of large-scale practical applications. All the experiments are either toy problems or applications on tiny-scale neural networks.", "other_comments_or_suggestions": "1. Is the \"non-standard linear algebra operations\" a well-known terminology? SVD and QR decompositions feel quite standard.\n2. In Sec 6.4, a six-layer network is used. The first four layers are column orthonormal, and the output layer is unconstrained. How about the missing layer?", "questions_for_authors": "1. What is the number of inner iterations used in the experiments?\n2. What is the reasonable range for choosing $r$, especially, how low could the ratio $r/p$ be?\n3. Since the random matrix $P_k$ changes every iteration, the momentum for matrix Y does not make sense, correct? I think the inability to use momentum in RSDM is a drawback compared to RGD. Furthermore, it will be more interesting to add RGD with momentum to the baselines in the experiments.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "qNfvzhqWj5", "reviewer": "Reviewer_N9y9", "summary": "In the work, the authors propose a randomized Riemannian submanifold approach for optimization on Stiefel manifolds. The authors prove that the proposed method converges under certain conditions.\nEmpirical results are included to support the effectiveness of the proposed method.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n\n**1. Whether the proposed idea requires a compatible Riemannian metric to effectively reduce the computational cost.**\n\nWe 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.", "claims_and_evidence": "The claims are sound and solid.", "methods_and_evaluation_criteria": "The proposed method makes sense. The authors properly evaluate their method on several problems.\nThe empirical results are convincing.", "theoretical_claims": "Convergence analysis is also provided to support the soundness of the proposed method.\nI do not check the proofs in the appendices.", "experimental_designs_or_analyses": "The experimental designs make sense.", "relation_to_broader_scientific_literature": "Optimization", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "This work is interesting and original.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "It seems that the proposed method requires an efficient computation of a Riemannian gradient on a submanifold.\nThe authors should comment on whether the proposed idea requires a compatible Riemannian metric to effectively reduce the computational cost.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a randomized Riemannian submanifold method for optimization on the Stiefel manifold and studies its convergence behavior under various assumptions on the objective function. While the reviewers find merits in the contribution, they find the numerical experiments somewhat limited. Also, there is no discussion on what applications will give rise to functions that satisfy the Riemannian PL condition, and so the relevance of the condition is unclear. Lastly, as suggested by a reviewer's comment, the term "non-standard linear algebra operations" is inappropriate. The examples cited by the author(s), i.e., matrix decomposition and matrix inverse, are by most account standard. Moreover, the computation of the matrix exponential has been extensively studied for decades; see, e.g., Higham: The Scaling and Squaring Method for the Matrix Exponential Revisited. SIMAX 26, 2005 and the references therein. Thus, the author(s) should simply state that the linear algebra operations they use have lower complexities than those used in other works, if that is indeed the case. Overall, the author(s) should take into account the reviewers' comments when they revise the paper.
Decision: Accept (poster)