Momentum Stiefel Optimizer, with Applications to Suitably-Orthogonal Attention, and Optimal Transport
Authors: Lingkai Kong
Yuqing Wang
Molei Tao
Venue: ICLR 2023 poster (2023)

Abstract: The problem of optimization on Stiefel manifold, i.e., minimizing functions of (not necessarily square) matrices that satisfy orthogonality constraints, has been extensively studied. Yet, a new approach is proposed based on, for the first time, an interplay between thoughtfully designed continuous and discrete dynamics. It leads to a gradient-based optimizer with intrinsically added momentum. This method exactly preserves the manifold structure but does not require additional operation to keep momentum in the changing (co)tangent space, and thus has low computational cost and pleasant accuracy. Its generalization to adaptive learning rates is also demonstrated. Notable performances are observed in practical tasks. For instance, we found that placing orthogonal constraints on attention heads of trained-from-scratch Vision Transformer (Dosovitskiy et al., 2020) could markedly improve its performance, when our optimizer is used, and it is better that each head is made orthogonal within itself but not necessarily to other heads. This optimizer also makes the useful notion of Projection Robust Wasserstein Distance (Paty and Cuturi, 2019; Lin et al., 2020) for high-dim. optimal transport even more effective.

Published as a conference paper at ICLR 2023
MOMENTUM STIEFEL OPTIMIZER, WITH APPLICA- TIONS TO SUITABLY-ORTHOGONAAL , AND OPTIMAL TRANSPORT
Lingkai Kong, Yuqing Wang, Molei Tao
Georgia Institute of Technology (1kong75, ywang3398, mtao) ) saatecch.eee
ABSTRACT
The problem of optimization on Stiefel manifold, i.c., minimizing functions of (not necessarily square) matrices that satisfy orthogonality constraints, has been extensively studied. Yet, new approach is proposed based on, for the first time, an interplay between thoughtfully designed continuous and discrete dynamics. It leads to gradient-based optimizer with intrinsically added momentum. This method exactly preserves the manifold structure but does not require additional operation to keep momentum in the changing (ca)tan space, and thus has low computational cost and pleasant accuracy. Its generalization to adaptive learning rates is also demonstrated. Notable performances observed in practical tasks. For instance, we found that placing orthogonal constraints on attention heads Vision Transformer (Dosovitskiy al.) 2020) could markedly improve its performance, when our optimizer is used, and it is better that each head is made orthogonal within itself but necessarily other heads. This optimizer also makes the useful notion of Projection Robust Wasserstein Distance (Paty and Cuturi 2019
Lin et 2020) for high-dim. optimal transport even more effective.
Code: httpt//sss.tttttaaccco///o//oola
INTRODUCTION
Matrices that satisfy orthogonal constraints play important roles in various areas including machine learning. A range of studies showed this both theoretically (Saxe et al. (2013) Xiao et al., 2018) and experimentally for example, orthogonality
boost the performances of many architectures" e.g.. MLP (Cisse et al.] 2017), CNN and ResNet (Li et al. 2020; Wang et al. 20200, RNN (Arjovsky et al.
2016 Wisdom al. [2066] Helfrich et al., 2018). Transformer (Zhang et al. 2021). The training of such models amounts to optimization under orthogonal constraints, whose applications to machine learning, however, are not restricted to just improving deep learning models. For example, such optimization helps construct optimal low-dim. projection of high-dim. data, which can be used for, e.g., robust and efficient approximation of Wasserstein distance (Paty and Cuturi 2019 Lin et al.. 2020). Therefore, this article will focus on this non-Euclidean Riemannian optimization problem.
In fact, matrices that contain orthogonal columns constitute a Riemannian manifold known as the Stiefel manifold. Given integers 2 m > 0, it is defined as St(ne, m) == (X E _2 X X ] (see Apdx. B.3 for the n = m special case which almost (0()).. Then, given C function f St(n,, m)) - R, we consider, under gradient oracle, the smooth optimization problem min f(X).
which equivalent to XEX=xm__X X= s(X).
min
(1)
X XESt(n,,)
Due to the importance of this problem, a wide range of approaches have been proposed. Methods based on regularizers, which approximate the constraints (and hence the manifold structure), for example, have been popular in the machine learning literature (Cisse et al. 2017, Bansal et al. 2018 Wang et al. 2020) Zhang et al. 2021). Meanwhile, efforts have been continuously made to construct algorithms that truly preserve the constraints (i.e. the manifold structure), although challenges such as computational scalability and how to add momentum still remain; see the 2nd next paragraph for how our method addresses them, and more discussions existing milestones in Sec. 1.1
Our strategy toward good Stiefel optimizer is the following: first, formulate a variational principle and use that construct an optimizing dynamics in continuous time, which described by a system
It helped computer vision applications prior the deep learning era well (e.g., Liu et al. 2003

---

Published as a conference paper at ICLR 2023
of ODEs corresponding to damped mechanical systems on constrained manifold; then, design a delicate time discretization of these ODEs, which yields optimization algorithms that precisely preserve the constraints and mimic the continuous dynamics.
Our optimizer has several pleasant properties: 1) It is exactly preserving the manifold structure, not only of the Stiefel manifold, but in fact of its tangent bundle. In other words, throughout the course of optimization, the position variable remains exactly on the Stiefel manifold, and the momentum variable remains exactly in the (co)tangent space. 2) Typically, in order to maintain the manifold map operation is needed, and since we
structure, some kind of prjjectionretaatiin//ppooooo have both position and momentum, such operation is needed for both variables (i.e. maintain the cotangent bundle structure). However, our carefully designed ODE and its discretization make the structure preservation of momentum automatic, meaning that no extra operation (projection, retraction, parallel transport, etc.) is needed for the momentum variable. This not only leads to improved computational efficiency, but also serves an indirect evidence of having a reduced overall (i.e. both position and momentum) local error. 3) We used a for our specific position retraction operation, which makes it fast. 4) Due to 2)+3), our per iteration computational complexity, (nmm²), has a small constant factor (see Sec. C for details). 5) Our discretization is also numerically stable so that it well preserves the structure even under low machine precision and numerous iterations, which are beneficial in machine learning contexts. 6) Because algorithm is derived from a variational framework that unify both Euclidean and Stiefel variables, the same hyperparameters can be used for both these parameters; see Sec SecBaan and note this difference from previous milestones (e.g.. Li et al.] (2020)) significantly reduces tuning efforts. 7) Our algorithm works for a range of Riemannian metrics, allowing extra flexibility in choosing suitable geometry to optimize the performance for a specific problem.
iterative solver
Selected (due to space) experimental tests of our optimizer are: (1) We consider the simple problem of leading eigenvalues, which is yet practically important in data sciences. It systematically investigates algorithmic performances under different parameters. (2) We show the elegant idea of approximating optimal transport distance in high-dim. via a good low-dim. projection (Paty and Cuturi, 2019 Lin al. 2020) can be made even more efficacious by our optimizer. (3) We note that Vision Transformer (ViT) can be further improved by imposing attention heads to be orthogonal; more precisely:
Consider training VIT from scratch. We discover that 1) requiring each head to be orthogonal in itself improves both training and testing accuracies the most. An important recent work by Zhang ctt al. (202I) applied orthogonality to transformers and demonstrated improved performance in NLP tasks. It concatenates each of the W.O W.K, WY matrices in attention across all heads, and applies orthogonal constraint each of the three, via regularizer. This makes each head (approximately) orthogonal, not only within itself, but also to other heads. Orthogonal constraint is also applied, via regularizer, to each weight matrix of ford-forward layers in their case. With our Stiefel optimizer which are not restricted to square matrices, we can now make each head exactly and only orthogonal within itself, which leads to further improvements orthogonality both in and across heads is found less effective than 1), but it is still better than requiring no orthogonality (i.e. vanilla ViT). No orthogonality on feed-forward layers was used in either 1) 2). In addition, 3) to achieve these improvements, our Stiefel optimizer needs to be used; methods that do not have momentum or not exactly preserve structure (e.g., regularizer- based) are seen not fully exploiting the benefit of orthogonality.
least in CV tasks. Meanwhile, 2) having
1.1 MORE ON RELATED WORK
Orthogonality in deep learning Initialization; Orthogonal initialization is both theoretically and experimentally shown good for deep neural networks (Saxe et al. 2013). CNN: Cisse et al. (2017) used regularizer to make weight matrices orthogonal when training MLP and CNN, and showed both improved accuracy and robustness. Rodriguez et al. (2017) showed that posing orthogonality to CNN reduces overfitting. Bansal et al. (2018) experimentally compared several orthogonal regularizers on ResNet. RNN: Arjovsky et al. (2016). Wisdom et al. (2016); Helfriihh et al.) (2018); Vorontsov et al. (2017): Lezcano-Casado and Martinez- Rubio (2019) showed orthogonalilty helps plain RNN achieve long term memory and even out-perform advanced models such as LSTM. Transformer: Despite of its success, transformer is actually a recent model; to the best of our knowledge, the community just started applying orthogonal constraint to Transformer for NLP tasks (Zhang et al. 2021). Our work is the first to apply it to Vision Transformer (Dosovitskiy et al. 2020).
Optimization algorithms Manifold optimization is a profound field. General approaches for 1st-order (i.e. gradient based) optimization on (Riemannian) manifold typically involve retraction
2

---

Published as a conference paper at ICLR 2023
(see Apdx Dj which is not present in Euclidean cases. Thus we categorize 1st-order Stiefel optimizers into 2 types: O1) retraction-basee based, and O2) non-retraction-base based.
01) Retraction- Given a point and a tangent vector, exponential map gives how the point moves along the tangent vector on the manifold. In practice, approximation of the exponential map, known as a retraction, is often used for reducing the computational cost. The exponential map and various retractions
Stiefel manifold well-studied. For Stiefel, the exponential map can be computed with complexity O (nm²² (Edelman et al. 1998), whose constant prefactor is however large as the map is essentially matrix exponentiation. A well-known retraction is Cayley map (Eq. (32)) and Wen and Yin (2013) proposed a gradient descent method with smartly lowered computational cost of Cayley map (will be referred to as unless Stiefel (S) GD). SVD can also help construct a retraction: Li et al. (2019) did so by first following the gradient out of the manifold and then projecting back via modifying the singular values, which is interesting but expensive. Meanwhile, a remaining challenge is to add momentum, which is nontrivial in manifold setting but helpful for improved speed of convergence (Zhang and Sra 2018 Ahn and Sra 2020: Alimisis et al. 2021). In this case, the geometry that needs to be maintained is that of the augmented state-space, known as the tangent bundle; i.e., position has to be in Stiefel manifold and momentum in its tangent space. One could follow a general Riemannian retraction idea (e.g.. Absii and Malick (2012) and use retraction that projects both position and momentum to the tangent bundle to ensure manifold preservation. This was done for example in Li et al. (2020) (will be referred to P Projected Stiefel SGD / Adam). Another cool technique to maintain the structure of momentum is parallel transport (e.g.. Béeigneul and Ganea (2019)), which is however computationally expensive and possibly inexact (Edelman et al.) 1998). J. Our work could be viewed as a retraction cetraetion-based approach, because our position variable is projected back to St at each step; however, doesn't 't require any retraction or projection on the momentum variable - thanks to our variational approach and delicate discretization, position and momentum intrinsically move together, and once the position X is placed on St, the momentum Q automatically finds itself in StSt. This leads to better computational efficiency, as well as improved accuracy because variables stay closer to the tangent bundle.
02) nonn-reraction-basee In some sense one can always call an exact manifold preserving method retraction- -bsedd but approximate structure preservation is also a possibiliity. Regularizers for instance could be used to approach X X = J. Between exact and approximate structure preserving, which one works better is theoretically still an open question, and task-dependent but insightful empirical results exist. For example, regularizer-ban baaed methods usually need more hyperparameter tuning, and the final result is likely sensitive to the regularization strength (Wang et al. 2020). Lai and Osher (2014) also discussed their slow convergence and possible improvements. Regularizer-ban typically have low computational costs, but due to trajectory not exactly on the manifold, they may converge to (the neighborhood of) a local min. different from the one an exact algorithm converges to for multimodal problems. Similar observation was made for non-regularizer-baaed approximate method too; e.g.. Vorontsov et al. (2017) suggested that at least in certain cases 'hard' constraint is better than 'soft'. The aforementioned Li et al. (2020) is actually inexact too, because in their practical implementation the retraction was inexactly performed for the sake of speed. Mcanwhile, we also note an interesting, non-regularizer-bon based recent result (Ablin and Peyré 2022) that converges to the manifold despite of being approximate prior to convergence.
methods
Additional. We also mention the existence of useful algorithms for electronic structure calculations that can be viewed as Stiefel optimizers (e.g.. ZZang et al. (2014): (y Dai et al.) (2017): Gao et al. (2018) 2019) Hu al. 2019);] Dai al. (2020)). They did not use momentum. Another inexact-manifold preserving but interesting result is Bu and Chang (002). The problem can also be approached via a constrained optimization formulation (e.g.. (Leimkuhler et al.) (021)). In addition, there are other interesting optimizers also based on discretizing damped mechanical systems but not via the splitting approach used here, e.g., [eee et al. (202T); Duruisseaux and Leok (202T); however they are mostly implicit and not always scalable to deep learning applications. Finally. it is possible to represent matrix by other matrices and optimize those matrices instead. A special case of 50(n) was considered by Lezcano- Casado and Martinez-Rubio 2001);; Helfrich et al. (2018) where X- = expm( (A) or X Cayley(A) was used and then one optimized skew symmetric A instead. Although these authors did not consider Stiefel generalization, this could be possible, but even without the generalization the computational cost is already high.
"Their setting was deterministic, but extension use stochastic gradient descents is straightforward.
3

---

Published as a conference paper at ICLR 2023
This paper considers exact structure preservation. For applications included here, this is either empirically more beneficial (Sec 3.2) or a necessity (SeeB.1) & Apdx ( See comparisons of most
relevant approaches in Tab. 2 and more on complexity in Apdx. C
2 DERIVATION OF THE OPTIMIZERS AND THEIR PROPERTIES
We represent Stiefel manifold as St(n. m) := {X E R"xm, XXX X = Imxm) where n 2 m i.e. matrices with orthonormal columns. This is a natural embedding in Euclidean space. Based on identity isomorphism between cotangent & tangent spaces (see Apdx. [.2). it also gives: Proposition 1.
The (co)tangent space of St at X E Tt := {A R x
0}: ; its taagent bundle is TSt = (X, : X E St, E St}.
Throughout this paper, we focus on family of Riemannian metrics on St defined in the following. Definition 1 (canonical-type metric). For fixed constant a 1, the canonical-type metric 9x St x St + R is defined as
gx gx(4,, (A1, 42) Tr axx)) AA A2 E St.
The following are two commonly used examples of the canonical -yype metric (Tagare 2011
When a 0, 0. A2) = TrATT A2), which corresponds to the Euclidean metric.
When a = (A,,, 22) = Tr - x )A2), which is known as the canonical metric.
of == ox (x) Denote by O Ø, Notation. Denote Euclidean gradient in ambient space by ax
iy" and A element-wise product, division, and cth power. Denote expm to be the matrix exponential.
2.1 OPTIMIZATION DYNAMICS IN CONTINUOUS TIME
In this section, we study the Stiefel optimization problem I from a variational perspective. The manifold setup yields constrained variational problem, which is hard handle due to nonlinear geometry of the function space. Thus we introduce an alternative variational formulation based on functional Lagrange multiplier bypass this difficulty. generates an ODE that never escapes from the manifold TtS and is guaranteed to converge to a local minimizer of the function f on St.
In Euclidean space, momentum has long been used accelerate gradient descent (Nesterov) 1983). More recently, continuous time limit brought useful tools to complement the classical analyses for such optimizers (e.g., Su et al. (2014). and variational formulation was then established unify a large family of momentum-basee can also provide a powerful and intrinsic way to generalize momentum-basee optimization to non- Euclidean settings. One demonstration is, for example, accelerated optimization on Lie groups (Tao and Ohsawal 2020). Following the same line of thoughts, we consider Stiefel optimization by first defining deeendennn Lagrangian L 7St R,
Euclidean optimizers (Wibisono et al. 2016).. This formulation
R as
T(X,X, = r(!) 59 [0555 (X, X) s(x) r()) (X ( -(x)]- (3)
Without r(t), this would be standard Lagrangian for mechanies manifold, and it would generate dynamics in which the kinetic energy 4. (X, X) and the potential energy s(X) continuously exchange with each other. Choosing r(1) be increasing function, however, will break time translational symmetry and introduce dissipation in the system; consequently, the total energy 1 +X+ will monotonically decrease in time, leading to the minimization of f. Given this Lagrangian, can define variational problem (VP) on manifold whose corresponding stationarity condition (known as Euler- Lagrange equation) yields an ODE that optimizes f in continuous time
Constrained VP: 5 /"" (((X(1, ((),, = 0, s.1. (X(1), (() E TSt, V T. (4)
However, to obtain concrete algorithms, this constrained VP is very difficult to handle because the variation of X x-)) has to keep in a nonlinearly constrained function space. Tao and Ohsawa (2020) applied tool of reduction via Lie algebra solve the case of Lie group manifold, but for Stiffel manifold St(n, m). unless n = m, there is no group structure and this technique no longer applies. Therefore, we transfer 44 into an unconstrained variational problem, using tools described in, e.g.. Chen et al. (2021). By adding a Lagrange multiplie E A(() E mmmm for the constraints
(() XTI) 0. we augment the Lagrangian to be
_ Tr ( (xT (xxT)x) - f(x) - (X D)), (5)
1
((t,
"The problem is harder when n. >> m. which thus will be our focus. For the n = m case, see Sec. 0 'I differs from the seminal work (Wn and Yin 2013) as our sctup (and the Lagrange multiplicr) is dynamical.
4

---

Published as a conference paper at ICLR 2023
where A(t) Rm is symmetric matrix. The benefit is that now X can be varied in an unconstrained, flat function space, corresponding to
Unconstrained VP: 8 ((X(0))(( 0. 0A st.
(6)
To sum up the above discussion, problem (T) can be resolved through (un)constrained variational
problem (4) (6. i.e., [I]
(6) in the sense that the solutions of ( solve problem
meanwhile, (6) can be explicitly solved, as detailed in the following (proof is in Apdx. G):
Theorem 1 (Optimization dynamics on TSt). To use problem (6) for solving problem ①
we have:
1. The solution of the unconstrained variational problem 66 is the following ODE
x== = Q
(7)
( -10-x0 - 30 (1- 11I-XX XX XQQ af ax + 14 XX of 8x 2 ax x
x as
momentum, 7(1) ((()-=())rr()
where (())) (X(t). x()) E R=xm x R is the tuple of state variable and its constant
Commentum.700
r(1)/r(1) is the scalar friction coefficient, and = a-I =
depending on the canonical-type metrie (2).
2. For any isolated local minimum X. E St of f, there exists a neighbourhood U c TSt of (X., 0)
s.t., for any initial condition (Xo, Qo) U. the solution (X(1), Q()) of the system converges
to (X., 0) as + CO.
One feature of ODE (7) is the preservation of constraints: even though the ODE is defined in the Euclidean space, as long as it starts on the manifold, it will never leave the manifold:
Theorem 2 (Constrained optimization with unconstrained dynamics). If the initial condition of 7 is on TSt, the cotangent bundle of Stiefel manifold, then dynamics 7 automatically stays on TSt, i.e., if (0) ( I mm ((0) Q(0) Q(0) (0) Omxmm then for all / >
(t) (t) = Imxms (t) T Q(t) + Q(t) X (t) 0mxm
Proof is in Apdx. H
2.2 STRUCTURE- PRESERVING DISCRETIZATION VIA VARIABLE DECOMPOSITION AND
OPERATOR SPLITTING
Although the continuous dynamics 7 preserves the constraints, such preservation is in general not guaranteed when time is discretized. The construction of our discretization briefly follows four steps: a geometric decomposition of momentum, a carefully designed operator splitting scheme to approximate the ODEs, prructure-prssevvv preserving relaxation that further reduces the computational cost. Details now follow.
approximations of the split system, and structure-
Preparation: from static decomposition of the tangent space to decomposition of Q dynamics. To retain the preservation of geometric structures through a time discretization, we first decompose the tangent space Tx St into X and X components (see more details in Tagare (2011)); more precisely, given a tangent vector Q represented by an X m matrix, we rewrite it as Q XY + V for Y IR mm V IR mx and use Y, V to replace Q. This transformation changes the constraint XTQ + x = 0 to {YT +Y = 0, X T = 0} instead (see Apdx F. Another advantage of
this decomposition is that Y, V naturally split the canonical-type metric (2) (see E.2).
Although a fixed tangent Q can be uniquely decomposed into and V. when we start to consider a dynamical Q(t) and make the decomposition at each t, we need to understand how the corresponding (t) and (t) evolve. This is not trivial question, but can be proved that the new dynamics is 0f
X XY
Y= -
1 -
X
Ōf
x).
2
0X
0x
3a - 2 VY - VV (1- x
0f
(8)
2
0x
with initial condition satisfying X (0) (() (.Y +Y 0, x ( (V = 0. This system
is equivalent to 7 via Q(t) = X X + (t) and preserving (() (t) = Y (t +Y()) 0, (t) (t))= 0 for all t (Thm&)). Moreover, it will be amenable good discretization and is thus the base for constructing our optimizer.
With this decomposition defined, we can finally make the phrase 'structure preservation' precise: Definition 2. Structure preservation means variables exactly satisfying all constraints for all time. That is, for 'XQ'- -syystmmm X(t) x(t) = (t) Q(t) = 0, Vt; for 'XYV' 0. Vt. In comparison. staving exactly on Stiefel, iec. ( xt)) = I is termed 'eaasibe' (Wen and Yin 2013 Ablin and Peyré 2022). but
((tt) [t) + Y(t)
X(t) ( V
have additional constraints due to momentum.
5

---

Published as a conference paper at ICLR 2023
Step I: operator splitting of the ODEs. To handle the high nonlinearity of ODE ( and maintain the preservation of constraints after discretization, we adopt an operator splitting method, based on general fact that a numerical discretization of ODE can be obtained by composing the (approximate) flow maps of split ODEs (Mcl achlan and Quispel 2002). More precisely, we split the vector field of (8) as a sum of three vector helds, each associated with one of the following ODEs: -XY Y
X
X V Y 0
-
Y 0
ط-b x 01
OL 8X
x
(9)
V -VV + 3
3y=
(10)
(II).
V -VVVV
=
-(I - xx
Define the corresponding time-h evolution maps Φ1, 22 d3 of Eq. ( PO to be 0e
(X(,), (t), [(t) + [X(t + h), (t + h), ((t + h)] for system j = 1. 3. Note 01 02. ds give the exact solutions of these split ODEs. Then
specific split honors all constraints:
Theorem 3. Φ1, 22 Φ3 are all structure preserving.
Proof is in Apdx. J
approximation of exact flow maps. Due to the nonlinearity, Φ1 and Φs do not admit analytical expressions; وय does have explicit expression (Eq. (30), but an approximation will still reduce computational costs while maintaining certain accuracy (see Fig. Therefore we first denote the sst-order approximation of the exact flow maps Φ1, t22, d2 to be 11, d2; da. where By [Xo; Yo, Vol + [X,, Yn, Vn] = (x([No, Yo, Vol) + ((2)).) = 1, 2,3. Then
Step II: structure-preserving
X = Xo expm)
X = Xo + hV Xo
Yn exp(-nh))o
X = X, x1) -1/2
[1-b)(1-exp( -h1)
(XT Bf
Xo)
(12) 03
YE = Vo
(13)
27
8Na
aXa
Vn Vo
Vh = Vo Tu-han -oooVo
3a 2 hVoo - h (T XxXX,
OS (Xo), Yn = Yoo (14)
2 X4. = No, V/2 = 0(1 Vh +
2
ox
Theorem 4. Φ1, 22, 03 are all structure preserving.
Proof is in Apdx. K
This shows not only the split flow maps but their decently designed discretizations maintain the constraints of X, YYY V. There are several specially designed features to enable the numerical stability of the integrators: 1) (X, X2)-1/2 in Ø3 is part of a nontrivial discretization scheme and the same as polar retraction (Absil et al. 2009). directly leads to the preservation of the geometry of the position variable by Ø3. i.c., even the input of Φ3 has error that X Xo * I, the output always satisfies X, X2 = I, but it will not impair the order of the ( local discretization error. See more about its connection to numerically stability to arithmetic errors in Apdx. [M.] 2) In the first equation of Φ3. X X, Xo used instead of Xo + Vo, even though most of the time X, Xo = I. This guarantees that even if X Xo # I, the constraint X, Vo= = 0 itself still leads to X, V.. = 0, which improves numerical stability. What's more, when combined with 1), this property will also enable us to use a cheaper ei in the following Step III. 3) The 'forward Euler"-like discretization for d3 is carefully chosen,and updating Vn, using X₁, instead of Xg. for example, will destroy its structure preservation. 4) No extra handling except forward Euler is applied to the momentum variables Y, V; Le. this discretization itself is beneficial enough to guarantee momentum structure preservation.
Step III: relaxation and composition of the operators. Our goal is to obtain structure preserving scheme for the original ODE (8) instead of requiring structure preservation of each operator for the split ODEs ( 1T. The latter will ensure the former as we have:
Theorem 5. The composition of any ordering of Φ2, Φ3. e.g., おi 062, is a structure-preservingg Proof is in Apdx. L
Ist-order (in h) numerical integrator of (8).
However, the latter (all operators preserving structure) is not needed for the former (their composition preserves structure). We can still eventually obtain a structure- presrevine integrator, however without the costly 'expm', by relaxing some of the structure preservation requirements for ø1.
Theorem 6. Consider a consistent approximation of Φ1 by Φr (i.e. क = σ1 + O(h))) Assume Φi satisfies that if initially X_ Xo = 7, YoT + Yo = 0, XJ Vo = O. then Y, + Y, = 0, X.X,,.V V. = O. Then the specific composition Ø3 Φ1 σb is structure preserving. Moreover, the composition of any ordering of Φ1, $2, dh is a Ist-order (in h)) integrator of (8).
Proof is in Apdx. M.3
Therefore, our default recommendation is to use da o Φ1 22,, where H1 is X, = Xo Yo. 0t Yn = - (1- - (XT Exo aXo Xo) Vn = Vo. mainly approximating the expm in ~1
as

---

Published as a conference paper at ICLR 2023
by forward Euler. The result is concretely summarized in Algo. In order to match commonly used notations in machine learning, we rescale the parameters: learning rate 17 == 1-exp -h)
momentum parameter μ == exp(- Z 1/1-xxp(- 111-xxp(-ph) and U Apdx M.4 for more information.
1 -xp( See
A few side remarks are listed in the following. If m is small, then क1 o or Ø3 (or any permutation) can also be used, although experimentally no clear advantage was observed (Fig. otherwise, the computational cost of expm can become prohibitive.
Moreover, the computational cost of matrix inverse square root in the polar retraction X+) should not be concern since it is only for X matrices instead of n x (n m). Meanwhile, it can be computed to machine precision rapidly using a quadratically convergent iterative method for computational complexity.
(Higham 1997); see Algo. in Apdx. M.I for details, and Apdx [
Algorithm I: Momentum (S)GD on St(n,, m) (SGD: /0x replaced by stochastic estimator) Hyperparameter::/ :.) E (0, +00), / E [0, 1). maximum number of iterations N
Initialization :Xo. Uo. Zo s.t. X Xo = 1, x, Uo = 0, Zo + Zo =0
for i = 0, .N - 1 do
Compute 'gradients': : fi 1-b
X.T
0f (X:)
X.X,) af T (x)) = -I- (x
2
8X
8x
Update 62: U = pl;: 2 n1UZii -9i
Update क1: Zi+1 uZz - fú: Xi+ 3 = X; + n2
Update Φ3: X; = Xit3 + +nU++ XTT 141 x Compute (X, x1)- using Algo.
X Xx, X+) :; 21U11 U1+1 nxi+tU UT 14 Us+
end
return XN
2.3 AN ADAPTIVE LEARNING RATE VERSION
Tuning for the best SGD learning rate can be labor intensive and computationally unafforable, and sometimes SGD even performs worse than adaptive methods (Zhang et al.) 2020). Thus in this section, we propose an adaptive version of our Stiefel optimizer. More precisely, we will establish, as an example, a Stiefel version of Adam (Kingma and Ba 2015). which estimates the 1st and 2nd moments of gradients to obtain element wise adaptive step sizes.
The algorithm is established via the following ideas. The 'gradient' in this Adam-version method is constructed from Stiefel SGD with momentum [ where the 'gradients' in Y/Z and V/U X) and (1 ( X x)- af xx respectively. The main difficulty of extending Stiefel SGD to Stiefel Adam that does not appear in Euclidean case is that element-wisg operation on momentum loses tangent vector structure. We solve this respectively: (1) For Y/Z-direction the skew semmeeyyyy is preserved after a symmetric element-wise operation Z 2 (p + +). (2) For dir-directio,, we apply projection I - X(xTx)=Xx* - to the element- clement-wisely rescaled momentum U @ (qº) E), making sure 'X V O'. Combining all the above, we obtain the Adam-Stiefe optimizer. Denote ø1, d2, 63 be the modification of 61, d2, d3 in the Adam version (see detailed expressions in Apdx. M.5). Then the integrator is defined as &o Φ1 2b.. The overall method is shown in Algo.
direction can be interpreted as ( of
0f
#X
BX
Theorem 7. The Adam-Stiefel de Φ1 02 is structure-preservingg.
Proof is in Apdx.
N
3 EXPERIMENTS
This section demonstrates our Stiefel optimizers on Projection Robust Wasserstein Distance (Pat and Cuturi 2019 Lin et al. 2020) and fromsss-fro--cratcc Vision Transformer (Dosovitsiiy et al. 2020). They will also be compared with other popular Stiefel optimizers summarized in Tab. 2 Canonical metric (i.e. 1/2 in Eq is used in both examples to show that the gained performance is due to algorithmic innovation but not an extra tuned knob. An additional experiment on leading eigenvalue problem is deferred to Apdx. Q where also compare the convergence rates and time consumptions of different algorithms, test various choices of metrics, and study why our retraction is tailored to our algorithm. Good performance is observed in all three examples.
7

---

Published as a conference paper at ICLR 2023
Algorithm 2: Adam on Stiefel manifold m)
Hyperparameter : E (0, +o0), B1 E [0,
1, number of iterations N
Initialization :Xo, Vo. Yo. Po. 4o s.t. X 1, Uo =0 0, Zo + Z 0., Po = Po
for i = o,... ,N 1 do
Compute 'gradients": S, 1 (X, 0f
0x (X,) of (
-I-IIXXX0 x,,,
of (X,)
8x
2nd-moment estimation: Pi+1 = Bop; + (1- 9:41 $2q: + + -2)9,²
Update 62: U = 31U,- - 2 2 - (1 (1- 31)g
Update 61: Zi+1 = 312z - (1- 81)
X X; m/ 32+1Xg 1x, (Zit1 O (pit) + +))
Update dz: Ú = / 1 - 1351(( (T X (X x ) - (4 (UL+ @ (dif1 + +c);
143
X X + n X. [t x X,- = X, (X; x)) x UH11 U nx+t U U
end
return XN
PRW distance between MNIST digits
PRW distance between Shakespeare plays
PRW/ distance mean
0.8
PRW distance mean
0.6
0.900
0.875
0.15
0.190
0.188
0.850
0.186
0.10
0.184
Strrtel
0,800
0.2
Nomennumies
48
Projected
0.05
0.0
10
20
30
40 50
10
20
30
40
50
Iter
iter
Figure I: Projection Robust Wasserstein Distance (PRW) tested on MNIST and Shakespeare plays. Data points are features extracted by a pre-trained model. The mean optimal transport value is taken among all digits or movie pairs; larger mean optimal transport value means more effective orthogonal projection. Our method makes PRW more effective by getting the best local minimum (largest optimal transport value) and fast convergence.
More experimental details are in Apdx. P
3.1 PROJECTION ROBUST WASSERSTEIN DISTANCE
Projection Robust Wasserstein Distance (PRW) (Paty and Cuturi 2019 Lin et al. 2020) is notion recently proposed to improve the robustness of standard Wasserstein metric, especially in high- dim settings. The idea is simultaneously look for a best projection from high-dim data (x,, Mj+ respectively with weights ri, (5) to low-dim ones, and an entropic regularized optimal transport plan between projected data, i.e.
max (ESt(d,)) TER
min
PusilUUe x. UT wsll + TT(m, log(T)) - 1,,,
E. May
This problem is geodesically nonconvex w.r.t. the Stiefel variable U U'JJoraa et al. 2022). , thus computationally ettra-challenging to Stiefel optimizers. Lin al. (2020) proposed effective method based on alternations between a full Sinkhorn step (Cuturi 2013). given the current projection, and optimization step that improves the projection. In particular, they developed Riemmanian optimizer with projection and retraction (RGAS) and its adaptive learning rate version (RAGAS), We replace RGAS and RAGAS by our optimizer and others, and test on the hardest experiments in Lin et al. 2020). namely MNIST and Shakespeare.
eonneiically-nononvexx
Results are in 1 Our method is observed to find the largest value of PRW and thus the best projection among the tested, which implies the best performance. Details of setup are in Apdx. P.2
3.2 How COULD ORTHOGONALITY IMPROVE VANILLA VISION TRANSFORMER (VIT)?
This section explores the possibility of making self-attention in Transformer models (Vaswani et al.. 2017) orthogonal. The intuition is, if we interpret each attention head characterization of interac- tions between tokens, maybe linearly dependent interactions shouldn't be double-counted so that the characterization can be more efficient and robust. In addition, there is unlikely just one way in which tokens relate to each other, and multiple heads handle different characterizations in parallel. So our second question is, should different heads also be orthogonal to each other? To formulate precisely, we follow the notation of Vaswani et al. (2017) and only review relevant parts: a Scaled Dot-Product Multihead Attention is given by MeltiHead Q, K, V) Concat (head, -.. head ) W' where
8

---

Published as a conference paper at ICLR 2023
VIT trained from scratch (6.3M parameters) test error (mean std tests, percentage)
Method
CIFAR 10
CIFAR 100
Orthogonality across heads
Stiefel SGD (ours, same as Tao and Ohsawa 2020)
8,85 E 0.35
31,51 E 0.39
in this case)
Projected Stiefel SGD Li et al.
2020
9.65 + 0.23
32.14 + 0,48
Regularizer SGD)
9.58 0.21
32.77 0.23
Stiefel Adam (ours)
8.81 + 0.22
32.48 + 0.24
Projected Stiefel Adam Li et al.
2020
11.58 + 0.24
33.98 + 0.12
Regularizer Adam
11.09 0.26
33.90 = 0.24
Orthogonality only within
Stiefel SGD (ours)
8.32 $:2:000 0.39
30.20
E 0.55
Projected Stiefel SGD (Li al. 2020
8.90
31.71 0.26
each head
Regularizer SGD
9.78 + 0.14
33.08 0.42
Momentumless Stiefel SGD (Wen and Yin 2013
12.90
Stiefel Adam (ours)
8.26
$1.04.::0
Projected Stiefel Adam
(Li et al.
2020
10.77
33.73 0.54
Regularizer Adam
11.53 0.28
34.72 0.67
No constraints
SGD
9.75 + 0.37
32.61 + 0.30
(baseline)
Adam
9.52 + 0.32
33.15 + 0.50
Adam W WiLoshchilov and Hutter 2017)
9.05 + 0.43
33.11 + 0.47
Advanced methods
PVT-T (12.8M parameters) Wang et al. (2021)
9.49
30.38
with known error
DeïT-T (5.3M parameters) Touvron et al. (2021)
11.62
32.48
o: only tested for "within each head' hecause it hecomes ortho. SGD Tao and Ohsawa 2020) with O momentum for 'across head".
Table 1: Orthogonal Vision Transformer (Stiefel-Vit ViT) trained from scratch for CIFAR. For each of the 4 classes ((across, only within) x ASGD,Addmm }), blue is best within that class. Underlined is the best over all classes, and bold faced are methods from this paper.
head;- = Attention (QW, KW., VW.), Attention( Q, L.K. V) softmax T.O E RRchnss xds, W.K E Relassa dx, W. E Xd and W R. x desds correspond to train- able parameters. The three input matrices Q, K and V all have dimension requence_lengll x dmodel- dy and d, are usually smaller than dmodel.
OK dd2
matrices
For orthogonality only within head, we require that W.O. W.K are in St(dmodel. dz).. This needs dmudel 2 dg, which holds in most cases. For orthogonality across heads, we need dmantel 2 Thead which is satisfied many popular models, and require Concat (W,2, i = I..., Theud), Concat W K 1..., Thead) to be in St(dmodel, (headda), which means it contains not only 'orthogonality only within head', but also extra cross- orthogonality.
We test these two types of constraints on vanilla ViT Dosovitskiy et al. 2020). trained from scratch. Results summarized in Tab. (and Fig. i2 Apdx. P.3 show: (1) requiring each head to be orthogonal just within itself leads to the most significant improvement of performance; (2) additionally requiring heads to be orthogonal to each other will actually be less effective, although that is still better than requiring no orthogonality; (3) the choice of optimizer matters: our methods give the best results and out-perform non-constrained baselines in all cases, but not all existing Stiefel optimizers can beat non-constrained baselines in all cases; (4) momentum significantly helps train orthogonal VIT; (5) simply imposing orthogonal constraint (with our optimizer) makes ViT outperform some carefully designed models of the same size that are trained from scratch (Tab. A3 in Zhang et al. (2022) ) Note models trained from scratch are very different from pre-trained ViT. To have some realistic expectation of the performance of trained-from-screne PVT-T (Wang et al., 202I), has 9.49% and 30.38% error on CIFAR 10 and CIFAR 100, using 12.8M parameters, and another improvement, DeiT-T (Touvron et al.) 2021), 1, achieves 11.22% and 32.48% respectively with 5.3M parameters. Our model in Tab (and Fig in Apdx. P3 uses 6.3M parameters. More details about this experiment can be found in Apdx. [.3
VIT, recall, for example, an improved ViT,
Remark (No additional hyperparameters needed for the Stiefel part). Our Stiefel-ViT model has both non-Euclinn non-Euclidean and Euclidean parameters, the latter due to the reason that we don't impose orthogonality on feedforward layers and W. Our optimizer can use the same learning rate for both sets of parameters. This is because the optimizer is the time discretization of ODEs, and synchronous discretization is sufficient for convergence. This contrasts with other approaches such as projected Steifel SGD/Adam (Li al. 2020), where learning rates for these two sets of parameters need be adjusted separately and could differ by 40 times in some tasks. See Apdx. P.T for more discussion.
9

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
We thank Tomoki Ohsawa for insightful discussion. We are grateful for partial support by NSF DMS- TALS-1844 1847802 (LK, YW and MT), NSF ECCS-193677) (MT), and Cullen-Peck Scholar Award (LK, YW and MT).
REFERENCES
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition scale. ICLR, 2020.
François-Pierre Paty and Marco Cuturi. Subspace robust wasserstein distances. In International conference on machine learning, pages 5072-5081. PMLR, 2019.
Tianyi Lin, Chenyou Fan, Nhat Ho, Marco Cuturi, and Michael Jordan. Projection robust wasserstein distance and riemannian optimization. Advances in neural information processing systems, 33: 9383-9397, 2020.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning in deep linear neural networks. arXiv preprint arXiv: 13 /2.66120, 2013.
Lechao Xiao, Yasaman Bahri, Jascha Sohl- Dick--Dccst,en, Samuel Schoenholz. and Jeffrey Pennington. Dynamical isometry and mean field theory of cnns: How train 10,000-layer vanilla convo- lutional neural networks. In International Conference on Machine Learning, pages 1399-5402. PMLR, 2018
Xiuwen Liu, Anuj Srivastava, and Kyle Gallivan. Optimal linear representations of images for object recognition. In 2003 IEEE Computer Society Conference on Computer Vision and Pattern Recognition, 2003. Proceedings., volume 1, pages I-1. IEEE, 2003.
Moustapha Cisse, Piotr Bojanowski, Edouard Grave, Yann Dauphin, and Nicolas Usunier. Parseval networks: Improving robustness adversarial examples. In International Conference on Machine Learning, pages 854-863. PMLR, 2017.
Jun Li, Fuxin Li, and Sinisa Todorovic. Efficient Riemannian optimization on the Stiefel manifold via the Cayley transform. ICLR, 2020.
Jiayun Wang, Yubei Chen, Rudrasis Chakraborty, and Stella X Yu. Orthogonal convolutional neural networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 1 1150-11515. 2020.
Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International Conference Machine Learning, pages 1120-118, PMLR, 2016.
Scott Wisdom, Thomas Powers, John Hershey, Jonathan Le Roux, and Les Atlas. Full-capacity unitary recurrent neural networks. Advances in neural information processing systems, 29, 2016.
Kyle Helfrich, Devin Willmott, and Qiang Ye. Orthogonal recurrent neural networks with scaled cayley transform. In International Conference Machine Learning, pages 1969-1978. PMLR, 2018.
Aston Zhang, Alvin Chan, Yi Tay, Jie Fu, Shuohang Wang, Shuai Zhang, Huajie Shao, Shuochao Yao, and Roy Ka-Wei Lee. On orthogonality constraints for transformers. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing ( (Volume 2: Short Papers), pages 375-382, 2021.
Nitin Bansal, Xiaohan Chen, and Zhangyang Wang. Can we gain more from orthogonality regu- larizations in training deep networks? Advances in Neural Information Processing Systems, 31, 2018.
10

---

Published as a conference paper at ICLR 2023
Pau Rodríguez, Jordi Gonzalez, Guillem Cucurull, Josep M Gonfaus, and Xavier Roca. Regularizing cnns with locally constrained decorrelations. ICLR, 2017.
Eugene Vorontsov, Chiheb Trabelsi, Samuel Kadoury, and Chris Pal. On orthogonality and learning recurrent networks with long term dependencies. In International Conference on Machine Learning, pages 3570-3578. PMLR, 2017.
Mario Lezcano-Casado and David Martinez-Rubio. Cheap orthogonal constraints in neural networks: A simple parametrization of the orthogonal and unitary group. In International Conference on Machine Learning, pages 3794-3803. PMLR, 2019.
Alan Edelman, Tomás A Arias, and Steven T Smith. The geometry of algorithms with orthogonality constraints. SIAM journal on Matrix Analysis and Applications, 20(2)3003-333, 1998.
Zaiwen Wen and Wotao Yin. A feasible method for optimization with orthogonality constraints. Mathematical Programming. 142(1):397-433, 2013.
Shuai Li, Kui Jia, Yuxin Wen, Tongliang Liu, and Dacheng Tao. Orthogonal deep neural networks.
IEEE transactions pattern analysis and machine intelligence, 43(4):1352-1136, 2019.
Hongyi Zhang and Suvrit Sra. An estimate sequence for geodesically convex optimization. In Conference On Learning Theory, pages 1703-1723. PMLR, 2018.
Kwangjun Ahn and Suvrit Sra. From nesterov's estimate sequence to riemannian acceleration. In Conference Learning Theory, pages 84-118. PMLR, 2020.
Foivos Alimisis, Antonio Orvieto, Gary Becigneul, and Aurelien Lucchi. Momentum improves optimization on riemannian manifolds. In International Conference on Artificial Intelligence and Statistics, pages 1351-1359, PMLR, 2021.
P-A Absil and Jérôme Malick. Projection-like retractions on matrix manifolds. SIAM Journal on Optimization, 22(1): 2012.
Gary Bécigneul and Octavian- Fugen Ganea. Riemannian adaptive optimization methods. ICLR, 2019.
Rongjie Lai and Stanley Osher. A splitting method for orthogonality constrained problems. Journal of Scientific Computing, 58:431-449 2014.
Pierre Ablin and Gabriel Peyré. Fast and accurate optimization on the orthogonal manifold without retraction. In International Conference Artificial Intelligecce and Statistics, pages 5636-5657. PMLR, 2022.
Xin Zhang, Jinwei Zhu, Zaiwen Wen, and Aihui Zhou. Gradient type optimization methods for electronic structure calculations. SIAM Journal on Scientific Computing, 36(3):C265-C289, 2014.
Xiaoying Dai, Zhuang Liu, Liwei Zhang, and Aihui Zhou conjugate gradient method for electronic structure calculations. SIAM Journal on Scientific Computing, 3966) 2017.
Bin Gao, Xin Liu, Xiaojun Chen, and Ya-xiang Yuan. new first-order algorithmic framework for optimization problems with orthogonality constraints. SIAM Journal on Optimization, 28(1): 302-332, 2018.
Bin Gao, Xin Liu, and Ya-xiang Yuan. Parallelizable algorithms for optimization problems with orthogonality constraints. SIAM Journal on Scientific Computing,
2019.
Jiang Hu, Bo Jiang, Lin Lin, Zaiwen Wen, and Ya-xiang Yuan. Structured quasi-newton methods for optimization with orthogonality constraints. SIAM Journal on Scientific Computing, 41(4): A2239-A2288 2019.
Xiaoying Dai, Qiao Wang, and Aihui Zhou. Gradient flow based -hanm-sham density functional theory model. Multiscale Modeling Simulation, 18(4):1621-163,, 2020.

---

Published as a conference paper at ICLR 2023
Fanchen Bu and Dong Eui Chang. Feedback gradient descent: Efficient and stable optimization with orthogonality for dnns. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 6106-6114, 2022.
Benedict Leimkuhler, Tiffany Vlaar, Timothée Pouchon, and Amos Storkey. Better training using weight-consstrineed stochastic dynamics. arXiv preprint arXiv:21006 (10666077) 10704, 2021.
Taeyoung Lee, Molei Tao, and Melvin Leok. Variational symplectic accelerated optimization on lie groups. Conference on Decision and Control (CDC), 2021.
Valentin Duruisseaux and Melvin Leok. A variational formulation of accelerated optimization on riemannian manifolds. arXiv preprint arXiv:2101.065522 2021.
Hemant D Tagare. Notes on optimization on stiefel manifolds. Yale University, New Haven, 2011.
Yurii E Nesterov. A method for solving the convex programming problem with convergence rate (11/k². In Dokl, akad. nauk Sasr, volume 269, pages 543-547, 1983.
Weijie Su, Stephen Boyd, and Emmanuel Candes. A differential equation for modeling nesterov's accelerated gradient method: theory and insights. Advances in neural information processing systems, 27, 2014.
Andre Wibisono, Ashia C Wilson, and Michael I Jordan. A variational perspective accelerated methods in optimization. Proceedings of the National Academy of Sciences, 11(44)):7331-E77388 2016.
Molei Tao and Tomoki Ohsawa. Variational optimization on lie groups, with examples of leading (generalized) eigenvalue problems. International Conference on Artificial Intelligence and Statistics (AISTATS), pages 4269-4280, 2020.
Renyi Chen, Gongjie Li, and Molei Tao. Grit: A package for pructure-presenvve simulations of gravitationally interacting rigid bodies. The Astrophysical Journal, 919(1):50, 2021.
Robert McLachlan and G Reinout W Quispel. Splitting methods. Acta Numerica, 11:341-434 2002.
P-A Absil, Robert Mahony, and Rodolphe Sepulchre. Optimization algorithms on matrix manifolds. Princeton University Press, 2009.
Nicholas Higham. Stable iterations for the matrix square root. Numerical Algorithms, 15(2): 227-242. 1997.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiy Kumar, and Suvrit Sra. Why adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383-15393, 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR, 2015.
Michael Jordan, Tianyi Lin, and Hanmanouil-Vasileia Vlatakis- Sanakis-Gkrrggkounis. First-order algo- rithms for min-max optimization in geodesic metric spaces. arXiv preprint (xXiv:2206000411 2022.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Zizhao Zhang, Han Zhang, Long Zhao, Ting Chen, Sercan Arik, and Tomas Pfister. Nested hierarchi- cal transformer: Towards accurate, data-efficient and interpretable visual understanding. AAAI, 2022.
12

---

Published as a conference paper at ICLR 2023
Wenhai Wang. Enze Xie, Xiang Li, Deng- Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 568-578, 2021.
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers distillation through attention. International Conference on Machine Learning, pages 10347-10357. PMLR, 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint
10X1v:1711.0500,, 2017.
Ashia C Wilson, Ben Recht, and Michael Jordan. A Iyapunov analysis of accelerated methods in optimization. Journal of Machine Learning Research, 22(113):1 2021.
Molei Tao. Explicit high-order symplectic integrators for charged particles in general electromagnetic fields. Journal of Computational Physics, 327:245-251, 2016.
Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. Autoaugment: Learning augmentation policies from data. arXiv preprint arXiv: 1805.09501, 2018.
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.039%: 2016.
13

---

Published as a conference paper at ICLR 2023
Notation
For simplicity, we denote G 0f
when there is confusion.
0x
A A QUICK SUMMARY OF OPTIMIZERS WE EXPERIMENTALLY COMPARED TO IN THIS PAPER
manifold preserving
Optimizer
Pros and Cons
Stiefel SGD/Adam (Ours)
allows a family of metrics
/ fast iterative method for position retraction
momentum needs no retraction
/ learning rate be made adaptive
same learning rate for Stiefel & Euclidean parameters
Projected XGD/Adam et al. 2020
Stiefel
/ learning rate be made adaptive
X slower iterative method for position retraction; fast if early
stopped but then manifold preservation is lost
X momentum needs projection every iteration
x different learning rates for Stiefel & Euclidean parameters
Momentumless
X no momentum
Stiefel (S)GI) (Wen
and Yin 2013)
Regulare anizrr zer
Cisse
fast computation
et al. 2017)
X sensitive to regularization strength
Table 2: summary of pros and cons of existing optimizers.
B TECHNICAL REMARKS ON GEOMETRIC MECHANICS
B. INTUITION OF COTANGENT BUNDLE
If we require exact Stiefel manifold preservation for all time, i.e., X(t) (t) = T, simple time differentiation gives 2nd-order constraint X(t) Q(t) + X(t) 0 for Q X. This is why momentum Q has the additional structure (+t) E (tt St.
B.2 WHAT EXACTLY MOMENTUM?
Through this paper, we abused notation and called Q : X momentum, following the common practice of the community. really should be called velocity instead, because although in canonical Euclidean spaces this doesn't 't make much difference, for mechanical systems manifolds they are the same thing. For example, for our case, the geometrically intrinsic momentum variable should be given by Legendre transform
ØL P == 0x - - axx )
(15)
instead, and velocity Q is in the tangent space while momentum P is in the cotangent space. In the paper we used the identity map for an isomorphism between the tangent and cotangent spaces, but we't d like to make our variational formulation more elegant by viewing the kinetic energy as a natural pairing between the velocity and momentum variables, then the isomorphism should be given by the metric as in [5).
None of these affect the correctness or the efficacy results in this paper. This discussion just about terminology.
B.3 THE SPECIAL CASE OF St(n, n) AND ITS RELATION WITH ((n) AND SO(n)
When n 2 m, St(n, m) z (m))/O(n - m); when n > m strictly, we also have St(n, m) s 50(m)/SO() m). However, in the special case of m, St(ni,) a O(n) is not connected
14

---

Published as a conference paper at ICLR 2023
unlike the >> cases. Since our optimizer is based on the discretization of continuous dynamics, it cannot make jumps and thus just optimizes on a connected component of St(n,, m). which means for the special case of m, it is, to be precise, optimizing SO(n) but not (n)), although a similar complication is nonexistent when n > m.
C DETAILS ABOUT THE PER-ITERATION COMPLEXITY AND COMPUTATIONAL COST FOR OUR ALGORITHMS
The most costly operation in our algorithms is the n X m matrix multiplication (note n m). The computation for the matrix exponential and square root of matrix inversion, is cautiously designed to only deal with matrices of dimension m x m (see Apdx. M. M.2). and thus admits at most ((²)) at each step (particularly, forward Fuler only has the complexity of ((m²) while Cayley map is ((m²)
For the inner loop of the square root of matrix inversion (Algo. 3. due to its quadratic convergence (Higham 1997), it takes Olog log(1 (1/1()))) /u)) number of steps to achieve the machine precision u. Note early stopping can be applied to this loop so that the complexity can be further reduced while the
order of the overall method (Ist-order) is still maintained. The combination of all the above gives complexity (nm²² + Om log log(1/u)
Commonly used optimizers on Stiefel manifold (e.g. Tab. also have per-iteration complexity of Omm²²). Although it is difficult to compare them with our method in terms of the constant prefactors in respective complexity bounds, can do some heuristics to estimate these prefactors, which will demonstrate the low-cost advantage of our method. For our St Stieel SGD (Algo. []. we count the numbers of matrix multiplications needed per iteration, each of which costs nm 2 Namely, step 2: 2 multiplications; step 3: I; step 4: I; step 5: 6. In step 6 we need to call Algo. 3 for computing matrix root inversion, and a closer look gives 3 m=-cost matrix multiplicattons in each iteration of Algo. 3] Since Fig. (5c) shows that 8 inner iteration is enough for this inner loop to converge under double precision, total, 10 nm2-cost matrix multiplications and 24 m²-cost matrix multiplications are needed in each (outer) iteration of our optimizer (Algo. D. We also count the number of matrix multiplications needed for other aforementioned algorithms. For Moment ess Stiefel (S) GD (Wen and Yin, 2013). the smartly designed Cayley map retraction takes about 10 times of ncos matrix multiplications and one inversion for (2m) x (2m)-sized matrix each iteration. In comparison, Pro jected Stiefel SGD Li et al. (2020) needs 9 nmm2-cost matrix multiplications for projecting momentum and 6 nm22cost matrix multiplications each iteration of Cayley loop. Since 8 iterations are needed to compute the Cayley transform to single precision (see fig. 5 note it is just single precision instead of double precision by our inner loop), a total of 57 nmm -cst matrix multiplications are needed per outer iteration.
The above estimation is just for the part on maintaining the structure of the position variable. Let's now also discuss the momentum variable. Our carefully designed way of introducing momentum bypasses the costly moving momentum while keeping it in tangent spaces. Note SGD first moves momentum in the Euclidean space and then uses a cleverly designed projection, which markedly improves computational efficiency, but it still devotes 9 matrix multiplications to just the momentum projections, let alone the fact that needing projection means less accuracy. To add to this discussion, parallel transport is another existing way of moving momentum, but it cannot even be solved with Onnm²) complexity.
Stiefel
Altogether, the above estimations heuristically show that our prefactor is close to that in Moment Comentumless amless Stiefel (8) GD but much smaller than those in existing methods that have momentum. This matches the experimental evidence in Fig.
D TERMINOLOGIES OF RIEMANNIAN OPTIMIZATION
In each iteration of optimization, gradient needs to be vector in tangent space, and the operation that maps nonintrinsic gradient to the tangent space is called projection. Then, the variable goes one step on manifold given the tangent vector, ideally via the exponential map, which is, however, expensive or even not admitting closed form computation, and approximation known as retraction is used.
stands for machine precision, which is very small number
15

---

Published as a conference paper at ICLR 2023
When momentum, tangent vector, is involved, at the same time update the position, tangent space is also changed. An intrinsic way to move tangent vectors is by parallel transport, which an isomorphism of the tangent spaces of two points connected by geodesic.
E ABOUT THE Y, V DECOMPOSITION
E.1 THE GEOMETRIC VERSION AND THE DYNAMICAL VERSION OF THEIR STRUCTURAL CONSTRAINTS
tangent vector E St can be decomposed, studied in Edelman al. (1998) Absil al.
(2099, into
Y 7
Q Q= (X )
= XY + X 2
(16)
where X E St(n, m), Y E R 7 E (m- (n- mm X- E (ax(n-m) is a matrix in the orthogonal
complement of i.e., X X mx (n-mn Let X- We have X V and
Y =
(17)
V XY X1-
The above representation immediately implies Y is skew-symmetric matrix since
x"
(18)
In our derivation, however, X and Q variables (they really should be X(t St, Q(t) E St) that change with time, and this poses a new challenge. More precisely, given the Q dynamics, i.c. Q there could infinitely many ways of assigning Y and V so that they together produce the given Q; however, in general they do not maintain the tangent decomposition structures. More precisely, starting with
Y(O) + (0)) = and X(O)TV(O)) =
(and hence ((0)TQ(0) ((000((0)(())))) 0), may not have
Y(t) Y = 0 and (() (tt) = 0 for t > 0
T()))((()))))))))))()(((())))))) (() 0 will still be guaranteed. However, we have found nontrivial choice of y and V (Eq 8), so that the (static) geometric constraint becomes dynamically true, i.e. (tt) T + (t) = and (t) = for all t. Therefore we can simply maintain the decomposition
despite that X(()TQ(() + Q(()
Y(1) -x(()
- x ((XX(()x()
E.2 THE ADVANTAGE OF THE Y, V,V DECOMPOSITION
The Y. decomposition is near 'orthogonal' in the sense that XY 1 V and it makes the metric separable. More precisely, consider our Y, V representation. VΛ, E St, 1, if we denote Ya,: = ATA, and VA (T- X-xXXT)A which means A; = XY + Va,i, then the metric 2 can be rewritten as
(2),)))) - VV.)) ( axx (XY XXXX) + VA,2)
=(1- (1- (110)TMYY,ya.. - Tr(V,,
This means the metric of the tangent space is linear combination of the metrics of and -ddircctions. Therefore, this gives an intuition why our continuous dynamics and numerical discretization in the following can handle the family of canonical-tyre metries.
16

---

Published as a conference paper at ICLR 2023
F CHOICE OF Y
Although monotonicity of r suffices for the convergence, choices of y f/r affect the convergence speed just like Euclidean cases (e.g. Wibisono et al. (2016); Wilson et al. (202I)). Popular choices include constant y and ytt) = 3/1t, but methods proposed here work for all y (e.g., = 3/t + ct
for variance reduction Tao and Ohsawa (222)). However, for simplicity we will use constant ? from now on.
G PROOF OF THEOREM
Proof of Theorem
4 Part 1. Given the Lagrangian
L(X, L(X, r(t)
-(x)) X -1-xxxxx)x - 2 TA
(A ( xxxx) 1)) (19)
Legendre transform gives momentum Pas
aT.
P ax r([ aX )
One can equivalently switch the Hamiltonian picture, and the corresponding Hamiltonian S* + R with H(X, P) == (r((PTX)
1 -1 bxx +r])))+(((()--)))))))) [X (Imxm
H(X, P) = 2r
(20)
where b - solves 1 -I- equations
6XxT) I. Hence get the following Hamilton's 's
0H = (I 1XXT)P
x
0P
r
(21)
0H b
0f
(PP x - -r -XA 0x
0x
2r
Since X 1, which gives X X + x = 0. So we have X T + P O. Take derivative
to get X P + P PT
+ PT X
0, and use the condition that symmetric, we can solve
A is
b + 2 PT p
b X - G
A
(22)
2r
2r
And (X, P) system becomes
x 0H 1
- -bXXT)P
OP
p
0H
PP X
- p+
1
_ xx xxx PP x rG Fxx
b
ax
2r
7
2r
(23)
By a coordinate change Q = [1 1xx P E St and define friction parameter == r/r,
Eq. 23
becomes
② - (1 -1- p [ bxx -
b - _ -
3a [I (I-
1+b
1-b xx
=-
G
XGTx x
2
x =Q
3a (1 - - x G b+b xx
1-b
(24)
i = 7Q -Q -
G
XG x
2
2
In order to prove the second part of Thm. need the following lemma.
17

---

Published as a conference paper at ICLR 2023
Lemma 1 (First-order stationary point). If s.t G 2 1 X 0, then AA E St, we have Tr(G" A) = 0, which means X is first-order stationary point of f.
Proof. Left multiply both side of G X
Ob XX", we have XXTG- =
XG X, further we also have G = XX
So we have
Tr(G" A) =T((TGXXT)
-T(X -XTTATXXTG)
=- =-TrAAG)
Thus have Tr(A G) = 0, AΔ St, which means is 1-order stationary point of f.
Proof of Theorem 7 Part 2. Let - (X(t),Q(t) be solution of Eq. Define the 'energy'
function TSt
as
B(X (rQ
(Q) f(X)
(25)
This gives Lyapunov function. So we have neighbourhood U of (X.., 0) such that E(X, Q)
2
f(X) 2 f(X.) for any (X, x) U. More over, since X X + Q X O, we have
dE dt [
0f T
+ Tr
X
(26)
0X
Using the fact that X + Q = 0, we have TT(QTxQTQ) and Tr(Q XXTG)
Tr(Q XG 0, which gives
dE dl (x(1, - ax QQ) <
(27)
Since we have monotonely increasing, which means y 0, Vt. Then we have the energy is decreasing monotonically, implying that Q = 0 when converged. By lemma I that the limiting point for (t) is a first order stationary point, which is X, since it is an isolated local minimum.
Remark 2. For the sake of length, of convergence in specific situations (e.g., under geodesic convexity) will not be quantified, but it should be obtainable via tools in Wilson et al. (2021); Duruisseaux and Leok (22I).
H
PROOF OF THEOREM 2
Proof. We can derive system of ODEs
xxxx)=xxQ++*>
d
(x=))=->x"
3a x xx=====x0Q0Q
XI-XX x
xx X
dt
-2"
-T 1+ x x
G 1-b x X
with viewing X G as matrix function of t following Eq. (7). By the uniqueness and existence of
ODE, we have that this system of ODE with variable (x X. TTQ) Q) and initial condition (I,0) has the unique solution X X
0.
18

---

Published as a conference paper at ICLR 2023
I XYV SYSTEM IS STRUCTURE PRESERVING
Theorem 8 (Constrained optimization with unconstrained dynamics), As long as the initial condition of 8 satisfies
(() (() Im xm:
(()) (()) =0mxm: X(O) V(0) mmxm
then the dynamics automatically satisfies the same constraint, i.e.,
(()) (t) Imxm
Y(t) + +(()
Omxx:
X(t) V(t)
0mxm:
(28)
Proof. By the uniqueness of solution of ODE, we know that the following ODE for (XT ",," Y Y, X VV,, derived from Eq. 8 has unique solution if view V V as an independent variable.
d
di X
dt (Y + + = -- = +Y)
(29)
V = 7x V + 3a dt
VY + (1 -1 X1xX
2
We can see that given initial condition (Xg Xo, Yu + Yo, X, Vo) (1,0, the unique solution
is (X ",, +Y, VV) 0). So we know that the constraint in Eq. (88 are preserved by continuous dynamics Eq. (8)
PROOF OF THEOREM 3
Proof. We assume the initial condition (Xo, Yo, Vo) satisfies constraint (28
For Eq. 9 we check X (t)X(t+) = 0, (() + YT (t) = 0 and X (t)V = 0.
d (Y +Y) = +Y
dt
+Y
with initial condition Yo + Yo = 0. we have Y TT =0
d =
di
=VTVV XV
with initial condition X Vn 0, we have X XV=0
di xxx=xTx+*77
=XTXY x
XX
Using the conclusion Y + T 0 that we have just proved and initial condition X Xo = 1, we
have X X
For Eq. IOP. we have the exact solution of d2 is given by X(t) = (000, = Y()), and with 2a-21 3a-2y(6 (t),
M 7 Imxm
V(t) (0) expm(- -Mt) - (1-))(1-X(000 (() of
ax (X(0))- -T-exp((-Mt) (30)
19

---

Published as a conference paper at ICLR 2023
For Eq. (11. need to check di (t)x((()))) = and d (X (t)V(tt) 0.
By the uniqueness of solution of ODE, know that the following ODE for (X V.VV),. derived from Eq. [1 has a unique solution
d
(x x) =X-
dt
ti (X V) = - =11-x
(31)
d (V = = - di
We can see that given initial condition (Xg Xx, Xx Vo. 10 Vo) = (I,0, V/ Vo). the unique solution
Vo.
(XTx, V,.'VVVTTV 01, Vo). So we know that the constraint x I, X 0 preserved by continuous dynamics Eq. 11
K PROOF OF THEOREM 4
Proof. Eq. is structure preserving. We use the idea from Tao and Ohsawa (2020). Assume the initial value Xo. Yo and Vi satisfies the constraint, i.e., X Xu 7, Yo 0 and X Vo 0. For the first step updating Y, due the special form of the derivative of that is always skew- symmetric, we can tell that Y, is also skew-symmetric and Yn + +Y, O. And the skew-symmetricitt of Yn gives us that
Xx, _ XA expm )x,, Xoexpm(hYx))
x V = expm(hY), )x Vo
so all 3 conditions in Eq. 28 are satisfied, meaning Eq.
12 is structure preserving.
Eq. (14) structure preserving.
Given initial condition X Xo = and X Vo = O, we check constraint X Vn, = 0.
3a xX, 1x, AxT (T XxXX) Of 0x
X (1) Vo
(Xo)
Eq. 11 is structure preserving, First, we show that the discretization is one order approximation of the exact solution. We get Vn, and X₁ by forward Euler. And since Xo and Vo satisfies the constraint 28 we have
x X, =(x, + AV)) (Xo + 1V))
=(X, Xo + hXj Vo + hV. X. + hVVV
=1 + 12v, Vo
which means X, = +X++000 0 = Xo + hVo + O(h2), indeed one order approximation of exact solution.
Next we show the numerical discretization structure preserving. For constraint T = 0, we can check that
XTVV V/ = XT Vo hX. VoV, Vo - hxj XOVo Vo - 12V. XoV, Vo =
which gives us X Vn = 0. For restriction X X = I, we have
X, = (x, x,) t x x xx)==
20

---

Published as a conference paper at ICLR 2023
L PROOF OF THEOREM 5
Proof. The composition of structure-preservinn maps is structure preserving. Φ1 o022od is a Ist-order integrator due operator splitting theory McLachlan and Quispel (2002), and its convergence order is kept after any Φ gets replaced by Φy as long as the difference between Φ and dj is higher-order Tao (016).
M
DISCUSSIONS ON OUR NUMERICAL INTEGRATOR
M.1 BENEFITS OF THE STEP (x)+x))*
To remain the position X the manifold, certain techniques that pull point back manifold used. For example, in Li et al.) (2020); Lin et al. (2020) for an n-by-m matrix X that is not on the manifold, they perform QR decomposition X QR and update the new X to be the first columns of Q. However, the QR decomposition is not unique and such step does not have a closed form expression. More importantly, it cannot help design structure- prreservnp scheme.
Instead, our algorithm has similarly functioned step via the square root of the inverse matrix X₁ = (X, XG) x 2, In this case, it follows that X, Xr. = I. This step is carefully designed from discretization ODE ensuring the structure of X, Y. V is preserved at the same time. Through this update, we able to obtain pireeture-preserving method with closed form expression (see the proof of Theorem 4 in Apdx. K). Note that the square root of the inverse matrix can be iteratively solved with low computational cost (see Algo. Higham (197)). was proved to be quadratically convergent, which means only couple of iterations are needed to reach machine precision, and only matrix multiplications are needed. This lead to our inner loop is more efficient in both cost per iteration and convergence speed comparing to Li ct al. (2020). See Sec. O for more details.
Algorithm 3: Algorithm for matrix root and matrix root inversion (Eq. (2.6) in Higham (1997))
Input: Symmetric m-by-m matrix A, tol
Initialization :Yo = A, Zo = T 7mxm; h = 0
while y
All 2 tol do
Yo+1 (3/ Yx)
Zx+1 7 ]3I Zz
A + k+1
end
return
Y& z A and 75 ~ A
Another benefit of this step is that it is more stable to the truncation error produced by finite machine precision (which can become significant issue if the model is trained in single accuracy on consumer graphic cards or even quantized). Although X, X2 follows the expression
X, X1 =(xo + (V)) (Xo + +Vv)
=(X, Xo + hx Vo hV Xo + 13v, Vo)
=1 + 11vv Vo.
the stability of X+(X. ((XXXX) Xx(XXX,)-1/2 11/2 is much better than X+(1 + 12V, Vo) -1/2, Particularly, when n = m, it is an identical map if there is no machine error.
M.2 SIMPLIFICATION OF MATRIX EXPONENTIAL
There steps in the discretizations 30 and 122 that the matrix exponential of m m matrices. Similar matrix exponential is also shown in Li et al. (2020); Wen and Yin (2013) but with matrices of much larger size n X n (note n m)) We introduce two ways of simplifying the computation of the matrix exponential. First, note that Cayley transform is a 2nd-order structure- preserving approximation of the matrix exponential, defined as follows
Cayley(hY) (T AY/2) hY/2) exp(hY) + 0(13).
(32)
21

---

Published as a conference paper at ICLR 2023
By applying the Cayley transform, the computation is reduced to inversion of matrix and matrix multiplication while it still keeps the variable on the manifold.
Additionally,
the first-order forward Euler to discretize the ODEs for the two steps, which also the first-order truncation of the matrix exponential. In this case, the X update in scheme (12) just X Xo hY, and is no longer structure-pred prreeeri-p but structure- preserving property of the overall algorithm not affected (see Thm. 6
For scheme
30 its Cayley map approximation is structure preserving and defined as follows
VA = Vo Cayley(-?hl) + 3a (V,%)
1 exp(->h)) (I Xxx 0s
(33)
2
Note the variable V will still satisfy the constraint Euler
28
even we change the update of V to forward
3a 2 2 hVoY
1- exp(-xxhh
0f
Vn Vo V0- V/Vo +
Xoxx OXo
(34)
7
In fact, the above two ways and the original matrix exponential share the same complexity per iteration. Also, we do not find any significant difference in the convergence numerical experiments (Apdx. Q). Hence we will always use forward Euler as default, which has the lowest computational cost.
M.3 PROOF OF THEOREM 6
Proof. The composition of Φ1, 62. 63 in any order will give a structure preserving scheme, because in Sec. K we proved that each of them is structure preserving.
When apply in specific order ø3 01 62, we prove when X x / is no longer preserved by 62, the assembled scheme is still structure preserving. We only need to prove that [x,, Y,, V/] (a Yo. Vu) satisfies the following: when initial condition satisfies Y + Y, and X Vo
we have X2, Xx I. Y + Yn 0 and X, VA.
O.
Since we have the step X₁ X+ x+)-1/2, 1122, we have X X = J. Y is not updating so Yn + Y
X. Vn X.. Vo hX, XoVo Vo hX, voVo Vo 12x, XoVo XoVo Vo =0
As result, V, Vn (x) V/2 =0
M.4
THE INTEGRATOR IN RESCALED COORDINATES USED IN ALGORITHM
X, =X0
X,, -X, + TTp
Zn =Zo
21
Z,g -170
1-b
X
05
0s
No
Un =utp +
3a 2
2 Zo
2
ax,
Ox,
of
Un -U0
- (1 - -xxX) 0Xo
X, =X, + NUxX Xn
X, =x((X X+)
Φz.
bn =2o
Un =U0- - Nouuo Uo
22

---

Published as a conference paper at ICLR 2023
M.5
THE INTEGRATOR IN RESCALED COORDINATES USED IN ALGORITHM 2
f 1- b
0f
0f
X;
(X,)
(Xt)
OX
0x
9i -(1 - x.X,) 0f (X.)
'gradient' and 2-order momentum
OX
P++1 =B2ps + (1 B2)
96+1 =B29: + (1 B2)gg² .
X X; n/ 1 B221 Zit1 O (pit1 +
Xi+1 X;
21
Z,+1
8,Z,- (1 -1
Zi41 Zi
U1+1=
Ul+1 =
BiU; 3a- 3n-2,115 11U - (1 (1- -B))
2
0 = 1- (1- -1- -=XXXxxx) 0 (0 +())
X1 =x; + NU,,
:3:
X++1 = (X X1) -1/2
Z++1 =2i
U111 =U:- nXXUUT U.:
N PROOF OF THEOREM 7
Proof. We will use the same notation as in algorithm Assume X, = I, Y, +Y,
x, V, = 0. p; = p..". Our initialization satisfies these conditions. We prove that these conditions are satisfied after mapped by de o Φi 02 the following. Step Due to the skew-symmetricity of fi, we have P++1 is symmetric.
Step 2: X,, Z, U,4 1 == Xx( Z, U2), where Uc+3 -.U,
2 "nUZZ (1 - -1)9.. Since
x, U, = 0 and x, gi = O, we can check
3a 2
x, - =
nxx U.Zi -1-11 94 =0
2
and X Step 3:X, 1. Zit1, U == Z, U,+z), where Z1+1 = BiZ, - -(1- -1- BB)ff
X; + + 1 3 Xi (z+ 2 (pit + +e)
Since we have X, UTU = 0 We can have directly XT U/+1 0. And due to Sa is symmetric, 1+ X = I no longer stands now but it satisfied until
Z1+1 is still skew-symmetric. Note that X i41 the algorithm finishes.
Step 4 X Z411, U4+1 = (X Z U443). )) First we can see from X,+1 = XxXXX,,,+
that X = I as long as X = X + nū X
i4g X is full rank, which is always true since
X, is full rank and 1 is small.
Since we have XT ú = 0 by the special form of matrix I - X (XT Xx+3)_1x71 -X 3 14
= X 0 UTU++
Since X,Tx+=
X+ = 1,
=
(U+4-nXx++) -ai--nxe nx NON++gUTU,00 (T+++) =0
So we proved all the constrain are satisfied again, which means XT X 1, Y,T, + Y:4 0,
XTT V4+1 0, Pi+1 P
23

---

Published as a conference paper at ICLR 2023
O OPTIMIZATION ON (O(n) AS A SPECIAL CASE FOR = ON St(n, m)
Our method for Stiefel optimization can naturally be applied on the special orthogonal group (0(m) (Apdx. B.3 and is defined as follows
SO(n) (X E ER*" x = I, det 1}.
(35)
Tao and Ohsawa 2020) proposed an efficient algorithm based on Lie group structure for the opti- mization on SO n although it cannot be generalized to Stiefel case. Our method restores the same integrator in Tao and Ohsawa 2020) while using different approach applied to a family of metrics.
In greater detail, for (O(n), the canonical-type metric (2)
degenerates to Euclidean metric up to constant scaling, i.e., g(A A²) = (1 - a)Tr(AA A2). Then the corresponding Lagrangian is
L(X, r(t)
|
a Tr x (x*)
f(x)
2
(36)
which results in the following position position-momentuuee momentum (X, Q) dynamics
(t) =Q(t)
Q(I = ( - ((000)000000000 Q(1) 0f
(37)
σx (x(()= 1+6 x()x (1)
Os
1-b x(1) of T()x(
Ox (1) (0x()
0x
with the same initialization as 77 Next, apply the same Y, decomposition of Q. Since - XXT O, we have V = 0, i.e., the momentum Q is purely in -drrecction. Hence the equivalent dynamics to (37) is the following
(1) -(())(()
1-b b Y(t) = -
Of (t) Ox
of T ax (xx())
(38)
X
(t)
2
Note when take b = (namely, 1/2). 38) and Ohsawa (020)) although Tao and Ohsawa (002 uses an intrinsic form which is coordinate-Free but not explicit enough.
is identical to the continuous dynamics in Tao
The numerical integrator corresponding to (S)GD is defined the same as Section M.4 and is summa- rized in Alg 4 Note Step[5]is optional in the case. It gives the identity map no arithmetic error due to machine precision is incurred. However, this step is beneficial under low precision (see Apdx. M.
Algorithm 4: Momentum SGD on (O(n) (generalization to Tao and Ohsawa 2020
Hyperparameter :7) E (0, +-0). 7 E [0, +00), a maximum number iteration
Initialization :oo, Vo, Yo s.t. X, Xo = 1, X, Vo 0, Yo + Yo = 0, b o-1
Output: Local minimum X of for i ..., N- do
ft 1با X. of (x.)
at (X:) ax
Y+1 = =YY;-
X₁ = X,expm(nXXx4))
X x, (x, (xxxx)=== X+)
end
return X
In addition, we also have an adaptive version which was absent in Tao and Ohsawa 2020
0.1 ADAM ON SO(n)
In this section, we extend the above optimizer to an Adam version. Note can be derived from either the special case of the Stiefel optimizer, or the structure of Lie algebra so(n), i.e., left-trivialization
24

---

Published as a conference paper at ICLR 2023
(see Tao and Ohsawa (2020) such that the momentum Y can be recognized element in (n))( the tangent space at the identity. In the latter interpretation, the momentum will always stay in the same tangent space (n) which passes
message that trivialization almost reproduce the convenience in Fuclidean space for 5O(n). Hence unlike the general Adam-Stiefel optimizer (Algo. 2. no extra technique, for example projection, is needed for Adam-SO(n). Detailed iterations are shown in Algo. S
Note this is also a structure preserving scheme (see Thm.
Algorithm 5: Adam on SO(n)
Hyperparameter : E (0, +oo), B1 E [0, 82 e [0,1), < 1, maximum number of iteration N Initialization :oo, Vo, Yo s.t. Xo Xo = I, X for i O.... N
Vo X,
0, Yo + Yo = 0, Pu
0, go = 0
do
ft 1-b [x] 8f (X,) ax
85 (X,) Xx)
2
ax
p++1 B2p; + +11 (1-32)
f22
Y+ = BiV; - 101--1-11 Bi)
X X,expm 1- 11---tt11-0 Y1+1 2 (p:p) +()
X Xx(XXXx))**
end return XN
P EXPERIMENTAL DETAILS
Note: codes are provided in supplementary materials.
Experiments are conducted on a high-performance computing cluster whose name shall be revealed post anonymous period. Single v100 GPU was used.
P.1 GENERAL DISCUSSION
Initialization For variable X with the initial Xo. in fact, any initialization with full rank is valid, including the non-orthogonal X, such that X, does not start on the manifold. This is due to the pressrvvn property of our algorithm (both SGD and Adam versions, see Apdx. M After just one iteration, X will stay on the manifold. Here we suggest one way of obtaining orthogonal Xu which is performed in Saxe et al. (201)). After randomly generating entry-wiss i.i.d. normal distributed n-by-m matrix, we can perform the QR decomposition. Denote the orthogonal matrix from QR decomposition as Q and denote the diagonal matrix which is the sign of the diagonal elements of R as R. Then we can take X, = QR and thus have X, Xo = I.
For the variables U, Z in Stiefel SGD and U. Z, p, in Stiefel Adam, we would use zero matrices as their initialization since zero matrices always satisfy these constraints (28).
No need to tune learning rate separately for constrained and unconstrained parameters For problems with both constrained and unconstrained parameters (Stiefel and Euclidean spaces), for example, the ViT test in Sec. 3 existing literatures using projection, retraction, and intermediate matrices (see Sec. I.I requires separate adjustments of learning rates for the two types of parameters. As an illustration, in projected SGD and Adam on Stiefel manifold Li al. (2020). learning rate may vary 1Os of times for constrained and unconstrained parameters. This can be intuitively understood as the result of applying different optimization methods for different parameters. However, using various learning rates may lead to the divergence of the algorithms.
In contrast, our Stiefel method can be established for both the constrained and unconstrained parame- ters with the same learning rate. Due to the same derivation of numerical methods for these two types, the effort in tuning hyperparameters is reduced while there is still good performance. In practice, the learning rate need to be adjusted in Stiefel SGD and momentum can be chosen as 0.9 for most machine learning tasks. For Stiefel Adam, we recommend to use 10 to be the learning rate and (B1,80 (2) (0.9, 0.999).
25

---

Published as a conference paper at ICLR 2023
Method
MNIST
Shakespear
Original (RGAS)
1) 8c -4
1f = 2
Original (RAGAS)
7 0.01, 3 = 0.8
= 0.08, 3 0.9
Stiefel SGD (ours))
1 le μ 0.5
"J = 15, = 0.5
2. 1r 0.5
Projected Stiefel SGDLi et al.
= Te 4, = 0.5
(2020)
Momentumless Stiefel SGD
6e 4
= 2
Wen and Yin (2013)
Table 3: Details for hyperparameters of optimizers in PRW test. 7) stands for learning rate, stands for momentum parameter and stands for the parameter that adjusts stepsize automatically.
How to choose between Adam-Stiefel and momentum SGD-Stiefel Generally, we recommend choosing momentum SGD-Stiefel for traditional optimization and small scale neural networks that hyperparameters can be carefully adjusted while for large scale neural networks, use Adam-Stiefel as default.
P.2 DETAILS OF PRW
For MNIST experiment, we use a pretrained CNN to extract 128-dim. features figures in MNIST. For Shakespeare's plays experiment, compute the PRW distances between all pairs of items in a corpus of eight Shakespeare's ' operas embedded into 200-dim. space using word2vec. For each digit or play pair, we use the extracted features point sets {æ;} and {u.]. We choose the dimension of the target space of the projection to be k = 2 in both experiments. The mean optimal transport values are taken among all digits or movie pairs the specific iteration. All the setting
as the original paper Lin et al. (2020) except that early termination is removed The hyperparameters for each optimizer are listed Tab.
P.3 DETAILS OF VIT EXPERIMENTS
The structure of ViT and training process exactly same as https://github... com/
but use different ViT implementation from https:::pp::::taaa
The hyperparameters are carefully adjusted for each opti- mizer to obtain the best performance. For projected Stiefel SGD/Adam Li al. (2020), tune the learning rates for constrained and unconstrained parameters separately. For regularizer SGD/Adam, the regularizer scaling parameter also carefully chosen. The largest learning rate for momentumless Stiefel SGD Wen and Yin (2013) is applied but momentum methods still admit faster convergence. Additional experimental results can seen Fig.
The images are cut into 4 patches. dmodet dl of 384, Thead 12. dy = dg = dg = 32 (which also equals m Rhead, so 'otthogonality across heads' constraint can be applied: see Section 3.2). The 'classification token' in Dosovitskiy et al. (2020) is used. Total number layer of our VIT model is 7.
structure structure
http::: the implementation of (Cubuk et al. 2018) from AutoAugment]
Data augmentation http://www.aa....oooo/.coooaaa/a
Training All the training the same scheduler, which is to let the learning rate linearly increase from to the target learning rate in 5 epochs, followed by cosine annealing learning rate Loshchiloy 0.01 times max learning rate. Label
and Hutter (2016) with restart, and minimum learning smoothing with parameter 0.1 is used, weight decay- = 5e - unless specified.
All the hyperparameters used for optimizers are listed in Tab. 4
"Adam-type methods not included since they suitable for this problem.
26

---

Published as a conference paper at ICLR 2023
0.4
0.105
0.3-
0.100
0.095
0.090
0.2
0.085
0,080
100
170
180
190
200
lype
170
190
method on CIFARIO
baseline (No constraints)
Mody Kerpeaapage
Projected (artho across)
0.1
Ours (ortho, across)
Adam-type method on CIFAR10
Regularizer (ortho. across)
Projected [ortho: on within)
25
50
75
100
125
150
175
200
Ours (ortho, only within) Regularizer (ortha. only within)
cooth
e99th
Momentumless (ortho, only within]
0.6
0.6
0.36
0.35
0.34
0.34
05
0.31
wrror
170
190
SGD-type method on CIFAR100
Adam-type method on CIFAR100
0.3
0.3
50
75
100
125
150
175
200
25
50
75
100
125
150
175
200
opach
epach
Figure 2: Test errors ViT on CIFAR. We can see that our optimizer has the best test accuracy in both cases. Mcanwhile, when our optimizer is used, orthogonality only within each head performs better than orthogonality across heads.
Method
Hyperparameters
Stiefel SGD (Algo. 1h)
foothh Thon moon-orth = 0.15, = 0.9 North = 1, moon-orth = 0.1, μr = 0.9
Projected Sticfel SGD et al.
2020
Regularizer SGD
0.1, 0.9, penalty weight = le-6
Momentumless Stiefel SGD
North 0.1. For non-orth parameters, Thon -oth 0.1, 1- 0.9
Wen and Yin 2013
Stiefel Adam (Algo.
North 1 0.001, 0.9, 0.999
Projected Stiefel Adam Li et al.
North 0.001, Thon = 0.01, S1 = 0.9, B3 0.999
2020)
Regularizer Adam
0.001, 81 = 0.9, B2 0.999, penalty weight = 1e-6
SGD
North 0.1, f 0.9
Adam Kingma and Ba 2015)
0.001, 81 = 0.9, B2 0.999 = 0.001, 31 = 0.9, B2 = 0.999, weight decay 0.2
AdamW [Loshchllov and Hutter
2017)
Table 4: Details for hyperparameters of optimizers in ViT training. v stands for learning rate, μ stands for momentum parameter and 31, B2 are parameters in Adam.
27

---

Published as a conference paper at ICLR 2023
Method Stiefel SGD (ours)
Hyperparameters
" 0.1, μ 0.9 = 0.001, 31 = 0.9, 32 = 0.999
Stiefel Adam (ours)
Projected Stiefel SGDLI al.
"1 0.2, μ 0.9
(2020)
Momentumless Stiefel
SGD
17 0.1
Wen and Yin (2013)
Table Details for hyperparameters of optimizers in leading eigenvalue problem. 1 stands for learning rate, J stands for momentum parameter and 31 and 82 stands for the hyperparamcters Adam.
Q ADDITIONAL EXPERIMENT: SYSTEMATIC TESTS ON THE LEADING EIGENVALUE PROBLEM
In this section, we consider the leading eigenvalue problem: given an n-by-n matrix A, find the m largest eigenvalues of This problem at the core of many data sciences tasks, where n can be very large but m remains small. Due to its relative simplicity and the existence of exact solution, it's possible to implement a large amount of experiments in order to systematically investigate the convergence, manifold perseverance, and time consumption of various approaches. Particularly, this problem can be formulated as an optimization problem on St(n,, m) as follows
arg max Tr(X (TX) XESt(n.m)
(39)
The idea is seeks optimal m-dimensional subspace, represented by X via orthonormal bases in IR" corresponding to its m columns, so that eigenvalues projected onto this subspace sums up maximum value.
Such formulation is not unique. For example, Tao and Ohsawa (2020) consider applying their memnentuu-acceleraan SO(n) optimizer to arg min Tr(E' R ARE), where the constant padding matrix E := mm: 0(n-m)xml However, their setting is computationally less efficient than our Stiefel simplification. In fact, our formulation will be particularly suitable to cases when m small but n is very large.
Our test uses a deterministic matrix A generated by A= (E of n-by-n matrix with i.i.d. random normal elements. The exact solution obtainable from matrix decomposition so that the еггог can be quantified. As is mentioned above, several aspects of our algorithm are examined, including convergence speed, manifold perseverance, and computational methods. We can conclude that
where E is an instance
complexity in Fig. 3 as well as other exact our (S)GD- method has the fastest convergence rate and exact manifold perseverance with the lowest computational complexity.
Additionally, in Fig. we test the influence of different canonical-type metric, i.e., different a in Def. [ and the three different ways to 'compute' matrix exponential in Apdx. M.2x i.e., the matrix exponential itself, Cayley map, and forward Euler. No significant difference is found leading eigenvalue test and as result, the most commonly used canonical metric and the cheapest forward Euler are chosen as default all the other experiments. The hyperparameter used for each algorithm listed in Tab.
We also experimentally prove that better manifold preservation leads to better convergence, for both algorithm and Projected Stiefel SGD (Li et al. 2020) in 5 What's more, this experiment shows our (inner) iterative solver is quadratically convergent (as opposed to the linearly convergent Cayley approximation in Li al. (2020)), and therefore smaller number of iteration is needed to reach machine precision.
"The time recorded excludes the corresponding part computing gradients. Both the algorithm in paper Li et al, (2020) and their code has O(n m) computational complexity per iteration. However, this can be improved by changing the order of computing matrix production by associative law. Also, this Onnmm is the complexity under fixed number of iterations for approximating Cayley map according to their setup. It will be changed to Onmm² (1 log(1/u) /u])) if Cayley map is computed to machine precision.
28

---

Published as a conference paper at ICLR 2023
10r
10:
Convergence
Manifold preservance
1-
10"*
10a5
deviaton
10
10
10-12-
10-11
10-13
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
Iter
Iter
10a
Time complexity against
time consuming CPU (te))
10-7-
Srietel kdam [0urs]
Scietel 5GD [Dur)
(fix m=10)
Projected Stirhd
Projected Sticfel Adam
Uime consuming CPU
10
Nomentuu Stiee sGD
10-2
10
complexity against (fix r/m==00
10-1
10a
10a
102
10
Figure 3: Comparison of exact manifold preserving methods on leading eigenvalue problem. The algorithms are performed CPU with learning of each adjusted best performance. The upper figures show that our Stiefel SGD optimizer
best manifold perseverance. The lower two figures experimentally prove that we have the lowest (ram²) complexity per iteration, which matches our complexity analysis in Apdx. Apdx.[Gand [andd table 2 Notice that our algorithm and Monent umless Stiefel SGD have almost coinciding curves when lower curves than Pro jrotected Stei fel
the fastest convergence
(d), while both algorithms have much when m + ac. This means that our algorithm has a similar constant factor as the complexity of Moment umless Stiefel (S) GD (Wen and Yin (013), indicating that the introduction of momentum in our algorithm is almost 'free' and meanwhile, Projected Steifel
in (c) and when
SGD
Li et al. 2020) has much larger prefactor. Please see Apdx. for more discussions.
101
Natrix Exp=Euclidea
10a
Matrix Exp+Eudlidesn
Exp Canonical
Exp Cananical
10-1
Cayley map+ Fudldean
Gayley map Buclideann
Daylty mup+Ganonical
Gayley map+Canonica
10-1
Forward Eullr-Euclidean
10-3-
Forward Euler+Euclidean
Forward Euler Cananical
EuleriCanonical
loss
10
10-7
10
10"
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
(a) SGD
(b) Adam
Figure 4: Comparision of different canonical-type metrics (Euclidean and canonical metrics are tested) and different ways to approximate matrix exponential in Apdx. M.2
Though our special retraction on tangent bundle (polar retraction for position and extra handling for momentum) is cheap, it must be used with our specially designed algorithm and cannot be applied other algorithms even the position will still stay on Stiefel manifold. The reason is that the loss of structure of momentum will lead to slow convergence. Fig 6 shows that projected Stiefel SGD with our retraction for tangent bundle convergences slower than not only our algorithm but also the original projected Stiefel SGD.
29

---

Published as a conference paper at ICLR 2023
Garvegerce
Outoo mn
10
31-22
30.2
06
(a) Projected Stiefel SGD et al. (2020) using different numbers Cayley loops
(b) Projected Stiefel SGD et (2020) using QR retractions in dif- forent frequencies
(c) Algo. 4 using different num- ber of inner loops for matrix root inversion
Projected Stiefel SGD, performing more inner iterations that leads a more accurate Cayley transform gives a better optimized function value; (b) shows that more frequently applying QR retractions, i.e. projecting the variable, provides hetter accuracy: (c) shows that more accurate matrix root inversion for polar retraction in algorithm Serrithummlllads
Figure (a) shows
performing 8 iterations (note this curve coincides 100 iterations) of matrix root inversion gives perfect accuracy our method, compared (a) where 16 iterations of Cayley transform only give an error ~v 10" times bigger.
accurate solution. We
Convergance
Manifold pridenkince
101
Ourretact ES LP SANAIS-O (anro)
Carmesscion Uugge
Stkit S0D-0ban)
10--
10-
38-4
2044
2012
1:2 17.0
2ar
6ar
Figure 6: Convergence and manifold preservation when our low-cost retraction (polar retraction for position and extra handling for momentum) applied projected Stiefel SGD (Li al. (020). Though our design also helps the projected Stiefel SGD preserve the manifold structure with only machine precision error, the convergence to minimum slower than both our algorithm and original projected Stiefel SGD.
30


Reviews:
[{"review_id": "PSRmX5Jwss", "reviewer": "Reviewer_y3Tv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 1, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_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": "", "main_review": "", "summary_of_the_paper": "This paper presents a Momentum Stiefel Optimizer with several pleasant properties. ", "summary_of_the_review": "The paper is well structured and well written; Since I was not familiar with this field and not able to confirm every single proof/algorithm, I am unable to assess this paper.", "strength_and_weaknesses": "The paper is well structured and well written; Since I was not familiar with this field and not able to confirm every single proof/algorithm, I am unable to assess this paper.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "HWm0wZb_hSA", "reviewer": "Reviewer_AYdG", "summary": "", "questions": "", "limitations": "", "rating": 10, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_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": "", "main_review": "", "summary_of_the_paper": "This paper develops a gradient-based solution for optimizations if objective functions on a Stiefel manifold. It uses the geometry of Sitefel manifolds along with momentum-based ideas for speeding up gradient searches to develop an intrinsic solution for the optimization. It is quite elegant in its use of geometry and in ensuring that the iterative updates stay on the required spaces (manifold and the tangent bundle). I did not follow all the details in the construction but trust that the paper is solid in its development. From an application point of view, the focus is on training neural networks where orthogonality is important. The experimental results involve training ViT from scratch using different orthogonality constraints and the proposed method is clearly found to be better. ", "summary_of_the_review": "\nA strong paper on solving general optimization problems on Stiefel manifolds. Impressive level of details in ", "strength_and_weaknesses": "Strengths: \n\nThis is a strong paper on exploiting the differental geometry of Stiefel manifolds and bringing some of the momentum-based improvements from Euclidean domains to this manifold. The experimental results are quite satisfactory. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "DE87rC6fJRU", "reviewer": "Reviewer_TSJE", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_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": "", "main_review": "", "summary_of_the_paper": "This papers considers the problem of optimizing a function on the Stiefel manifold. The main objective of the paper is to build algorithms with momentum, which is often difficult / expensive in a Riemannian context. The key idea of the authors is to consider an ODE with momentum that converges to the solutions of the problem, and then to have a judicious discretization process that produces feasible iterates, while maintaining a relatively low cost (only one projection on the manifold should be computed at each iteration). The authors also extend the celebrated adam algorithm in this setting.\n\nThe authors then demonstrate the promises of their approach to train vision transformers with orthogonal attention heads and for computing projection robust wasserstein distances.", "summary_of_the_review": "This is a paper that this easy and pleasant to read, that proposes a new method for an important problem in ML and that convincingly demonstrates its utility in practice. However, the lack of any theoretical result is problematic to me.", "strength_and_weaknesses": "# Strengths \n- This paper considers a timely problem : optimization on the stiefel manifold arises naturally in many machine learning problems and deep learning problems, and considering the success of momentum-based techniques  / adam in the Euclidean case, it is important to have extensions of these methods to the stiefel manifold\n\n- the proposed algorithm seems sound, and is computationally quite cheap (on par with classical Riemannian SGD since we have 1 projection per iteration).\n\n- The numerical evaluation is satisfactory, the advantage of the proposed algorithm is clearly illustrated.\n\n# Weakness\n\n- The striking weakness of this paper is the lack of any theoretical result regarding the actual convergence of the proposed method. This is an optimization paper and yet there is no convergence result (while the method seems to work well in practice). This is a clear hole in the plot of the paper: what can we theoretically expect from this method ? Do the iterates of algorithm 1 decrease the loss function?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "-36nl9wPvF", "reviewer": "Reviewer_RJLh", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "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": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_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": "", "main_review": "", "summary_of_the_paper": "The authors considered the problem of optimization on Stiefel manifold. Heuristically, the proposed method can be divided into the following steps: 1) formulate a variational principle which naturally constraints the trajectory on a Stiefel manifold, hence it exactly preserves the manifold structure without extra operations such as projection / retraction 2) With a smart discretization in time of the derived ODEs, the authors came to an discrete-time optimization scheme.\n\nThe provided thorough theoretical analysis, including: 1) detailed derivation of the variational principle 2) proof of the structure preserving property 3) discretization scheme\n\nFinally, the authors studied the benefits of orthogonality, including orthogonality at initialization, within an attention head, and across attention heads. The author validated the proposed method in experiments.", "summary_of_the_review": "Overall the paper proposed a novel method to solve the Stiefel optimization problem without extra steps like projection / retraction. The idea is novel and well validated.", "strength_and_weaknesses": "Strengths:\n1) The proposed method looks pretty interesting and novel to me.The authors had a clear presentation from the general idea of constraint optimization on the manifold as a continuous-time process with continuous-time constraints, then to a proper discretization scheme.\n2) The problem to be studied is important. The authors have clearly demonstrated that orthogonality helps with the generalization performance of machine learning models, and have demonstrated the benefits of the proposed method with extensive studies.\n3) The theoretical analysis and proofs are solid. I only loosely checked the theoretical analysis and it looks pretty convincing to me.\n\nWeakness:\n1) It might be better to start with toy examples to introduce the idea of Stiefel optimization, especially for readers like me who are not familiar with this specific field.\n2) I would suggest the authors conduct experiments on more challenging datasets rather than Cifar.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster