Original: M OMENTUM S TIEFEL O PTIMIZER , WITH A PPLICATIONS TO S UITABLY -O RTHOGONAL A TTENTION , AND O PTIMAL T RANSPORT
Revised: M OMENTUM S TIEFEL O PTIMIZER , WITH A PPLICATIONS TO S UITABLY -O RTHOGONAL A TTENTION , AND O PTIMAL T RANSPORT

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: The problem of optimization on Stiefel manifold, i.e., minimizing functions of (not necessarily square) matrices that satisfy orthogonality constraints, has been extensively studied.
Revised: The problem of optimization on Stiefel manifold, i.e., minimizing functions of (not necessarily square) matrices that satisfy orthogonality constraints, has been extensively studied.

Original: Yet, a new approach is proposed based on, for the first time, an interplay between thoughtfully designed continuous and discrete dynamics.
Revised: Yet, a new approach is proposed based on, for the first time, an interplay between thoughtfully designed continuous and discrete dynamics.

Original: It leads to a gradient-based optimizer with intrinsically added momentum.
Revised: It leads to a gradient-based optimizer with intrinsically added momentum.

Original: 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.
Revised: 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.

Original: Its generalization to adaptive learning rates is also demonstrated.
Revised: Its generalization to adaptive learning rates is also demonstrated.

Original: Notable performances are observed in practical tasks.
Revised: Notable performances are observed in practical tasks.

Original: 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.
Revised: 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.

Original: This optimizer also makes the useful notion of Projection Robust Wasserstein Distance (Paty and
Revised: This optimizer also makes the useful notion of Projection Robust Wasserstein Distance (Paty and

Original: Cuturi, 2019; Lin et al., 2020) for high-dim.
Revised: Cuturi, 2019; Lin et al., 2020) for high-dim.

Original: optimal transport even more effective.
Revised: optimal transport even more effective.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Matrices that satisfy orthogonal constraints are known to play important roles in various areas including machine learning.
Revised: Matrices that satisfy orthogonal constraints are known to play important roles in various areas including machine learning.

Original: A wide range of studies have proven this both theoretically (Saxe et al., 2013; Xiao et al., 2018) and experimentally – for example, orthogonality can boost the performances of a wide range of architectures, e.g., MLP (Cisse et al., 2017), CNN and ResNet (Li et al., 2020; Wang et al., 2020), RNN (Arjovsky et al., 2016; Wisdom et al., 2016; Helfrich et al., 2018),and Transformer (Zhang et al., 2021).
Revised: A wide range of studies have proven this both theoretically (Saxe et al., 2013; Xiao et al., 2018) and experimentally – for example, orthogonality can boost the performances of a wide range of architectures 1 , e.g., MLP (Cisse et al., 2017), CNN and ResNet (Li et al., 2020; Wang et al., 2020), RNN (Arjovsky et al., 2016; Wisdom et al., 2016; Helfrich et al., 2018), Transformer (Zhang et al., 2021).

Original: 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.
Revised: 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.

Original: 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
Revised: 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

Original: Wasserstein distance (Paty and Cuturi, 2019; Lin et al., 2020).
Revised: Wasserstein distance (Paty and Cuturi, 2019; Lin et al., 2020).

Original: Therefore, this article will focus on this non-Euclidean / Riemannian optimization problem.
Revised: Therefore, this article will focus on this non-Euclidean / Riemannian optimization problem.

Original: In fact, matrices that contain orthogonal columns constitute a Riemannian manifold known as the
Revised: In fact, matrices that contain orthogonal columns constitute a Riemannian manifold known as the

Original: Stiefel manifold.
Revised: Stiefel manifold.

Original: Given integers n ≥ m > 0 , it is defined as St ( n, m ) :
Revised: Given integers n ≥ m > 0 , it is defined as St ( n, m ) :

Original: = { X ∈ R n × m : X ⊤ X =
Revised: = { X ∈ R n × m : X ⊤ X =

Original: I m × m } (see Apdx.
Revised: I m × m } (see Apdx.

Original: B.3 for the n = m special case which is almost SO ( n ) ).
Revised: B.3 for the n = m special case which is almost SO ( n ) ).

Original: Then, given a Cfunction f : St ( n, m ) → R , we consider, under gradient oracle, the smooth optimization problem
Revised: Then, given a Cfunction f : St ( n, m ) → R , we consider, under gradient oracle, the smooth optimization problem

Original: [Equation]
Revised: [Equation]

Original: which is equivalent to
Revised: which is equivalent to

Original: [Equation]
Revised: [Equation]

Original: Due to the importance of this problem, a wide range of approaches have been proposed.
Revised: Due to the importance of this problem, a wide range of approaches have been proposed.

Original: 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;
Revised: 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;

Original: Wang et al., 2020; Zhang et al., 2021).
Revised: Wang et al., 2020; Zhang et al., 2021).

Original: 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 on existing milestones in Sec. 1.1.
Revised: 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 on existing milestones in Sec. 1.1.

Original: Our strategy toward a good Stiefel optimizer is the following: first, formulate a variational principle and use that to construct an optimizing dynamics in continuous time, which is described by a system of ODEs corresponding to damped mechanical systems on a 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.
Revised: Our strategy toward a good Stiefel optimizer is the following: first, formulate a variational principle and use that to construct an optimizing dynamics in continuous time, which is described by a system of ODEs corresponding to damped mechanical systems on a 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.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Typically, in order to maintain the manifold structure, some kind of projection/retraction/exponential-map operation is needed, and since we have both position and momentum, such operation is needed for both variables (i.e. to maintain the cotangent bundle structure).
Revised: Typically, in order to maintain the manifold structure, some kind of projection/retraction/exponential-map operation is needed, and since we have both position and momentum, such operation is needed for both variables (i.e. to maintain the cotangent bundle structure).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: We used a quadratic-convergent iterative solver for our specific position retraction operation, which makes it fast.
Revised: We used a quadratic-convergent iterative solver for our specific position retraction operation, which makes it fast.

Original: Due to 2)+3), our per iteration computational complexity, O ( nm 2 ) , has a small constant factor (see Sec.
Revised: Due to 2)+3), our per iteration computational complexity, O ( nm 2 ) , has a small constant factor (see Sec.

Original: C for details).
Revised: C for details).

Original: 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.
Revised: 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.

Original: 6) Because our 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.3 and note this difference from previous milestones (e.g., Li et al.
Revised: 6) Because our 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.3 and note this difference from previous milestones (e.g., Li et al.

Original: (2020)) significantly reduces tuning efforts.
Revised: (2020)) significantly reduces tuning efforts.

Original: Our algorithm works for a range of Riemannian metrics, allowing extra flexibility in choosing suitable geometry to optimize the performance for a specific problem.
Revised: Our algorithm works for a range of Riemannian metrics, allowing extra flexibility in choosing suitable geometry to optimize the performance for a specific problem.

Original: 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.
Revised: 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.

Original: We use it to systematically investigate algorithmic performances under different parameters.
Revised: It systematically investigates algorithmic performances under different parameters.

Original: (2) We show the elegant idea ofefficiently approximating optimal transport distance in high-dim.
Revised: (2) We show the elegant idea of approximating optimal transport distance in high-dim.

Original: via a good low-dim. projection (Paty and Cuturi, 2019; Lin et al., 2020) can be made even more efficacious by our optimizer.
Revised: via a good low-dim. projection (Paty and Cuturi, 2019; Lin et al., 2020) can be made even more efficacious by our optimizer.

Original: We note that Vision Transformer (ViT) can be further improved by imposing attention heads to be orthogonal; more precisely:
Revised: We note that Vision Transformer (ViT) can be further improved by imposing attention heads to be orthogonal; more precisely:

Original: Consider training ViT from scratch.
Revised: Consider training ViT from scratch.

Original: We discover that 1) requiring each head to be orthogonal in itself improves both training and testing accuracies the most.
Revised: We discover that 1) requiring each head to be orthogonal in itself improves both training and testing accuracies the most.

Original: An important recent work by Zhang et al.
Revised: An important recent work by Zhang et al.

Original: (2021) applied orthogonality to transformers and demonstrated improved performance in NLP tasks.
Revised: (2021) applied orthogonality to transformers and demonstrated improved performance in NLP tasks.

Original: It concatenates each of the W Qi , W Ki , W Vi matrices in attention across all heads, and applies orthogonal constraint to each of the three, via regularizer.
Revised: It concatenates each of the W Qi , W Ki , W Vi matrices in attention across all heads, and applies orthogonal constraint to each of the three, via regularizer.

Original: This makes each head (approximately) orthogonal, not only within itself, but also to other heads.
Revised: This makes each head (approximately) orthogonal, not only within itself, but also to other heads.

Original: Orthogonal constraint is also applied, via regularizer, to each weight matrix of feed-forward layers in their case.
Revised: Orthogonal constraint is also applied, via regularizer, to each weight matrix of feed-forward layers in their case.

Original: 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 at least in CV tasks.
Revised: 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 at least in CV tasks.

Original: Meanwhile, 2) having 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).
Revised: Meanwhile, 2) having 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).

Original: No orthogonality on feed-forward layers was used in either 1) or 2).
Revised: No orthogonality on feed-forward layers was used in either 1) or 2).

Original: 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.
Revised: 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.

Original: 1. M ORE ON RELATED WORK
Revised: 1. M ORE ON RELATED WORK

Original: Orthogonality in deep learning Initialization: Orthogonal initialization is both theoretically and experimentally shown good for deep neural networks (Saxe et al., 2013).
Revised: Orthogonality in deep learning Initialization: Orthogonal initialization is both theoretically and experimentally shown good for deep neural networks (Saxe et al., 2013).

Original: CNN: Cisse et al.
Revised: CNN: Cisse et al.

Original: (2017) used regularizer to make weight matrices orthogonal when training MLP and CNN, and showed both improved accuracy and robustness.
Revised: (2017) used regularizer to make weight matrices orthogonal when training MLP and CNN, and showed both improved accuracy and robustness.

Original: Rodríguez et al.
Revised: Rodríguez et al.

Original: (2017) showed that posing orthogonality to CNN reduces overfitting.
Revised: (2017) showed that posing orthogonality to CNN reduces overfitting.

Original: Bansal et al.
Revised: Bansal et al.

Original: (2018) experimentally compared several orthogonal regularizers on
Revised: (2018) experimentally compared several orthogonal regularizers on

Original: ResNet.
Revised: ResNet.

Original: RNN: Arjovsky et al.
Revised: RNN: Arjovsky et al.

Original: ; Wisdom et al.
Revised: ; Wisdom et al.

Original: ; Helfrich et al.
Revised: ; Helfrich et al.

Original: ; Vorontsov et al.
Revised: ; Vorontsov et al.

Original: (2017); Lezcano-Casado and Martınez-Rubio (2019) showed orthogonalilty helps plain RNN achieve long term memory and even out-perform advanced models such as LSTM.
Revised: (2017); Lezcano-Casado and Martınez-Rubio (2019) showed orthogonalilty helps plain RNN achieve long term memory and even out-perform advanced models such as LSTM.

Original: Transformer:
Revised: Transformer:

Original: 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).
Revised: 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).

Original: Our work is the first to apply it to Vision Transformer (Dosovitskiy et al., 2020).
Revised: Our work is the first to apply it to Vision Transformer (Dosovitskiy et al., 2020).

Original: Optimization algorithms Manifold optimization is a profound field.
Revised: Optimization algorithms Manifold optimization is a profound field.

Original: General approaches for 1st-order (i.e. gradient based) optimization on (Riemannian) manifold typically involve retraction
Revised: General approaches for 1st-order (i.e. gradient based) optimization on (Riemannian) manifold typically involve retraction

Original: (see Apdx.D) which is not present in Euclidean cases.
Revised: (see Apdx.D) which is not present in Euclidean cases.

Original: Thus we categorize 1st-order Stiefel optimizers into 2 types: O1) retraction-based, and O2) non-retraction-based.
Revised: Thus we categorize 1st-order Stiefel optimizers into 2 types: O1) retraction-based, and O2) non-retraction-based.

Original: O1) Retraction-based.
Revised: O1) Retraction-based.

Original: Given a point and a tangent vector, exponential map gives how the point moves along the tangent vector on the manifold.
Revised: Given a point and a tangent vector, exponential map gives how the point moves along the tangent vector on the manifold.

Original: In practice, an approximation of the exponential map, known as a retraction, is often used for reducing the computational cost.
Revised: In practice, an approximation of the exponential map, known as a retraction, is often used for reducing the computational cost.

Original: The exponential map and various retractions on Stiefel manifold are well-studied.
Revised: The exponential map and various retractions on Stiefel manifold are well-studied.

Original: For Stiefel, the exponential map can be computed with complexity O
Revised: For Stiefel, the exponential map can be computed with complexity O

Original: nm 2 )
Revised: nm 2 )

Original: (Edelman et al., 1998), whose constant prefactor is however large as the map is essentially matrix exponentiation.
Revised: (Edelman et al., 1998), whose constant prefactor is however large as the map is essentially matrix exponentiation.

Original: A well-known retraction is Cayley map (Eq.
Revised: A well-known retraction is Cayley map (Eq.

Original: and Wen and Yin (2013) proposed a gradient descent method 1 with smartly lowered computational cost of Cayley map (will be referred to as Momentumless Stiefel (S)GD ).
Revised: and Wen and Yin (2013) proposed a gradient descent method 2 with smartly lowered computational cost of Cayley map (will be referred to as Momentumless Stiefel (S)GD ).

Original: SVD can also help construct a retraction; Li et al.
Revised: SVD can also help construct a retraction; Li et al.

Original: (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.
Revised: (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.

Original: Meanwhile, a remaining challenge is to add momentum, which is nontrivial in a manifold setting but helpful for improved speed of convergence (Zhang and Sra, 2018;
Revised: Meanwhile, a remaining challenge is to add momentum, which is nontrivial in a manifold setting but helpful for improved speed of convergence (Zhang and Sra, 2018;

Original: Ahn and Sra, 2020; Alimisis et al., 2021).
Revised: Ahn and Sra, 2020; Alimisis et al., 2021).

Original: 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.
Revised: 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.

Original: One could follow a general Riemannian retraction idea (e.g., Absil and Malick (2012)) and use retraction that projects both position and momentum to the tangent bundle to ensure manifold preservation.
Revised: One could follow a general Riemannian retraction idea (e.g., Absil and Malick (2012)) and use retraction that projects both position and momentum to the tangent bundle to ensure manifold preservation.

Original: This was done for example in Li et al.
Revised: This was done for example in Li et al.

Original: (2020) (will be referred to as Projected
Revised: (2020) (will be referred to as Projected

Original: Stiefel SGD/Adam ).
Revised: Stiefel SGD/Adam ).

Original: Another cool technique to maintain the structure of momentum is parallel transport (e.g., Bécigneul and Ganea (2018)), which is however computationally expensive and possibly inexact (Edelman et al., 1998).
Revised: Another cool technique to maintain the structure of momentum is parallel transport (e.g., Bécigneul and Ganea (2019)), which is however computationally expensive and possibly inexact (Edelman et al., 1998).

Original: Our work could be viewed as a retraction-based approach, because our position variable is projected back to St at each step; however, it doesn’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 T X St .
Revised: Our work could be viewed as a retraction-based approach, because our position variable is projected back to St at each step; however, it doesn’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 T X St .

Original: This leads to better computational efficiency, as well as improved accuracy because variables stay closer to the tangent bundle.
Revised: This leads to better computational efficiency, as well as improved accuracy because variables stay closer to the tangent bundle.

Original: O2) Non-retraction-based.
Revised: O2) Non-retraction-based.

Original: In some sense one can always call an exact manifold preserving method retraction-based, but approximate structure preservation is also a possibility.
Revised: In some sense one can always call an exact manifold preserving method retraction-based, but approximate structure preservation is also a possibility.

Original: Regularizers for instance could be used to approach X ⊤ X
Revised: Regularizers for instance could be used to approach X ⊤ X

Original: = I .
Revised: = I .

Original: Between exact and approximate structure preserving, which one works better is theoretically still an open question, and task-dependent but insightful empirical results exist.
Revised: Between exact and approximate structure preserving, which one works better is theoretically still an open question, and task-dependent but insightful empirical results exist.

Original: For example, regularizer-based methods usually need more hyperparameter tuning, and the final result is likely sensitive to the regularization strength (Wang et al., 2020).
Revised: For example, regularizer-based methods usually need more hyperparameter tuning, and the final result is likely sensitive to the regularization strength (Wang et al., 2020).

Original: They typically have low computational costs, but due to trajectory not exactly on the manifold, they may converge to (the neighborhood of) a local min.
Revised: They typically have low computational costs, but due to trajectory not exactly on the manifold, they may converge to (the neighborhood of) a local min.

Original: different from the one an exact algorithm converges to for multimodal problems.
Revised: different from the one an exact algorithm converges to for multimodal problems.

Original: Similar observation was made for non-regularizer-based approximate method too; e.g., Vorontsov et al.
Revised: Similar observation was made for non-regularizer-based approximate method too; e.g., Vorontsov et al.

Original: (2017) suggested that at least in certain cases ‘hard’ constraint is better than ‘soft’.
Revised: (2017) suggested that at least in certain cases ‘hard’ constraint is better than ‘soft’.

Original: The aforementioned Li et al.
Revised: The aforementioned Li et al.

Original: (2020) is actually inexact too, because in their practical implementation the retraction was inexactly performed for the sake of speed.
Revised: (2020) is actually inexact too, because in their practical implementation the retraction was inexactly performed for the sake of speed.

Original: Meanwhile, we also note an interesting, non-regularizer-based recent result (Ablin and Peyré, 2022) that converges to the manifold despite of being approximate prior to convergence.
Revised: Meanwhile, we also note an interesting, non-regularizer-based recent result (Ablin and Peyré, 2022) that converges to the manifold despite of being approximate prior to convergence.

Original: Additional.
Revised: Additional.

Original: We also mention the existence of useful algorithms for electronic structure calculations that can be viewed as Stiefel optimizers (e.g., Zhang et al.
Revised: We also mention the existence of useful algorithms for electronic structure calculations that can be viewed as Stiefel optimizers (e.g., Zhang et al.

Original: ; Dai et al.
Revised: ; Dai et al.

Original: ; Gao et al.
Revised: ; Gao et al.

Original:  Hu et al.
Revised: (2018; 2019); Hu et al.

Original: Dai et al.
Revised: Dai et al.

Original: They did not use momentum.
Revised: They did not use momentum.

Revised: The problem can also be approached via a constrained optimization formulation (e.g., (Leimkuhler et al., 2021)).

Original: In addition, there are other interesting optimizers also based on discretizing damped mechanical systems but not via the splitting approach used here, e.g., Lee et al.
Revised: In addition, there are other interesting optimizers also based on discretizing damped mechanical systems but not via the splitting approach used here, e.g., Lee et al.

Original: ; Duruisseaux and Leok (2021); however they are mostly implicit and not always scalable to deep learning applications.
Revised: ; Duruisseaux and Leok (2021); however they are mostly implicit and not always scalable to deep learning applications.

Original: Finally, it is possible to represent a matrix by other matrices and optimize those matrices instead.
Revised: Finally, it is possible to represent a matrix by other matrices and optimize those matrices instead.

Original: A special case of SO ( n ) was considered by Lezcano-Casado and Martınez-Rubio (2019);
Revised: A special case of SO ( n ) was considered by Lezcano-Casado and Martınez-Rubio (2019);

Original: Helfrich et al.
Revised: Helfrich et al.

Original: where X = expm ( A ) or 
Revised: where X = expm ( A ) or X

Original: X =
Revised: =

Original: Cayley ( A ) was used and then one optimized skew symmetric A instead.
Revised: Cayley ( A ) was used and then one optimized skew symmetric A instead.

Original: Although these authors did not consider a Stiefel generalization, this could be possible, but even without the generalization the computational cost is already high.
Revised: Although these authors did not consider a Stiefel generalization, this could be possible, but even without the generalization the computational cost is already high.

Original: This paper considers exact structure preservation.
Revised: This paper considers exact structure preservation.

Original: For applications included here, this is either empirically more beneficial (Sec.3.2) or a necessity (Sec.3.1 & Apdx.Q).
Revised: For applications included here, this is either empirically more beneficial (Sec.3.2) or a necessity (Sec.3.1 & Apdx.Q).

Original: See comparisons of most relevant approaches in Tab. 2 and more on complexity in Apdx.
Revised: See comparisons of most relevant approaches in Tab. 2 and more on complexity in Apdx.

Original: C.
Revised: C.

Original:  D ERIVATION OF THE O PTIMIZERS AND T HEIR P ROPERTIES
Revised:  D ERIVATION OF THE O PTIMIZERS AND T HEIR P ROPERTIES

Original: We represent Stiefel manifold as St ( n, m ) :
Revised: We represent Stiefel manifold as St ( n, m ) :

Original: = { X ∈ R n × m : X ⊤ X
Revised: = { X ∈ R n × m : X ⊤ X

Original: = I m × m } where n ≥ m 2 , i.e. matrices with orthonormal columns.
Revised: = I m × m } where n ≥ m 3 , i.e. matrices with orthonormal columns.

Original: This is a natural embedding in Euclidean space.
Revised: This is a natural embedding in Euclidean space.

Original: Based on identity isomorphism between cotangent & tangent spaces (see Apdx.
Revised: Based on identity isomorphism between cotangent & tangent spaces (see Apdx.

Original: B.2), it also gives:
Revised: B.2), it also gives:

Original: Proposition 1.
Revised: Proposition 1.

Original: The (co)tangent space of St at X ∈ St is T X St := { ∆ ∈ R n × m : X ⊤ ∆ + ∆ ⊤ X = 0 } ; its (co)tangent bundle is T St := { ( X, ∆) : X ∈ St , ∆ ∈ T X St } .
Revised: The (co)tangent space of St at X ∈ St is T X St := { ∆ ∈ R n × m : X ⊤ ∆ + ∆ ⊤ X = 0 } ; its (co)tangent bundle is T St := { ( X, ∆) : X ∈ St , ∆ ∈ T X St } .

Original: Throughout this paper, we focus on a family of Riemannian metrics on St defined in the following.
Revised: Throughout this paper, we focus on a family of Riemannian metrics on St defined in the following.

Original: Definition 1 (canonical-type metric) .
Revised: Definition 1 (canonical-type metric) .

Original: For a fixed constant a < 1 , the canonical-type metric g X :
Revised: For a fixed constant a < 1 , the canonical-type metric g X :

Original: T X St × T X
Revised: T X St × T X

Original: St → R is defined as
Revised: St → R is defined as

Original: [Equation]
Revised: [Equation]

Original: The following are two commonly used examples of the canonical-type metric (Tagare, 2011):
Revised: The following are two commonly used examples of the canonical-type metric (Tagare, 2011):

Original: • When a = 0 , g X (∆ 1 , ∆ 2 ) =
Revised: • When a = 0 , g X (∆ 1 , ∆ 2 ) =

Original: Tr (∆ ⊤ 1 ∆ 2 ) , which corresponds to the Euclidean metric.
Revised: Tr (∆ ⊤ 1 ∆ 2 ) , which corresponds to the Euclidean metric.

Original: • When a = 1 / 2 , g X (∆ 1 , ∆ 2 ) =
Revised: • When a = 1 / 2 , g X (∆ 1 , ∆ 2 ) =

Original: Tr (∆ ⊤ 1 ( I − 12 XX ⊤ )∆ 2 ) , which is known as the canonical metric.
Revised: Tr (∆ ⊤ 1 ( I − 12 XX ⊤ )∆ 2 ) , which is known as the canonical metric.

Original: Notation.
Revised: Notation.

Original: Denote Euclidean gradient in ambient space by ∂f∂X :
Revised: Denote Euclidean gradient in ambient space by ∂f∂X :

Original: = (cid:0) ∂f∂X ij ( X ) (cid:1) ij .
Revised: = (cid:0) ∂f∂X ij ( X ) (cid:1) ij .

Original: Denote by ⊙ , ⊘ , and A ◦ c element-wise product, division, and c th power.
Revised: Denote by ⊙ , ⊘ , and A ◦ c element-wise product, division, and c th power.

Original: Denote expm to be the matrix exponential.
Revised: Denote expm to be the matrix exponential.

Original: 2. O PTIMIZATION DYNAMICS IN CONTINUOUS TIME
Revised: 2. O PTIMIZATION DYNAMICS IN CONTINUOUS TIME

Original: In this section, we study the Stiefel optimization problem (1) from a variational perspective.
Revised: In this section, we study the Stiefel optimization problem (1) from a variational perspective.

Original: The manifold setup yields a constrained variational problem, which is hard to handle due to nonlinear geometry of the function space.
Revised: The manifold setup yields a constrained variational problem, which is hard to handle due to nonlinear geometry of the function space.

Original: Thus we introduce an alternative variational formulation based on functional Lagrange multiplier to bypass this difficulty.
Revised: Thus we introduce an alternative variational formulation based on functional Lagrange multiplier to bypass this difficulty.

Original: It generates an ODE that never escapes from the manifold T St and is guaranteed to converge to a local minimizer of the function f on St .
Revised: It generates an ODE that never escapes from the manifold T St and is guaranteed to converge to a local minimizer of the function f on St .

Original: In Euclidean space, momentum has long been used to accelerate gradient descent (Nesterov, 1983).
Revised: In Euclidean space, momentum has long been used to accelerate gradient descent (Nesterov, 1983).

Original: More recently, continuous time limit brought useful tools to complement the classical analyses for such optimizers (e.g., Su et al.
Revised: More recently, continuous time limit brought useful tools to complement the classical analyses for such optimizers (e.g., Su et al.

Original: (2014)), and a variational formulation was then established to unify a large family of momentum-based Euclidean optimizers (Wibisono et al., 2016).
Revised: (2014)), and a variational formulation was then established to unify a large family of momentum-based Euclidean optimizers (Wibisono et al., 2016).

Original: This formulation can also provide a powerful and intrinsic way to generalize momentum-based optimization to nonEuclidean settings.
Revised: This formulation can also provide a powerful and intrinsic way to generalize momentum-based optimization to nonEuclidean settings.

Original: One demonstration is, for example, accelerated optimization on Lie groups (Tao and Ohsawa, 2020).
Revised: One demonstration is, for example, accelerated optimization on Lie groups (Tao and Ohsawa, 2020).

Original: Following the same line of thoughts, we consider Stiefel optimization by first defining a time-dependent Lagrangian L : T St × R + → R as
Revised: Following the same line of thoughts, we consider Stiefel optimization by first defining a time-dependent Lagrangian L : T St × R + → R as

Original: [Equation]
Revised: [Equation]

Original: Without r ( t ) , this would be a standard Lagrangian for mechanics on manifold, and it would generate dynamics in which the kinetic energy 12 g X ( ˙ X, ˙ X ) and the potential energy f
Revised: Without r ( t ) , this would be a standard Lagrangian for mechanics on manifold, and it would generate dynamics in which the kinetic energy 12 g X ( ˙ X, ˙ X ) and the potential energy f

Original: ( X ) continuously exchange with each other.
Revised: ( X ) continuously exchange with each other.

Original: Choosing r ( t ) to be an increasing function, however, will break time translational symmetry and introduce dissipation in the system; consequently, the total energy 12 g X ( ˙ X, ˙ X )+ f ( X ) will monotonically decrease in time, leading to the minimization of f .
Revised: Choosing r ( t ) to be an increasing function, however, will break time translational symmetry and introduce dissipation in the system; consequently, the total energy 12 g X ( ˙ X, ˙ X )+ f ( X ) will monotonically decrease in time, leading to the minimization of f .

Original: Given this
Revised: Given this

Original: Lagrangian, we can define a variational problem (VP) on manifold whose corresponding stationarity condition (known as Euler-Lagrange equation) yields an ODE that optimizes f in continuous time
Revised: Lagrangian, we can define a variational problem (VP) on manifold whose corresponding stationarity condition (known as Euler-Lagrange equation) yields an ODE that optimizes f in continuous time

Original: Constrained
Revised: Constrained

Original: [Equation]
Revised: [Equation]

Original: However, to obtain concrete algorithms, this constrained VP is very difficult to handle because the variation of X ( · ) has to keep it in a nonlinearly constrained function space.
Revised: However, to obtain concrete algorithms, this constrained VP is very difficult to handle because the variation of X ( · ) has to keep it in a nonlinearly constrained function space.

Original: Tao and Ohsawa (2020) applied a tool of reduction via Lie algebra to solve the case of Lie group manifold, but for
Revised: Tao and Ohsawa (2020) applied a tool of reduction via Lie algebra to solve the case of Lie group manifold, but for

Original: Stiefel manifold St ( n, m ) , unless n = m , there is no group structure and this technique no longer applies.
Revised: Stiefel manifold St ( n, m ) , unless n = m , there is no group structure and this technique no longer applies.

Original: Therefore, we transfer (4) into an unconstrained variational problem, using tools described in, e.g., Chen et al.
Revised: Therefore, we transfer (4) into an unconstrained variational problem, using tools described in, e.g., Chen et al.

Original: By adding a Lagrange multiplier 3 Λ( t ) ∈ R m × m for the constraints
Revised: By adding a Lagrange multiplier 4 Λ( t ) ∈ R m × m for the constraints

Original: X ( t ) ⊤ X ( t ) − I = 0 , we augment the Lagrangian to be
Revised: X ( t ) ⊤ X ( t ) − I = 0 , we augment the Lagrangian to be

Original: [Equation]
Revised: [Equation]

Original: where Λ( t ) ∈
Revised: where Λ( t ) ∈

Original: R m × m is a symmetric matrix.
Revised: R m × m is a symmetric matrix.

Original: The benefit is that now X can be varied in an unconstrained, flat function space, corresponding to
Revised: The benefit is that now X can be varied in an unconstrained, flat function space, corresponding to

Original: [Equation]
Revised: [Equation]

Original: To sum up the above discussion, problem (1) can be resolved through (un)constrained variational problem (4),(6), i.e., (1) ⇐ = (4) ⇐⇒ (6), in the sense that the solutions of (4),(6) solve problem (1); meanwhile, (6) can be explicitly solved, as detailed in the following (proof is in Apdx.
Revised: To sum up the above discussion, problem (1) can be resolved through (un)constrained variational problem (4),(6), i.e., (1) ⇐ = (4) ⇐⇒ (6), in the sense that the solutions of (4),(6) solve problem (1); meanwhile, (6) can be explicitly solved, as detailed in the following (proof is in Apdx.

Original: G):
Revised: G):

Original: Theorem 1 (Optimization dynamics on T St ) .
Revised: Theorem 1 (Optimization dynamics on T St ) .

Original: To use problem (6) for solving problem (1), we have:
Revised: To use problem (6) for solving problem (1), we have:

Original: The solution of the unconstrained variational problem (6) is the following ODE (cid:40)
Revised: The solution of the unconstrained variational problem (6) is the following ODE (cid:40)

Original: [Equation]
Revised: [Equation]

Original: where ( X ( t ) , Q ( t ))
Revised: where ( X ( t ) , Q ( t ))

Original: =
Revised: =

Original: ( X ( t )
Revised: ( X ( t )

Original: ˙ X ( t ))
Revised: ˙ X ( t ))

Original: ∈ R n × m × R n × m is the tuple of state variable and its momentum, γ ( t ) := ˙
Revised: ∈ R n × m × R n × m is the tuple of state variable and its momentum, γ ( t ) := ˙

Original: r ( t ) /r ( t ) is the scalar friction coefficient, and b :=
Revised: r ( t ) /r ( t ) is the scalar friction coefficient, and b :=

Original: aa − 1 is a constant depending on the canonical-type metric (2).
Revised: aa − 1 is a constant depending on the canonical-type metric (2).

Original: For any isolated local minimum X ∗ ∈ St of f , there exists a neighbourhood U ⊂ T St of ( X ∗ , 0 ) s.t., for any initial condition ( X 0 , Q 0 )
Revised: For any isolated local minimum X ∗ ∈ St of f , there exists a neighbourhood U ⊂ T St of ( X ∗ , 0 ) s.t., for any initial condition ( X 0 , Q 0 )

Original: ∈ U , the solution ( X ( t ) , Q ( t )) of the system (7) converges to ( X ∗ , 0 ) as t → ∞ .
Revised: ∈ U , the solution ( X ( t ) , Q ( t )) of the system (7) converges to ( X ∗ , 0 ) as t → ∞ .

Original: One feature of ODE (7) is the preservation of constraints: even though the ODE is defined in the
Revised: One feature of ODE (7) is the preservation of constraints: even though the ODE is defined in the

Original: Euclidean space, as long as it starts on the manifold, it will never leave the manifold:
Revised: Euclidean space, as long as it starts on the manifold, it will never leave the manifold:

Original: Theorem 2 (Constrained optimization with unconstrained dynamics) .
Revised: Theorem 2 (Constrained optimization with unconstrained dynamics) .

Original: If the initial condition of (7) is on T St , the cotangent bundle of Stiefel manifold, then dynamics (7) automatically stays on T St , i.e., if X (0) ⊤ X (0)
Revised: If the initial condition of (7) is on T St , the cotangent bundle of Stiefel manifold, then dynamics (7) automatically stays on T St , i.e., if X (0) ⊤ X (0)

Original: =
Revised: =

Original: I m × m , X (0)
Revised: I m × m , X (0)

Original: ⊤ Q (0) + Q (0)
Revised: ⊤ Q (0) + Q (0)

Original: ⊤ X (0) = 

Revised: ⊤ X (0) =

Original: m × m , , then for all t ≥
Revised: m × m , , then for all t ≥

Original: X ( t ) ⊤ X ( t ) =
Revised: X ( t ) ⊤ X ( t ) =

Original: I m × m , X ( t )
Revised: I m × m , X ( t )

Original: ⊤ Q ( t ) + Q ( t )
Revised: ⊤ Q ( t ) + Q ( t )

Original: ⊤ X ( t ) = 

Revised: ⊤ X ( t ) =

Original: m × m .
Revised: m × m .

Original: Proof is in Apdx.
Revised: Proof is in Apdx.

Original: H .
Revised: H .

Original: 2. S TRUCTURE - PRESERVING DISCRETIZATION VIA VARIABLE DECOMPOSITION AND OPERATOR SPLITTING
Revised: 2. S TRUCTURE - PRESERVING DISCRETIZATION VIA VARIABLE DECOMPOSITION AND OPERATOR SPLITTING

Original: Although the continuous dynamics (7) preserves the constraints, such preservation is in general not guaranteed when time is discretized.
Revised: Although the continuous dynamics (7) preserves the constraints, such preservation is in general not guaranteed when time is discretized.

Original: The construction of our discretization briefly follows four steps: a geometric decomposition of momentum, a carefully designed operator splitting scheme to approximate the ODEs, structure-preserving approximations of the split system, and a structurepreserving relaxation that further reduces the computational cost.
Revised: The construction of our discretization briefly follows four steps: a geometric decomposition of momentum, a carefully designed operator splitting scheme to approximate the ODEs, structure-preserving approximations of the split system, and a structurepreserving relaxation that further reduces the computational cost.

Original: Details now follow.
Revised: Details now follow.

Original: Preparation: from a static decomposition of the tangent space to a decomposition of Q dynamics.
Revised: Preparation: from a static decomposition of the tangent space to a decomposition of Q dynamics.

Original: To retain the preservation of geometric structures through a time discretization, we first decompose the tangent space T X St into X and X ⊥ components (see more details in Tagare (2011)); more precisely, given a tangent vector Q represented by an n × m matrix, we rewrite it as Q = XY + V for Y ∈ R m × m , V ∈ R n × m , and use Y, V to replace Q .
Revised: To retain the preservation of geometric structures through a time discretization, we first decompose the tangent space T X St into X and X ⊥ components (see more details in Tagare (2011)); more precisely, given a tangent vector Q represented by an n × m matrix, we rewrite it as Q = XY + V for Y ∈ R m × m , V ∈ R n × m , and use Y, V to replace Q .

Original: This transformation changes the constraint
Revised: This transformation changes the constraint

Original: X ⊤ Q + Q ⊤ X
Revised: X ⊤ Q + Q ⊤ X

Original: = 0 to { Y ⊤ + Y = 0 , X ⊤ V = 0 } instead (see Apdx E.1).
Revised: = 0 to { Y ⊤ + Y = 0 , X ⊤ V = 0 } instead (see Apdx E.1).

Original: Another advantage of this decomposition is that Y, V naturally split the canonical-type metric (2) (see Apdx E.2).
Revised: Another advantage of this decomposition is that Y, V naturally split the canonical-type metric (2) (see Apdx E.2).

Original: Although a fixed tangent Q can be uniquely decomposed into Y 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
Revised: Although a fixed tangent Q can be uniquely decomposed into Y 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

Original: Y ( t ) and V ( t ) evolve.
Revised: Y ( t ) and V ( t ) evolve.

Original: This is not a trivial question, but it can be proved that the new dynamics is
Revised: This is not a trivial question, but it can be proved that the new dynamics is

Original: [Equation]
Revised: [Equation]

Original: with initial condition satisfying X (0)
Revised: with initial condition satisfying X (0)

Original: ⊤ X (0)
Revised: ⊤ X (0)

Original: =
Revised: =

Original: I, Y (0) ⊤ + Y (0) =
Revised: I, Y (0) ⊤ + Y (0) =

Original: 0 , X (0)
Revised: 0 , X (0)

Original: ⊤ V (0)
Revised: ⊤ V (0)

Original: = 0 .
Revised: = 0 .

Original: This system is equivalent to (7) via Q ( t ) =
Revised: This system is equivalent to (7) via Q ( t ) =

Original: X ( t ) Y ( t )
Revised: X ( t ) Y ( t )

Original: + V ( t ) and preserving X ( t )
Revised: + V ( t ) and preserving X ( t )

Original: ⊤ X ( t )
Revised: ⊤ X ( t )

Original: =
Revised: =

Original: I, Y ( t ) ⊤ + Y ( t ) =
Revised: I, Y ( t ) ⊤ + Y ( t ) =

Original: , X ( t )
Revised: , X ( t )

Original: ⊤ V ( t ) = 

Revised: ⊤ V ( t ) =

Original: for all t (Thm.8).
Revised: for all t (Thm.8).

Original: Moreover, it will be amenable to a good discretization and is thus the base for constructing our optimizer.
Revised: Moreover, it will be amenable to a good discretization and is thus the base for constructing our optimizer.

Original: With this decomposition defined, we can finally make the phrase ‘structure preservation’ precise:
Revised: With this decomposition defined, we can finally make the phrase ‘structure preservation’ precise:

Original: Definition 2.
Revised: Definition 2.

Original: Structure preservation means variables exactly satisfying all constraints for all time.
Revised: Structure preservation means variables exactly satisfying all constraints for all time.

Original: That is, for ‘ XQ ’-system, X ( t ) ⊤ X ( t ) =
Revised: That is, for ‘ XQ ’-system, X ( t ) ⊤ X ( t ) =

Original: I, X ( t )
Revised: I, X ( t )

Original: ⊤ Q ( t ) = 0 , ∀ t ; for ‘ XY V ’-system,
Revised: ⊤ Q ( t ) = 0 , ∀ t ; for ‘ XY V ’-system,

Original: X ( t ) ⊤ X ( t )
Revised: X ( t ) ⊤ X ( t )

Original: =
Revised: =

Original: I, Y ( t ) ⊤ + Y ( t ) =
Revised: I, Y ( t ) ⊤ + Y ( t ) =

Original: , X ( t )
Revised: , X ( t )

Original: ⊤ V ( t ) = 0 , ∀ t .
Revised: ⊤ V ( t ) = 0 , ∀ t .

Original: In comparison, staying exactly on
Revised: In comparison, staying exactly on

Original: Stiefel, i.e. X ( t ) ⊤ X ( t ) =
Revised: Stiefel, i.e. X ( t ) ⊤ X ( t ) =

Original: I is termed ‘feasible’ (Wen and Yin, 2013; Ablin and Peyré, 2022), but we have additional constraints due to momentum.
Revised: I is termed ‘feasible’ (Wen and Yin, 2013; Ablin and Peyré, 2022), but we have additional constraints due to momentum.

Original: Step I: operator splitting of the ODEs.
Revised: Step I: operator splitting of the ODEs.

Original: To handle the high nonlinearity of ODE (8) and maintain the preservation of constraints after discretization, we adopt an operator splitting method, based on a general fact that a numerical discretization of an ODE can be obtained by composing the (approximate) flow maps of split ODEs (McLachlan and Quispel, 2002).
Revised: To handle the high nonlinearity of ODE (8) and maintain the preservation of constraints after discretization, we adopt an operator splitting method, based on a general fact that a numerical discretization of an ODE can be obtained by composing the (approximate) flow maps of split ODEs (McLachlan and Quispel, 2002).

Original: More precisely, we split the vector field of (8) as a sum of three vector fields, each associated with one of the following ODEs:
Revised: More precisely, we split the vector field of (8) as a sum of three vector fields, each associated with one of the following ODEs:

Original: [Equation]
Revised: [Equation]

Original: Define the corresponding time-
Revised: Define the corresponding time-

Original: h evolution maps ϕ 1 , ϕ 2 , ϕ 3 of Eq.(9)(10)(11) to be ϕ j : [ X ( t ) , Y ( t ) , V ( t )]
Revised: h evolution maps ϕ 1 , ϕ 2 , ϕ 3 of Eq.(9)(10)(11) to be ϕ j : [ X ( t ) , Y ( t ) , V ( t )]

Original: (cid:55)→ [ X ( t + h ) , Y ( t + h ) , V ( t + h )] for system j = 1 , 2 , 3 .
Revised: (cid:55)→ [ X ( t + h ) , Y ( t + h ) , V ( t + h )] for system j = 1 , 2 , 3 .

Original: Note ϕ 1 , ϕ 2 , ϕgive the exact solutions of these split ODEs.
Revised: Note ϕ 1 , ϕ 2 , ϕgive the exact solutions of these split ODEs.

Original: Then we see our specific split honors all constraints:
Revised: Then we see our specific split honors all constraints:

Original: Theorem 3.
Revised: Theorem 3.

Original: ϕ 1 , ϕ 2 , ϕ 3 are all structure preserving.
Revised: ϕ 1 , ϕ 2 , ϕ 3 are all structure preserving.

Original: Proof is in Apdx.
Revised: Proof is in Apdx.

Original: J .
Revised: J .

Original: Step II: structure-preserving approximation of exact flow maps.
Revised: Step II: structure-preserving approximation of exact flow maps.

Original: Due to the nonlinearity, ϕand ϕ 3 do not admit analytical expressions; ϕ 2 does have an explicit expression (Eq.
Revised: Due to the nonlinearity, ϕand ϕ 3 do not admit analytical expressions; ϕ 2 does have an explicit expression (Eq.

Original: but an approximation will still reduce computational costs while maintaining certain accuracy (see Fig.
Revised: but an approximation will still reduce computational costs while maintaining certain accuracy (see Fig.

Original: Therefore we first denote the 1st-order approximation of the exact flow maps ϕ 1 , ϕ 2 , ϕ 3 to be ¯ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 , where ¯ ϕ j : [ X 0 , Y 0 , V 0 ]
Revised: Therefore we first denote the 1st-order approximation of the exact flow maps ϕ 1 , ϕ 2 , ϕ 3 to be ¯ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 , where ¯ ϕ j : [ X 0 , Y 0 , V 0 ]

Original: (cid:55)→ [ X h , Y h , V h ] =
Revised: (cid:55)→ [ X h , Y h , V h ] =

Original: ϕ j
Revised: ϕ j

Original: ([ X 0 , Y 0 , V 0 ])
Revised: ([ X 0 , Y 0 , V 0 ])

Original: + O ( h 2 ) , j = 1 , 2 , 3 .
Revised: + O ( h 2 ) , j = 1 , 2 , 3 .

Original: Then ¯ ϕ 1 :
Revised: Then ¯ ϕ 1 :

Original: [Equation]
Revised: [Equation]

Original: Theorem 4.
Revised: Theorem 4.

Original: ¯ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 are all structure preserving.
Revised: ¯ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 are all structure preserving.

Original: Proof is in Apdx.
Revised: Proof is in Apdx.

Original: K .
Revised: K .

Original: This shows not only the split flow maps but their decently designed discretizations maintain the constraints of X, Y, V .
Revised: This shows not only the split flow maps but their decently designed discretizations maintain the constraints of X, Y, V .

Original: There are several specially designed features to enable the numerical stability of the integrators: 1) ‘ X † ( X ⊤† X † ) − 1 / 2 ’ in ¯ ϕ 3 is part of a nontrivial discretization scheme and the same as polar retraction (Absil et al., 2009).
Revised: There are several specially designed features to enable the numerical stability of the integrators: 1) ‘ X † ( X ⊤† X † ) − 1 / 2 ’ in ¯ ϕ 3 is part of a nontrivial discretization scheme and the same as polar retraction (Absil et al., 2009).

Original: It directly leads to the preservation of the geometry of the position variable by ¯ ϕ 3 , i.e., even the input of ¯ ϕ 3 has error that X ⊤ 

Revised: It directly leads to the preservation of the geometry of the position variable by ¯ ϕ 3 , i.e., even the input of ¯ ϕ 3 has error that X ⊤

Original: X 0 ̸
Revised: X 0 ̸

Original: = I , the output always satisfies X ⊤ h X h
Revised: = I , the output always satisfies X ⊤ h X h

Original: = I , but it will not impair the order of the O ( h 2 ) local discretization error.
Revised: = I , but it will not impair the order of the O ( h 2 ) local discretization error.

Original: See more about its connection to numerically stability to arithmetic errors in Apdx.
Revised: See more about its connection to numerically stability to arithmetic errors in Apdx.

Original: M.1.
Revised: M.1.

Original: In the first equation of ¯ ϕ 3 , X 0 + hV 0 X ⊤ 0 X 0 is used instead of X 0 + hV 0 , even though most of the time X ⊤ 0 X 

Revised: In the first equation of ¯ ϕ 3 , X 0 + hV 0 X ⊤ 0 X 0 is used instead of X 0 + hV 0 , even though most of the time X ⊤ 0 X

Original: = I .
Revised: = I .

Original: This guarantees that even if X ⊤ 0 X 0 ̸
Revised: This guarantees that even if X ⊤ 0 X 0 ̸

Original: = I , the constraint X ⊤ 0 V 0 = 0 itself still leads to X ⊤ h V h = 0 , which improves numerical stability.
Revised: = I , the constraint X ⊤ 0 V 0 = 0 itself still leads to X ⊤ h V h = 0 , which improves numerical stability.

Original: What’s more, when combined with 1), this property will also enable us to use a cheaper ˜ ϕ 1 in the following Step III.
Revised: What’s more, when combined with 1), this property will also enable us to use a cheaper ˜ ϕ 1 in the following Step III.

Original: The ‘forward Euler’-like discretization for ϕis carefully chosen,and updating V h using X h instead of X 0 , for example, will destroy its structure preservation.
Revised: The ‘forward Euler’-like discretization for ϕis carefully chosen,and updating V h using X h instead of X 0 , for example, will destroy its structure preservation.

Original: No extra handling except forward Euler is applied to the momentum variables Y, V ; i.e. this discretization itself is beneficial enough to guarantee momentum structure preservation.
Revised: No extra handling except forward Euler is applied to the momentum variables Y, V ; i.e. this discretization itself is beneficial enough to guarantee momentum structure preservation.

Original: Step III: relaxation and composition of the operators.
Revised: Step III: relaxation and composition of the operators.

Original: Our goal is to obtain a structure preserving scheme for the original ODE (8) instead of requiring structure preservation of each operator for the split ODEs (9)(10)(11).
Revised: Our goal is to obtain a structure preserving scheme for the original ODE (8) instead of requiring structure preservation of each operator for the split ODEs (9)(10)(11).

Original: The latter will ensure the former as we have:
Revised: The latter will ensure the former as we have:

Original: Theorem 5.
Revised: Theorem 5.

Original: The composition of any ordering of ¯ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 , e.g., ¯ ϕ 1 ◦ ¯ ϕ 2 ◦ ¯ ϕ 3 , is a structure-preserving, 1st-order (in h ) numerical integrator of (8).
Revised: The composition of any ordering of ¯ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 , e.g., ¯ ϕ 1 ◦ ¯ ϕ 2 ◦ ¯ ϕ 3 , is a structure-preserving, 1st-order (in h ) numerical integrator of (8).

Original: Proof is in Apdx.
Revised: Proof is in Apdx.

Original: L .
Revised: L .

Original: However, the latter (all operators preserving structure) is not needed for the former (their composition preserves structure).
Revised: However, the latter (all operators preserving structure) is not needed for the former (their composition preserves structure).

Original: We can still eventually obtain a structure-preserving integrator, however without the costly ‘expm’, by relaxing some of the structure preservation requirements for ¯ ϕ 1 .
Revised: We can still eventually obtain a structure-preserving integrator, however without the costly ‘expm’, by relaxing some of the structure preservation requirements for ¯ ϕ 1 .

Original: Theorem 6.
Revised: Theorem 6.

Original: Consider a consistent approximation of ¯ ϕ 1 by ˜ ϕ
Revised: Consider a consistent approximation of ¯ ϕ 1 by ˜ ϕ

Original: 1 (i.e. ˜ ϕ 

Revised: 1 (i.e. ˜ ϕ

Original: = ¯ ϕ 

Revised: = ¯ ϕ

Original: + O ( h 2 ) ).
Revised: + O ( h 2 ) ).

Original: Assume ˜
Revised: Assume ˜

Original: ϕ
Revised: ϕ

Original: 1 satisfies that if initially X ⊤ 0 X 

Revised: 1 satisfies that if initially X ⊤ 0 X

Original: =
Revised: =

Original: I, Y ⊤ 0 + Y 0 = 0 , X ⊤ 0 V 0 = 0 , then Y ⊤ h
Revised: I, Y ⊤ 0 + Y 0 = 0 , X ⊤ 0 V 0 = 0 , then Y ⊤ h

Original: + Y h
Revised: + Y h

Original: = 0 , X ⊤ h V h = 0 .
Revised: = 0 , X ⊤ h V h = 0 .

Original: Then the specific composition ¯ ϕ 3 ◦ ˜ ϕ
Revised: Then the specific composition ¯ ϕ 3 ◦ ˜ ϕ

Original: 1 ◦ ¯ ϕ 2 is structure preserving.
Revised: 1 ◦ ¯ ϕ 2 is structure preserving.

Original: Moreover, the composition of any ordering of ˜ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 is a 1st-order (in h ) integrator of (8).
Revised: Moreover, the composition of any ordering of ˜ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 is a 1st-order (in h ) integrator of (8).

Original: Proof is in Apdx.
Revised: Proof is in Apdx.

Original: M.
Revised: M.

Original: Therefore, our default recommendation is to use ¯ ϕ 3 ◦ ˜ ϕ
Revised: Therefore, our default recommendation is to use ¯ ϕ 3 ◦ ˜ ϕ

Original: 1 ◦ ¯ ϕ 2 , where ˜ ϕ 1 is X h
Revised: 1 ◦ ¯ ϕ 2 , where ˜ ϕ 1 is X h

Original: = X 0 + hX 

Revised: = X 0 + hX

Original: Y h ,
Revised: Y h ,

Original: Y h
Revised: Y h

Original: =
Revised: =

Original: (1 − γh )
Revised: (1 − γh )

Original: Y 0 − 1 − b 2 h (cid:16) X ⊤ 0 ∂f ∂X 0 − ∂f ∂X 0 ⊤ X 0 (cid:17) , V h
Revised: Y 0 − 1 − b 2 h (cid:16) X ⊤ 0 ∂f ∂X 0 − ∂f ∂X 0 ⊤ X 0 (cid:17) , V h

Original: = V 0 , mainly approximating the expm in ¯ ϕ by forward Euler.
Revised: = V 0 , mainly approximating the expm in ¯ ϕ by forward Euler.

Original: The result is concretely summarized in Algo.
Revised: The result is concretely summarized in Algo.

Original: In order to match commonly used notations in machine learning, we rescale the parameters: learning rate η := 1 − exp( − γh ) γ h , momentum parameter µ := exp( − γh ) ,
Revised: In order to match commonly used notations in machine learning, we rescale the parameters: learning rate η := 1 − exp( − γh ) γ h , momentum parameter µ := exp( − γh ) ,

Original: Z :=
Revised: Z :=

Original: Y/ 1 − exp( − γh ) γ , and U :=
Revised: Y/ 1 − exp( − γh ) γ , and U :=

Original: V/ 1 − exp( − γh ) γ .
Revised: V/ 1 − exp( − γh ) γ .

Original: See
Revised: See

Original: Apdx M.4 for more information.
Revised: Apdx M.4 for more information.

Original: A few side remarks are listed in the following.
Revised: A few side remarks are listed in the following.

Original: If m is small, then ¯ ϕ 1 ◦ ϕ 2 ◦ ¯ ϕ 3 (or any permutation) can also be used, although experimentally no clear advantage was observed (Fig. 4); otherwise, the computational cost of expm can become prohibitive.
Revised: If m is small, then ¯ ϕ 1 ◦ ϕ 2 ◦ ¯ ϕ 3 (or any permutation) can also be used, although experimentally no clear advantage was observed (Fig. 4); otherwise, the computational cost of expm can become prohibitive.

Original: Moreover, the computational cost of matrix inverse square root in the polar retraction ( X ⊤† X † ) − 1 /should not be a concern since it is only for m × m matrices instead of n × n ( n ≥ m ).
Revised: Moreover, the computational cost of matrix inverse square root in the polar retraction ( X ⊤† X † ) − 1 /should not be a concern since it is only for m × m matrices instead of n × n ( n ≥ m ).

Original: Meanwhile, it can be computed to machine precision rapidly using a quadratically convergent iterative method (Higham, 1997); see Algo. 3 in Apdx.
Revised: Meanwhile, it can be computed to machine precision rapidly using a quadratically convergent iterative method (Higham, 1997); see Algo. 3 in Apdx.

Original: M.1 for details, and Apdx.
Revised: M.1 for details, and Apdx.

Original: C for computational complexity.
Revised: C for computational complexity.

Original: Algorithm 1: Momentum (S)GD on St ( n, m ) (SGD: ∂f/∂X replaced by a stochastic estimator)
Revised: Algorithm 1: Momentum (S)GD on St ( n, m ) (SGD: ∂f/∂X replaced by a stochastic estimator)

Original: Hyperparameter : η ∈ (0 , + ∞ ) , µ ∈ [0 , 1) , maximum number of iterations N
Revised: Hyperparameter : η ∈ (0 , + ∞ ) , µ ∈ [0 , 1) , maximum number of iterations N

Original: Initialization : X 0 , U 0 , Z 0 s.t.
Revised: Initialization : X 0 , U 0 , Z 0 s.t.

Original: X ⊤ 

Revised: X ⊤

Original: X 

Revised: X

Original: =
Revised: =

Original: I , X ⊤ 0 U 0 = 0 , Z 0 + Z ⊤ 

Revised: I , X ⊤ 0 U 0 = 0 , Z 0 + Z ⊤

Original: =1 for i = 0 , · · · , N − 1 do
Revised: =1 for i = 0 , · · · , N − 1 do

Original: Compute ‘gradients’:
Revised: Compute ‘gradients’:

Original: f i = 1 − b 2 (cid:16) X ⊤ i ∂f∂X ( X i )
Revised: f i = 1 − b 2 (cid:16) X ⊤ i ∂f∂X ( X i )

Original: − ∂f ∂X ⊤ ( X i )
Revised: − ∂f ∂X ⊤ ( X i )

Original: X i (cid:17) ; g i =
Revised: X i (cid:17) ; g i =

Original: ( I − X i X ⊤ i )
Revised: ( I − X i X ⊤ i )

Original: ∂f∂X ( X i )
Revised: ∂f∂X ( X i )

Original: Update ¯ ϕ 2 : U i + 12 = µU i − 3 a − 2 2 ηU i Z i − g
Revised: Update ¯ ϕ 2 : U i + 12 = µU i − 3 a − 2 2 ηU i Z i − g

Original: i
Revised: i

Original: Update ˜ ϕ
Revised: Update ˜ ϕ

Original: 1 : Z i +1 =
Revised: 1 : Z i +1 =

Original: µZ i − f i ; X i + 

Revised: µZ i − f i ; X i +

Original: =
Revised: =

Original: X i + ηX
Revised: X i + ηX

Original: i
Revised: i

Original: Z
Revised: Z

Original: i
Revised: i

Original: Update ¯ ϕ 3 : X † =
Revised: Update ¯ ϕ 3 : X † =

Original: X i + 12 + ηU i + 12 X ⊤ i + 12 X
Revised: X i + 12 + ηU i + 12 X ⊤ i + 12 X

Original: i + 12 ; Compute ( X ⊤† X † ) − 12 using Algo. 3;
Revised: i + 12 ; Compute ( X ⊤† X † ) − 12 using Algo. 3;

Original: X i +

Revised: X i +

Original: = X † ( X ⊤† X † ) − 12 ; U i +

Revised: = X † ( X ⊤† X † ) − 12 ; U i +

Original: =
Revised: =

Original: U i + 12 − ηX
Revised: U i + 12 − ηX

Original: i
Revised: i

Original: + 12 U ⊤ i
Revised: + 12 U ⊤ i

Original: + 12 U
Revised: + 12 U

Original: i +6 end 7 return X N
Revised: i +6 end 7 return X N

Original: 2. A N ADAPTIVE LEARNING RATE VERSION
Revised: 2. A N ADAPTIVE LEARNING RATE VERSION

Original: 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).
Revised: 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).

Original: Thus in this section, we propose an adaptive version of our Stiefel optimizer.
Revised: Thus in this section, we propose an adaptive version of our Stiefel optimizer.

Original: 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.
Revised: 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.

Original: The algorithm is established via the following ideas.
Revised: The algorithm is established via the following ideas.

Original: The ‘gradient’ in this Adam-version method is constructed from our Stiefel SGD with momentum (Alg.1) where the ‘gradients’ in Y/Z and V/U direction can be interpreted as 1 − b 2 (cid:0) X ⊤ ∂f∂X − ∂f∂X ⊤ X (cid:1) and ( I − X ( X ⊤ X ) − 1 X ⊤ ) ∂f∂X respectively.
Revised: The ‘gradient’ in this Adam-version method is constructed from our Stiefel SGD with momentum (Alg.1) where the ‘gradients’ in Y/Z and V/U direction can be interpreted as 1 − b 2 (cid:0) X ⊤ ∂f∂X − ∂f∂X ⊤ X (cid:1) and ( I − X ( X ⊤ X ) − 1 X ⊤ ) ∂f∂X respectively.

Original: The main difficulty of extending Stiefel SGD to Stiefel Adam that does not appear in Euclidean case is that element-wise operation on momentum loses tangent vector structure.
Revised: The main difficulty of extending Stiefel SGD to Stiefel Adam that does not appear in Euclidean case is that element-wise operation on momentum loses tangent vector structure.

Original: We solve this respectively: (1) For Y/Z -direction, the skew-symmetry is preserved after a symmetric element-wise operation Z ⊘ ( p ◦ 12 + ϵ ) .
Revised: We solve this respectively: (1) For Y/Z -direction, the skew-symmetry is preserved after a symmetric element-wise operation Z ⊘ ( p ◦ 12 + ϵ ) .

Original: (2) For V/U -direction, we apply a projection I − X
Revised: (2) For V/U -direction, we apply a projection I − X

Original: ( X ⊤ X ) − 1 X ⊤ to the element-wisely rescaled momentum U ⊘ ( q ◦ 12 + ϵ ) , making sure ‘ X ⊤ V = 0 ’.
Revised: ( X ⊤ X ) − 1 X ⊤ to the element-wisely rescaled momentum U ⊘ ( q ◦ 12 + ϵ ) , making sure ‘ X ⊤ V = 0 ’.

Original: Combining all the above, we obtain the Adam-Stiefel optimizer.
Revised: Combining all the above, we obtain the Adam-Stiefel optimizer.

Original: Denote ˆ ϕ 1 , ˆ ϕ 2 , ˆ ϕ 3 to be the modification of ϕ 1 , ϕ 2 , ϕin the Adam version (see detailed expressions in Apdx. M.5).
Revised: Denote ˆ ϕ 1 , ˆ ϕ 2 , ˆ ϕ 3 to be the modification of ϕ 1 , ϕ 2 , ϕin the Adam version (see detailed expressions in Apdx. M.5).

Original: Then the integrator is defined as ˆ ϕ 3 ◦ ˆ ϕ 1 ◦ ˆ ϕ 2 .
Revised: Then the integrator is defined as ˆ ϕ 3 ◦ ˆ ϕ 1 ◦ ˆ ϕ 2 .

Original: The overall method is shown in Algo.
Revised: The overall method is shown in Algo.

Original: Theorem 7.
Revised: Theorem 7.

Original: The Adam-Stiefel ˆ ϕ 3 ◦ ˆ ϕ 1 ◦ ˆ ϕ 2 is structure-preserving.
Revised: The Adam-Stiefel ˆ ϕ 3 ◦ ˆ ϕ 1 ◦ ˆ ϕ 2 is structure-preserving.

Original: Proof is in Apdx.
Revised: Proof is in Apdx.

Original: N.
Revised: N.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: This section demonstrates our Stiefel optimizers on Projection Robust Wasserstein Distance (Paty and Cuturi, 2019; Lin et al., 2020) and trained-from-scratch Vision Transformer (Dosovitskiy et al., 2020).
Revised: This section demonstrates our Stiefel optimizers on Projection Robust Wasserstein Distance (Paty and Cuturi, 2019; Lin et al., 2020) and trained-from-scratch Vision Transformer (Dosovitskiy et al., 2020).

Original: They will also be compared with other popular Stiefel optimizers summarized in Tab. 2.
Revised: They will also be compared with other popular Stiefel optimizers summarized in Tab. 2.

Original: Canonical metric (i.e. a = 1 / 2 in Eq.2) is used in both examples to show that the gained performance is due to the algorithmic innovation but not an extra tuned knob.
Revised: Canonical metric (i.e. a = 1 / 2 in Eq.2) is used in both examples to show that the gained performance is due to the algorithmic innovation but not an extra tuned knob.

Original: An additional experiment on leading eigenvalue problem is deferred to Apdx.
Revised: An additional experiment on leading eigenvalue problem is deferred to Apdx.

Original: Q, where we also compare the rates of convergence and time consumptions of different algorithms, test different choices of metrics, and study why our retraction is tailored to our algorithm.
Revised: Q, where we also compare the rates of convergence and time consumptions of different algorithms, test different choices of metrics, and study why our retraction is tailored to our algorithm.

Original: Good performance is observed in all three examples.
Revised: Good performance is observed in all three examples.

Original: More experimental details are in Apdx.
Revised: More experimental details are in Apdx.

Original: P.
Revised: P.

Original: Hyperparameter : η ∈ (0 , + ∞ ) , β 1 ∈ [0 , 1) , β 2 ∈ [0 , 1) , 0 < ϵ ≪ 1 , number of iterations N
Revised: Hyperparameter : η ∈ (0 , + ∞ ) , β 1 ∈ [0 , 1) , β 2 ∈ [0 , 1) , 0 < ϵ ≪ 1 , number of iterations N

Original: Initialization : X 0 , V 0 , Y 0 , p 0 , q 0 s.t.
Revised: Initialization : X 0 , V 0 , Y 0 , p 0 , q 0 s.t.

Original: X ⊤ 

Revised: X ⊤

Original: X 

Revised: X

Original: =
Revised: =

Original: I , X ⊤ 0 U 0 = 0 , Z 0 + Z ⊤ 

Revised: I , X ⊤ 0 U 0 = 0 , Z 0 + Z ⊤

Original: = 0 , p 0 = p ⊤1 for i = 0 , · · · , N − 1 do
Revised: = 0 , p 0 = p ⊤1 for i = 0 , · · · , N − 1 do

Original: Compute ‘gradients’:
Revised: Compute ‘gradients’:

Original: f i = 1 − b 2 (cid:0) X ⊤ i ∂f∂X ( X i )
Revised: f i = 1 − b 2 (cid:0) X ⊤ i ∂f∂X ( X i )

Original: − ∂f∂X ( X i )
Revised: − ∂f∂X ( X i )

Original: ⊤ X i (cid:1) ; g i
Revised: ⊤ X i (cid:1) ; g i

Original: = ( I − X
Revised: = ( I − X

Original: i X ⊤ i )
Revised: i X ⊤ i )

Original: ∂f∂X ( X i ) 2nd-moment estimation: p
Revised: ∂f∂X ( X i ) 2nd-moment estimation: p

Original: i +1 =
Revised: i +1 =

Original: β 2 p
Revised: β 2 p

Original: i + (1 − β 2 ) f ◦ 2 i ; q i +1 =
Revised: i + (1 − β 2 ) f ◦ 2 i ; q i +1 =

Original: β 2 q i
Revised: β 2 q i

Original: + (1 − β 2 ) g ◦ 

Revised: + (1 − β 2 ) g ◦

Original: i
Revised: i

Original: Update ˆ ϕ 2 : U i + 12 = β 1 U i − 3 a − 2 2 ηU
Revised: Update ˆ ϕ 2 : U i + 12 = β 1 U i − 3 a − 2 2 ηU

Original: i Z i − (1 − β 1 ) g
Revised: i Z i − (1 − β 1 ) g

Original: i
Revised: i

Original: Update ˆ ϕ 1 : Z i +1 =
Revised: Update ˆ ϕ 1 : Z i +1 =

Original: β 1 Z i − (1 − β 1 ) f
Revised: β 1 Z i − (1 − β 1 ) f

Original: i ;
Revised: i ;

Original: X i + 

Revised: X i +

Original: =
Revised: =

Original: X i + η (cid:113) 1 − β
Revised: X i + η (cid:113) 1 − β

Original: i
Revised: i

Original: +12 X i (cid:16) Z i +1 ⊘ ( p ◦ 

Revised: +12 X i (cid:16) Z i +1 ⊘ ( p ◦

Original: i +1 + ϵ ) (cid:17)
Revised: i +1 + ϵ ) (cid:17)

Original: Update ˆ ϕ 3 :
Revised: Update ˆ ϕ 3 :

Original: ˜ U = (cid:113) 1 − β
Revised: ˜ U = (cid:113) 1 − β

Original: i +

Revised: i +

Original: ( I − X i + 

Revised: ( I − X i +

Original: ( X ⊤ i + 12 X
Revised: ( X ⊤ i + 12 X

Original: i + 12 ) − 1 X ⊤ i + 12 )( U i + 12 ⊘ ( q ◦ 

Revised: i + 12 ) − 1 X ⊤ i + 12 )( U i + 12 ⊘ ( q ◦

Original: i +1 + ϵ )) ;
Revised: i +1 + ϵ )) ;

Original: X †
Revised: X †

Original: =
Revised: =

Original: X i + 12 + η ˜
Revised: X i + 12 + η ˜

Original: UX ⊤ i +
Revised: UX ⊤ i +

Original: 12 X
Revised: 12 X

Original: i + 12 ; X i +

Revised: i + 12 ; X i +

Original: = X † ( X ⊤† X † ) − 12 ; U i +

Revised: = X † ( X ⊤† X † ) − 12 ; U i +

Original: =
Revised: =

Original: U i + 12 − ηX
Revised: U i + 12 − ηX

Original: i
Revised: i

Original: + 12 ˜
Revised: + 12 ˜

Original: U ⊤ U
Revised: U ⊤ U

Original: i +7 end 8 return X N
Revised: i +7 end 8 return X N

Original: Figure
Revised: Figure

Original: 3. P ROJECTION R OBUST W ASSERSTEIN D ISTANCE
Revised: 3. P ROJECTION R OBUST W ASSERSTEIN D ISTANCE

Original: Projection Robust Wasserstein Distance (PRW) (Paty and Cuturi, 2019; Lin et al., 2020) is a notion recently proposed to improve the robustness of standard Wasserstein metric, especially in highdim settings.
Revised: Projection Robust Wasserstein Distance (PRW) (Paty and Cuturi, 2019; Lin et al., 2020) is a notion recently proposed to improve the robustness of standard Wasserstein metric, especially in highdim settings.

Original: The idea is to simultaneously look for a best projection from high-dim data ( x i , y j , respectively with weights r i , c j ) to low-dim ones, and an entropic regularized optimal transport plan between projected data, i.e.
Revised: The idea is to simultaneously look for a best projection from high-dim data ( x i , y j , respectively with weights r i , c j ) to low-dim ones, and an entropic regularized optimal transport plan between projected data, i.e.

Original: [Equation]
Revised: [Equation]

Original: This problem is geodesically-nonconvex w.r.t.
Revised: This problem is geodesically-nonconvex w.r.t.

Original: the Stiefel variable U (Jordan et al., 2022), thus computationally extra-challenging to Stiefel optimizers.
Revised: the Stiefel variable U (Jordan et al., 2022), thus computationally extra-challenging to Stiefel optimizers.

Original: Lin et al.
Revised: Lin et al.

Original: (2020) proposed an effective method based on alternations between a full Sinkhorn step (Cuturi, 2013), given the current projection, and an optimization step that improves the projection.
Revised: (2020) proposed an effective method based on alternations between a full Sinkhorn step (Cuturi, 2013), given the current projection, and an optimization step that improves the projection.

Original: In particular, they developed Riemmanian optimizer with projection and retraction (RGAS) and its adaptive learning rate version (RAGAS).
Revised: In particular, they developed Riemmanian optimizer with projection and retraction (RGAS) and its adaptive learning rate version (RAGAS).

Original: We replace RGAS and RAGAS by our optimizer and others, and test on the hardest experiments in Lin et al.
Revised: We replace RGAS and RAGAS by our optimizer and others, and test on the hardest experiments in Lin et al.

Original: (2020), namely MNIST and Shakespeare.
Revised: (2020), namely MNIST and Shakespeare.

Original: Results are in Fig.1.
Revised: Results are in Fig.1.

Original: Our method is observed to find the largest value of PRW and thus the best projection among the tested, which implies the best performance.
Revised: Our method is observed to find the largest value of PRW and thus the best projection among the tested, which implies the best performance.

Original: Details of setup are in Apdx.
Revised: Details of setup are in Apdx.

Original: P.2.
Revised: P.2.

Original: H OW COULD ORTHOGONALITY IMPROVE VANILLA V ISION T RANSFORMER (V I T)?
Revised: H OW COULD ORTHOGONALITY IMPROVE VANILLA V ISION T RANSFORMER (V I T)?

Original: This section explores the possibility of making self-attention in Transformer models (Vaswani et al., 2017) orthogonal.
Revised: This section explores the possibility of making self-attention in Transformer models (Vaswani et al., 2017) orthogonal.

Original: The intuition is, if we interpret each attention head as a characterization of interactions between tokens, maybe linearly dependent interactions shouldn’t be double-counted so that the characterization can be more efficient and robust.
Revised: The intuition is, if we interpret each attention head as a characterization of interactions between tokens, maybe linearly dependent interactions shouldn’t be double-counted so that the characterization can be more efficient and robust.

Original: In addition, there is unlikely just one way in which tokens relate to each other, and multiple heads handle different characterizations in parallel.
Revised: In addition, there is unlikely just one way in which tokens relate to each other, and multiple heads handle different characterizations in parallel.

Original: So our second question is, should different heads also be orthogonal to each other?
Revised: So our second question is, should different heads also be orthogonal to each other?

Original: To formulate precisely, we follow the notation of Vaswani et al.
Revised: To formulate precisely, we follow the notation of Vaswani et al.

Original: (2017) and only review relevant parts: a Scaled Dot-Product
Revised: (2017) and only review relevant parts: a Scaled Dot-Product

Original: Multihead Attention is given by MultiHead ( Q, K, V ) =
Revised: Multihead Attention is given by MultiHead ( Q, K, V ) =

Original: Concat ( head 1 , ..., head n head ) W O , where head i =
Revised: Concat ( head 1 , ..., head n head ) W O , where head i =

Original: Attention ( QW Qi , KW K i , V W V i ) , Attention ( ˜ Q, ˜ K, ˜ V ) = softmax ( ˜ Q ˜ K ⊤ √ d k ) ˜ V , matrices
Revised: Attention ( QW Qi , KW K i , V W V i ) , Attention ( ˜ Q, ˜ K, ˜ V ) = softmax ( ˜ Q ˜ K ⊤ √ d k ) ˜ V , matrices

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: correspond to trainable parameters.
Revised: correspond to trainable parameters.

Original: The three input matrices Q , K and V all have dimension sequence _ length ×
Revised: The three input matrices Q , K and V all have dimension sequence _ length ×

Original: d model .
Revised: d model .

Original: d k and d v are usually smaller than d model .
Revised: d k and d v are usually smaller than d model .

Original: For orthogonality only within head, we require that W Qi , W Ki are in St ( d model , d k ) .
Revised: For orthogonality only within head, we require that W Qi , W Ki are in St ( d model , d k ) .

Original: This needs d model ≥ d k , which holds in most cases.
Revised: This needs d model ≥ d k , which holds in most cases.

Original: For orthogonality across heads, we need d model ≥ n head d k , which is satisfied in many popular models, and require Concat ( W Qi , i = 1 ..., n head ) , Concat ( W Ki , i = 1 ..., n head ) to be in St ( d model , n head d k ) , which means it contains not only ‘orthogonality only within head’, but also extra cross-head orthogonality.
Revised: For orthogonality across heads, we need d model ≥ n head d k , which is satisfied in many popular models, and require Concat ( W Qi , i = 1 ..., n head ) , Concat ( W Ki , i = 1 ..., n head ) to be in St ( d model , n head d k ) , which means it contains not only ‘orthogonality only within head’, but also extra cross-head orthogonality.

Original: We test these two types of constraints on vanilla ViT (Dosovitskiy et al., 2020), trained from scratch.
Revised: We test these two types of constraints on vanilla ViT (Dosovitskiy et al., 2020), trained from scratch.

Original: Results summarized in Tab. 1 (and Fig. 2 in Apdx.
Revised: Results summarized in Tab. 1 (and Fig. 2 in Apdx.

Original: 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.
Revised: 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.

Original: Note models trained from scratch are very different from pre-trained ViT. To have some realistic expectation of the performance of trained-from-scratch ViT, recall, for example, an improved ViT,
Revised: Note models trained from scratch are very different from pre-trained ViT. To have some realistic expectation of the performance of trained-from-scratch ViT, recall, for example, an improved ViT,

Original: PVT-T (Wang et al., 2021), 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), achieves 11.62% and 32.48% respectively with 5.3M parameters.
Revised: PVT-T (Wang et al., 2021), 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), achieves 11.62% and 32.48% respectively with 5.3M parameters.

Original: Our model in Tab.1 (and Fig.2 in Apdx.
Revised: Our model in Tab.1 (and Fig.2 in Apdx.

Original: P.3) uses 6.3M parameters.
Revised: P.3) uses 6.3M parameters.

Original: More details about this experiment can be found in Apdx.
Revised: More details about this experiment can be found in Apdx.

Original: P.3.
Revised: P.3.

Original: Remark 1 (No additional hyperparameters needed for the Stiefel part) .
Revised: Remark 1 (No additional hyperparameters needed for the Stiefel part) .

Original: Our Stiefel-ViT model has both non-Euclidean and Euclidean parameters, the latter due to the reason that we don’t impose orthogonality on feedforward layers and W Vi .
Revised: Our Stiefel-ViT model has both non-Euclidean and Euclidean parameters, the latter due to the reason that we don’t impose orthogonality on feedforward layers and W Vi .

Original: Our optimizer can use the same learning rate for both sets of parameters.
Revised: Our optimizer can use the same learning rate for both sets of parameters.

Original: This is because the optimizer is the time discretization of ODEs, and synchronous discretization is sufficient for convergence.
Revised: This is because the optimizer is the time discretization of ODEs, and synchronous discretization is sufficient for convergence.

Original: This contrasts with other approaches such as projected
Revised: This contrasts with other approaches such as projected

Original: Steifel SGD/Adam (Li et al., 2020), where learning rates for these two sets of parameters need to be adjusted separately and could differ by 40 times in some tasks.
Revised: Steifel SGD/Adam (Li et al., 2020), where learning rates for these two sets of parameters need to be adjusted separately and could differ by 40 times in some tasks.

Original: See Apdx.
Revised: See Apdx.

Original: P.1 for more discussion.
Revised: P.1 for more discussion.

Original: A A QUICK SUMMARY OF OPTIMIZERS WE EXPERIMENTALLY COMPARED TO IN THIS PAPER
Revised: A A QUICK SUMMARY OF OPTIMIZERS WE EXPERIMENTALLY COMPARED TO IN THIS PAPER

Original: [Table]
Revised: [Table]

Original: B T ECHNICAL R EMARKS ON G EOMETRIC M ECHANICS B. I NTUITION OF THE COTANGENT BUNDLE
Revised: B T ECHNICAL R EMARKS ON G EOMETRIC M ECHANICS B. I NTUITION OF THE COTANGENT BUNDLE

Original: If we require exact Stiefel manifold preservation for all time, i.e., X ( t ) ⊤ X ( t )
Revised: If we require exact Stiefel manifold preservation for all time, i.e., X ( t ) ⊤ X ( t )

Original: =
Revised: =

Original: I , a simple time differentiation gives a 2nd-order constraint X ( t )
Revised: I , a simple time differentiation gives a 2nd-order constraint X ( t )

Original: ⊤ Q ( t ) + Q ( t )
Revised: ⊤ Q ( t ) + Q ( t )

Original: ⊤ X ( t ) = 0 for Q = ˙
Revised: ⊤ X ( t ) = 0 for Q = ˙

Original: X .
Revised: X .

Original: This is why momentum Q has
Revised: This is why momentum Q has

Original: the additional structure Q ( t ) ∈ T ∗ X ( t )
Revised: the additional structure Q ( t ) ∈ T ∗ X ( t )

Original: St
Revised: St

Original: . B. W HAT E XACTLY IS M OMENTUM ?
Revised: . B. W HAT E XACTLY IS M OMENTUM ?

Original: Through this paper
Revised: Through this paper

Original: [Table]
Revised: [Table]

Original: we abused notation and called Q := ˙
Revised: we abused notation and called Q := ˙

Original: X momentum
Revised: X momentum

Original: [Table]
Revised: [Table]

Original: following the common practice of the community.
Revised: following the common practice of the community.

Original: It really should be called velocity instead, because although in canonical
Revised: It really should be called velocity instead, because although in canonical

Original: Euclidean spaces this doesn’t make much a difference, for mechanical systems on manifolds they are not the same thing.
Revised: Euclidean spaces this doesn’t make much a difference, for mechanical systems on manifolds they are not the same thing.

Original: For example, for our case, the geometrically intrinsic momentum variable should be given by Legendre transform
Revised: For example, for our case, the geometrically intrinsic momentum variable should be given by Legendre transform

Original: [Equation]
Revised: [Equation]

Original: instead, and velocity Q is in the tangent space while momentum P is in the cotangent space.
Revised: instead, and velocity Q is in the tangent space while momentum P is in the cotangent space.

Original: In the paper we used the identity map for an isomorphism between the tangent and cotangent spaces, but if we’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 (15).
Revised: In the paper we used the identity map for an isomorphism between the tangent and cotangent spaces, but if we’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 (15).

Original: None of these affect the correctness or the efficacy of results in this paper.
Revised: None of these affect the correctness or the efficacy of results in this paper.

Original: This discussion is just about terminology.
Revised: This discussion is just about terminology.

Original: B. T HE SPECIAL CASE OF St ( n, n ) AND ITS RELATION WITH O ( n ) AND SO ( n )
Revised: B. T HE SPECIAL CASE OF St ( n, n ) AND ITS RELATION WITH O ( n ) AND SO ( n )

Original: When n ≥ m , St ( n, m ) ∼ = O ( n ) / O ( n − m ) ; when n > m strictly, we also have St ( n, m ) ∼ =
Revised: When n ≥ m , St ( n, m ) ∼ = O ( n ) / O ( n − m ) ; when n > m strictly, we also have St ( n, m ) ∼ =

Original: SO ( n ) / SO ( n − m ) .
Revised: SO ( n ) / SO ( n − m ) .

Original: However, in the special case of n = m , St ( n, n ) ∼ = O ( n ) is not connected unlike the n > m cases.
Revised: However, in the special case of n = m , St ( n, n ) ∼ = O ( n ) is not connected unlike the n > m cases.

Original: 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 n = m , it is, to be precise, optimizing on SO ( n ) but not O ( n ) , although a similar complication is nonexistent when n > m .
Revised: 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 n = m , it is, to be precise, optimizing on SO ( n ) but not O ( n ) , although a similar complication is nonexistent when n > m .

Original: C D ETAILS ABOUT THE PER - ITERATION COMPLEXITY AND COMPUTATIONAL COST FOR OUR ALGORITHMS
Revised: C D ETAILS ABOUT THE PER - ITERATION COMPLEXITY AND COMPUTATIONAL COST FOR OUR ALGORITHMS

Original: The most costly operation in our algorithms is the n × m matrix multiplication (note n > m ).
Revised: The most costly operation in our algorithms is the n × m matrix multiplication (note n > m ).

Original: The computation for the matrix exponential and square root of matrix inversion, is cautiously designed to only deal with matrices of dimension m × m (see Apdx. M.1M.2), and thus admits at most O ( m 3 ) at each step (particularly, forward Euler only has the complexity of O ( m 2 ) while Cayley map is O ( m 3 ) ).
Revised: The computation for the matrix exponential and square root of matrix inversion, is cautiously designed to only deal with matrices of dimension m × m (see Apdx. M.1M.2), and thus admits at most O ( m 3 ) at each step (particularly, forward Euler only has the complexity of O ( m 2 ) while Cayley map is O ( m 3 ) ).

Original: For the inner loop of the square root of matrix inversion (Algo. 3), due to its quadratic convergence (Higham, 1997), it takes O (log log(1 / u )) number of steps to achieve the machine precision u .
Revised: For the inner loop of the square root of matrix inversion (Algo. 3), due to its quadratic convergence (Higham, 1997), it takes O (log log(1 / u )) number of steps to achieve the machine precision u .

Original: Note early stopping can be applied to this loop so that the complexity can be further reduced while the order of the overall method (1st-order) is still maintained.
Revised: Note early stopping can be applied to this loop so that the complexity can be further reduced while the order of the overall method (1st-order) is still maintained.

Original: The combination of all the above gives complexity O ( nm 2 ) + O ( m 3 log log(1 / u )) 4 .
Revised: The combination of all the above gives complexity O ( nm 2 ) + O ( m 3 log log(1 / u )) 5 .

Original: Commonly used optimizers on Stiefel manifold (e.g. Tab. 2) also have per-iteration complexity of
Revised: Commonly used optimizers on Stiefel manifold (e.g. Tab. 2) also have per-iteration complexity of

Original: O ( nm 2 ) .
Revised: O ( nm 2 ) .

Original: Although it is difficult to compare them with our method in terms of the constant prefactors in respective complexity bounds, we can do some heuristics to estimate these prefactors, which will demonstrate the low-cost advantage of our method.
Revised: Although it is difficult to compare them with our method in terms of the constant prefactors in respective complexity bounds, we can do some heuristics to estimate these prefactors, which will demonstrate the low-cost advantage of our method.

Original: For our Stiefel SGD (Algo. 1), we count the numbers of matrix multiplications needed per iteration, each of which costs nm 2 .
Revised: For our Stiefel SGD (Algo. 1), we count the numbers of matrix multiplications needed per iteration, each of which costs nm 2 .

Original: Namely, step 2: 2 multiplications; step 3: 1; step 4: 1; step 5: 6.
Revised: Namely, step 2: 2 multiplications; step 3: 1; step 4: 1; step 5: 6.

Original: In step 6 we need to call Algo.
Revised: In step 6 we need to call Algo.

Original: 3 for computing matrix root inversion, and a closer look gives 3 m 3 -cost matrix multiplications in each iteration of Algo.
Revised: 3 for computing matrix root inversion, and a closer look gives 3 m 3 -cost matrix multiplications in each iteration of Algo.

Original: Since Fig. 5(c) shows that 8 inner iteration is enough for this inner loop to converge under double precision, in total, 10 nm 2 -cost matrix multiplications and 24 m 3 -cost matrix multiplications are needed in each (outer) iteration of our optimizer (Algo. 1).
Revised: Since Fig. 5(c) shows that 8 inner iteration is enough for this inner loop to converge under double precision, in total, 10 nm 2 -cost matrix multiplications and 24 m 3 -cost matrix multiplications are needed in each (outer) iteration of our optimizer (Algo. 1).

Original: We also count the number of matrix multiplications needed for other aforementioned algorithms.
Revised: We also count the number of matrix multiplications needed for other aforementioned algorithms.

Original: For Momentumless
Revised: For Momentumless

Original: Stiefel (S)GD (Wen and Yin, 2013), the smartly designed Cayley map retraction takes abouttimes of nm 2 -cost matrix multiplications and one inversion for a (2 m ) ×
Revised: Stiefel (S)GD (Wen and Yin, 2013), the smartly designed Cayley map retraction takes abouttimes of nm 2 -cost matrix multiplications and one inversion for a (2 m ) ×

Original: (2 m ) -sized matrix in each iteration.
Revised: (2 m ) -sized matrix in each iteration.

Original: In comparison, Projected Stiefel SGD Li et al.
Revised: In comparison, Projected Stiefel SGD Li et al.

Original: (2020) needs 9 nm 2 -cost matrix multiplications for projecting momentum and 6 nm 2 -cost matrix multiplications each iteration of
Revised: (2020) needs 9 nm 2 -cost matrix multiplications for projecting momentum and 6 nm 2 -cost matrix multiplications each iteration of

Original: Cayley loop.
Revised: Cayley loop.

Original: Since 8 iterations are needed to compute the Cayley transform to single precision (see fig.
Revised: Since 8 iterations are needed to compute the Cayley transform to single precision (see fig.

Original: 5; note it is just single precision instead of double precision by our inner loop), a total ofnm 2 -cost matrix multiplications are needed per outer iteration.
Revised: 5; note it is just single precision instead of double precision by our inner loop), a total ofnm 2 -cost matrix multiplications are needed per outer iteration.

Original: The above estimation is just for the part on maintaining the structure of the position variable.
Revised: The above estimation is just for the part on maintaining the structure of the position variable.

Original: Let’s now also discuss the momentum variable.
Revised: Let’s now also discuss the momentum variable.

Original: Our carefully designed way of introducing momentum bypasses the costly moving momentum while keeping it in tangent spaces.
Revised: Our carefully designed way of introducing momentum bypasses the costly moving momentum while keeping it in tangent spaces.

Original: Note Projected Stiefel
Revised: Note Projected Stiefel

Original: 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.
Revised: 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.

Original: To add to this discussion, parallel transport is another existing way of moving momentum, but it cannot even be solved with O ( nm 2 ) complexity.
Revised: To add to this discussion, parallel transport is another existing way of moving momentum, but it cannot even be solved with O ( nm 2 ) complexity.

Original: Altogether, the above estimations heuristically show that our prefactor is close to that in
Revised: Altogether, the above estimations heuristically show that our prefactor is close to that in

Original: Momentumless Stiefel (S)GD but much smaller than those in existing methods that have momentum.
Revised: Momentumless Stiefel (S)GD but much smaller than those in existing methods that have momentum.

Original: This matches the experimental evidence in Fig.
Revised: This matches the experimental evidence in Fig.

Original: D T ERMINOLOGIES OF R IEMANNIAN OPTIMIZATION
Revised: D T ERMINOLOGIES OF R IEMANNIAN OPTIMIZATION

Original: In each iteration of optimization, gradient needs to be a vector in tangent space, and the operation that maps nonintrinsic gradient to the tangent space is called projection.
Revised: In each iteration of optimization, gradient needs to be a vector in tangent space, and the operation that maps nonintrinsic gradient to the tangent space is called projection.

Original: 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 a closed form computation, and approximation known as retraction is used.
Revised: 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 a closed form computation, and approximation known as retraction is used

Original: When momentum, a tangent vector, is involved, at the same time we update the position, tangent space is also changed.
Revised: When momentum, a tangent vector, is involved, at the same time we update the position, tangent space is also changed.

Original: An intrinsic way to move tangent vectors is by parallel transport, which is an isomorphism of the tangent spaces of two points connected by a geodesic.
Revised: An intrinsic way to move tangent vectors is by parallel transport, which is an isomorphism of the tangent spaces of two points connected by a geodesic.

Original: E A BOUT THE Y, V DECOMPOSITION E. T HE GEOMETRIC VERSION AND THE DYNAMICAL VERSION OF THEIR STRUCTURAL CONSTRAINTS
Revised: E A BOUT THE Y, V DECOMPOSITION E. T HE GEOMETRIC VERSION AND THE DYNAMICAL VERSION OF THEIR STRUCTURAL CONSTRAINTS

Original: A tangent vector Q ∈ T X St can be decomposed, as studied in Edelman et al.
Revised: A tangent vector Q ∈ T X St can be decomposed, as studied in Edelman et al.

Original: (1998); Absil et al.
Revised: (1998); Absil et al.

Original: (2009), into
Revised: (2009), into

Original: [Equation]
Revised: [Equation]

Original: 16) where X ∈ St ( n, m ) ,
Revised: 16) where X ∈ St ( n, m ) ,

Original: Y ∈ R m × m , Z ∈ R ( n − m ) × m ; X ⊥ ∈ R n × ( n − m ) is a matrix in the orthogonal complement of X , i.e., X ⊤ · X ⊥ =
Revised: Y ∈ R m × m , Z ∈ R ( n − m ) × m ; X ⊥ ∈ R n × ( n − m ) is a matrix in the orthogonal complement of X , i.e., X ⊤ · X ⊥ =

Original: 0 m × ( n − m ) .
Revised: 0 m × ( n − m ) .

Original: Let V =
Revised: Let V =

Original: X ⊥ Z .
Revised: X ⊥ Z .

Original: We have X ⊥ V and
Revised: We have X ⊥ V and

Original: [Equation]
Revised: [Equation]

Original: The above representation immediately implies Y is a skew-symmetric matrix since
Revised: The above representation immediately implies Y is a skew-symmetric matrix since

Original: [Equation]
Revised: [Equation]

Original: In our derivation, however, X and Q are variables (they really should be X ( t )
Revised: In our derivation, however, X and Q are variables (they really should be X ( t )

Original: ∈
Revised: ∈

Original: St , Q ( t ) ∈ T X ( t ) St ) that change with time, and this poses a new challenge.
Revised: St , Q ( t ) ∈ T X ( t ) St ) that change with time, and this poses a new challenge.

Original: More precisely, given the Q dynamics, i.e. ˙ Q , there could be infinitely many ways of assigning ˙ Y and ˙ V so that they together produce the given ˙
Revised: More precisely, given the Q dynamics, i.e. ˙ Q , there could be infinitely many ways of assigning ˙ Y and ˙ V so that they together produce the given ˙

Original: Q ; however, in general they do not maintain the tangent decomposition structures.
Revised: Q ; however, in general they do not maintain the tangent decomposition structures.

Original: More precisely, starting with
Revised: More precisely, starting with

Original: Y
Revised: Y

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: 0 and X (0)
Revised: 0 and X (0)

Original: ⊤ V (0) =(and hence X (0)
Revised: ⊤ V (0) =(and hence X (0)

Original: ⊤ Q (0) + Q (0)
Revised: ⊤ Q (0) + Q (0)

Original: ⊤ X (0) = 0 ), one may not have
Revised: ⊤ X (0) = 0 ), one may not have

Original: Y ( t ) ⊤ + Y ( t ) = 0 and X ( t )
Revised: Y ( t ) ⊤ + Y ( t ) = 0 and X ( t )

Original: ⊤ V ( t ) = 

Revised: ⊤ V ( t ) =

Original: for t >despite that X ( t )
Revised: for t >despite that X ( t )

Original: ⊤ Q ( t ) + Q ( t )
Revised: ⊤ Q ( t ) + Q ( t )

Original: ⊤ X ( t ) = 0 will still be guaranteed.
Revised: ⊤ X ( t ) = 0 will still be guaranteed.

Original: However, we have found a nontrivial choice of ˙ Y and ˙ V (Eq.8), so that the (static) geometric constraint becomes dynamically true, i.e. Y ( t ) ⊤
Revised: However, we have found a nontrivial choice of ˙ Y and ˙ V (Eq.8), so that the (static) geometric constraint becomes dynamically true, i.e. Y ( t ) ⊤

Original: + Y ( t ) = 

Revised: + Y ( t ) =

Original: and X ( t )
Revised: and X ( t )

Original: ⊤ V ( t ) = 0 for all t .
Revised: ⊤ V ( t ) = 0 for all t .

Original: Therefore we can simply maintain the decomposition
Revised: Therefore we can simply maintain the decomposition

Original: [Equation]
Revised: [Equation]

Original: E. T HE ADVANTAGE OF THE Y, V DECOMPOSITION
Revised: E. T HE ADVANTAGE OF THE Y, V DECOMPOSITION

Original: The Y, V decomposition is near ‘orthogonal’ in the sense that XY ⊥ V
Revised: The Y, V decomposition is near ‘orthogonal’ in the sense that XY ⊥ V

Original: and it makes the metric separable.
Revised: and it makes the metric separable.

Original: More precisely, consider our Y, V representation.
Revised: More precisely, consider our Y, V representation.

Original: ∀ ∆ i ∈ T X St, i = 1 , 2 , if we denote
Revised: ∀ ∆ i ∈ T X St, i = 1 , 2 , if we denote

Original: Y ∆ ,i = X ⊤ ∆ i and V ∆ , 2 =
Revised: Y ∆ ,i = X ⊤ ∆ i and V ∆ , 2 =

Original: ( I − XX ⊤ )∆ i , which means ∆ i = XY ∆ ,i + V ∆ ,i , then the metric (2) can be rewritten as
Revised: ( I − XX ⊤ )∆ i , which means ∆ i = XY ∆ ,i + V ∆ ,i , then the metric (2) can be rewritten as

Original: [Equation]
Revised: [Equation]

Original: This means the metric of the tangent space is a linear combination of the metrics of Y and V -directions.
Revised: This means the metric of the tangent space is a linear combination of the metrics of Y and V -directions.

Original: Therefore, this gives an intuition why our continuous dynamics and numerical discretization in the following can handle the family of canonical-type metrics.
Revised: Therefore, this gives an intuition why our continuous dynamics and numerical discretization in the following can handle the family of canonical-type metrics.

Original: F C HOICE OF γ
Revised: F C HOICE OF γ

Original: Although monotonicity of r suffices for the convergence, choices of γ = ˙ r/r affect the convergence speed just like in Euclidean cases (e.g.,Wibisono et al. (2016)
Revised: Although monotonicity of r suffices for the convergence, choices of γ = ˙ r/r affect the convergence speed just like in Euclidean cases (e.g.,Wibisono et al. (2016)

Original: ; Wilson et al.
Revised: ; Wilson et al.

Original: Popular choices include constant γ and γ ( t ) = 3 /t , but methods proposed here work for all γ ’s (e.g., γ = 3 /t + ct p for variance reduction Tao and Ohsawa (2020)).
Revised: Popular choices include constant γ and γ ( t ) = 3 /t , but methods proposed here work for all γ ’s (e.g., γ = 3 /t + ct p for variance reduction Tao and Ohsawa (2020)).

Original: However, for simplicity we will use constant γ from now on.
Revised: However, for simplicity we will use constant γ from now on.

Original: G P ROOF OF T HEOREM
Revised: G P ROOF OF T HEOREM

Original: Proof of Theorem 1, Part 1.
Revised: Proof of Theorem 1, Part 1.

Original: Given the Lagrangian
Revised: Given the Lagrangian

Original: [Equation]
Revised: [Equation]

Original: Legendre transform gives momentum P as
Revised: Legendre transform gives momentum P as

Original: [Equation]
Revised: [Equation]

Original: One can equivalently switch the Hamiltonian picture, and the corresponding Hamiltonian is H :
Revised: One can equivalently switch the Hamiltonian picture, and the corresponding Hamiltonian is H :

Original: T
Revised: T

Original: [Equation]
Revised: [Equation]

Original: St → R with H ( X, P
Revised: St → R with H ( X, P

Original: [Equation]
Revised: [Equation]

Original: :=
Revised: :=

Original: [Equation]
Revised: [Equation]

Original: (20) where b := aa − 1 solves ( I − aXX ⊤ )
Revised: (20) where b := aa − 1 solves ( I − aXX ⊤ )

Original: ( I − bXX ⊤ )
Revised: ( I − bXX ⊤ )

Original: =
Revised: =

Original: I .
Revised: I .

Original: Hence we get the following Hamilton’s equations
Revised: Hence we get the following Hamilton’s equations

Original: [Equation]
Revised: [Equation]

Original: Since X ⊤ X ≡ I , which gives ˙ X ⊤ X + X ⊤ ˙ X = 0 .
Revised: Since X ⊤ X ≡ I , which gives ˙ X ⊤ X + X ⊤ ˙ X = 0 .

Original: So we have X ⊤ P + P ⊤ X = 0 .
Revised: So we have X ⊤ P + P ⊤ X = 0 .

Original: Take derivative to get ˙
Revised: Take derivative to get ˙

Original: X ⊤ P + X ⊤ ˙
Revised: X ⊤ P + X ⊤ ˙

Original: P + ˙ P ⊤ X + P ⊤ ˙ X = 0 , and use the condition that Λ is symmetric, we can solve Λ is
Revised: P + ˙ P ⊤ X + P ⊤ ˙ X = 0 , and use the condition that Λ is symmetric, we can solve Λ is

Original: [Equation]
Revised: [Equation]

Original: And ( X, P ) system becomes
Revised: And ( X, P ) system becomes

Original: [Equation]
Revised: [Equation]

Original: By a coordinate change Q := 1 r ( I − bXX ⊤ ) P ∈ T X St and define friction parameter γ := ˙ r/r ,
Revised: By a coordinate change Q := 1 r ( I − bXX ⊤ ) P ∈ T X St and define friction parameter γ := ˙ r/r ,

Original: Eq.
Revised: Eq.

Original: (23) becomes
Revised: (23) becomes

Original: [Equation]
Revised: [Equation]

Original: In order to prove the second part of Thm. 1, we need the following lemma.
Revised: In order to prove the second part of Thm. 1, we need the following lemma.

Original: Lemma 1 (First-order stationary point) .
Revised: Lemma 1 (First-order stationary point) .

Original: If X ∈ St s.t.
Revised: If X ∈ St s.t.

Original: G − 1+ b 2 XX ⊤ G − 1 − b 2 XG ⊤ X = 0 , then ∀ ∆ ∈ T X St , we have Tr ( G ⊤ ∆)
Revised: G − 1+ b 2 XX ⊤ G − 1 − b 2 XG ⊤ X = 0 , then ∀ ∆ ∈ T X St , we have Tr ( G ⊤ ∆)

Original: = 0 , which means X is a first-order stationary point of f .
Revised: = 0 , which means X is a first-order stationary point of f .

Original: Proof.
Revised: Proof.

Original: Left multiply both side of G − 1+ b 2 XX ⊤ G − 1 − b 2 XG ⊤ X = 0 by XX ⊤ , we have XX ⊤ G =
Revised: Left multiply both side of G − 1+ b 2 XX ⊤ G − 1 − b 2 XG ⊤ X = 0 by XX ⊤ , we have XX ⊤ G =

Original: XG ⊤ X , further we also have G =
Revised: XG ⊤ X , further we also have G =

Original: XG ⊤ X .
Revised: XG ⊤ X .

Original: So we have
Revised: So we have

Original: [Equation]
Revised: [Equation]

Original: Thus we have Tr (∆ ⊤ G ) = 0 , ∀ ∆ ∈ T X St , which means X is a 1-order stationary point of f .
Revised: Thus we have Tr (∆ ⊤ G ) = 0 , ∀ ∆ ∈ T X St , which means X is a 1-order stationary point of f .

Original: Proof of Theorem 1, Part 2.
Revised: Proof of Theorem 1, Part 2.

Original: Let t → ( X ( t ) , Q ( t )) be a solution of Eq.
Revised: Let t → ( X ( t ) , Q ( t )) be a solution of Eq.

Original: Define the ‘energy’ function E : T St → R as
Revised: Define the ‘energy’ function E : T St → R as

Original: [Equation]
Revised: [Equation]

Original: This gives a Lyapunov function.
Revised: This gives a Lyapunov function.

Original: So we have a neighbourhood U of ( X ∗ , 0) such that E ( X, Q ) ≥ f
Revised: So we have a neighbourhood U of ( X ∗ , 0) such that E ( X, Q ) ≥ f

Original: ( X ) ≥ f
Revised: ( X ) ≥ f

Original: ( X ∗ ) for any ( X, Q ) ∈ U .
Revised: ( X ∗ ) for any ( X, Q ) ∈ U .

Original: More over, since X ⊤ Q + Q ⊤ X ≡ 0 , we have dE
Revised: More over, since X ⊤ Q + Q ⊤ X ≡ 0 , we have dE

Original: [Equation]
Revised: [Equation]

Original: Using the fact that X ⊤ Q + Q ⊤ X ≡ 0 , we have Tr ( Q ⊤ XQ ⊤ Q ) = 0 and Tr ( Q ⊤ XX ⊤ G )
Revised: Using the fact that X ⊤ Q + Q ⊤ X ≡ 0 , we have Tr ( Q ⊤ XQ ⊤ Q ) = 0 and Tr ( Q ⊤ XX ⊤ G )

Original: +
Revised: +

Original: Tr ( Q ⊤ XG ⊤ X ) = 0 , which gives dE
Revised: Tr ( Q ⊤ XG ⊤ X ) = 0 , which gives dE

Original: [Equation]
Revised: [Equation]

Original: Since we have r monotonely increasing, which means γ = r ′ /r > 0 , ∀ t .
Revised: Since we have r monotonely increasing, which means γ = r ′ /r > 0 , ∀ t .

Original: Then we have the energy is decreasing monotonically, implying that Q = 0 when converged.
Revised: Then we have the energy is decreasing monotonically, implying that Q = 0 when converged.

Original: By lemma 1 that the limiting point for X ( t ) is a first order stationary point, which is X ∗ since it is an isolated local minimum.
Revised: By lemma 1 that the limiting point for X ( t ) is a first order stationary point, which is X ∗ since it is an isolated local minimum.

Original: Remark 2.
Revised: Remark 2.

Original: For the sake of length, rate 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.
Revised: For the sake of length, rate 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.

Original: Duruisseaux and Leok (2021).
Revised: Duruisseaux and Leok (2021).

Original: H P ROOF OF T HEOREM
Revised: H P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: We can derive a new system of ODEs as
Revised: We can derive a new system of ODEs as

Original: [Equation]
Revised: [Equation]

Original: with viewing X ⊤ G as a matrix function of t following Eq.
Revised: with viewing X ⊤ G as a matrix function of t following Eq.

Original: By the uniqueness and existence of
Revised: By the uniqueness and existence of

Original: ODE, we have that this system of ODE with variable ( X ⊤ X, X ⊤ Q ) and initial condition ( I, 0) has the unique solution X ⊤ X ≡ I, X ⊤ Q ≡ 0 .
Revised: ODE, we have that this system of ODE with variable ( X ⊤ X, X ⊤ Q ) and initial condition ( I, 0) has the unique solution X ⊤ X ≡ I, X ⊤ Q ≡ 0 .

Original: I XYV SYSTEM IS STRUCTURE PRESERVING
Revised: I XYV SYSTEM IS STRUCTURE PRESERVING

Original: Theorem 8 (Constrained optimization with unconstrained dynamics) .
Revised: Theorem 8 (Constrained optimization with unconstrained dynamics) .

Original: As long as the initial condition of (8) satisfies
Revised: As long as the initial condition of (8) satisfies

Original: [Equation]
Revised: [Equation]

Original: then the dynamics automatically satisfies the same constraint, i.e.,
Revised: then the dynamics automatically satisfies the same constraint, i.e.,

Original: [Equation]
Revised: [Equation]

Original: Proof.
Revised: Proof.

Original: By the uniqueness of solution of ODE, we know that the following ODE for ( X ⊤ X, Y ⊤ +
Revised: By the uniqueness of solution of ODE, we know that the following ODE for ( X ⊤ X, Y ⊤ +

Original: Y, X ⊤ V ) , derived from Eq.
Revised: Y, X ⊤ V ) , derived from Eq.

Original: (8), has a unique solution if view V ⊤ V as an independent variable.
Revised: (8), has a unique solution if view V ⊤ V as an independent variable.

Original: [Equation]
Revised: [Equation]

Original: We can see that given initial condition ( X ⊤ 0 X 0 , Y ⊤ 

Revised: We can see that given initial condition ( X ⊤ 0 X 0 , Y ⊤

Original: + Y
Revised: + Y

Original: [Equation]
Revised: [Equation]

Original: , X ⊤
Revised: , X ⊤

Original: [Equation]
Revised: [Equation]

Original: ) =
Revised: ) =

Original: ( I, 0 , 0) , the unique solution is ( X ⊤ X, Y ⊤ + Y, X ⊤ V ) ≡ ( I, 0 , 0) .
Revised: ( I, 0 , 0) , the unique solution is ( X ⊤ X, Y ⊤ + Y, X ⊤ V ) ≡ ( I, 0 , 0) .

Original: So we know that the constraint in Eq.
Revised: So we know that the constraint in Eq.

Original: (28) are preserved by continuous dynamics Eq.
Revised: (28) are preserved by continuous dynamics Eq.

Original: J P ROOF OF T HEOREM
Revised: J P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: We assume the initial condition ( X 0 , Y 0 , V 0 )
Revised: We assume the initial condition ( X 0 , Y 0 , V 0 )

Original: satisfies constraint (28).
Revised: satisfies constraint (28).

Original: For Eq.
Revised: For Eq.

Original: (9), we check X ⊤ ( t ) X ( t ) ≡ 0 , Y ( t )
Revised: (9), we check X ⊤ ( t ) X ( t ) ≡ 0 , Y ( t )

Original: + Y ⊤ ( t ) ≡ 0 and X ⊤ ( t ) V ( t ) ≡ 0 .
Revised: + Y ⊤ ( t ) ≡ 0 and X ⊤ ( t ) V ( t ) ≡ 0 .

Original: d dt
Revised: d dt

Original: [Equation]
Revised: [Equation]

Original: Y

Original: [Equation]

Original: Y
Revised: Y

Original: [Equation]
Revised: [Equation]

Original: with initial condition Y 0 + Y ⊤ 

Revised: with initial condition Y 0 + Y ⊤

Original: = 0 , we have Y + Y ⊤ ≡
Revised: = 0 , we have Y + Y ⊤ ≡

Original: [Equation]
Revised: [Equation]

Original: with initial condition X ⊤ 0 V 

Revised: with initial condition X ⊤ 0 V

Original: =
Revised: =

Original: 0 , we have X ⊤ V ≡
Revised: 0 , we have X ⊤ V ≡

Original: [Equation]
Revised: [Equation]

Original: Using the conclusion Y + Y ⊤ ≡ 0 that we have just proved and initial condition X ⊤ 0 X 

Revised: Using the conclusion Y + Y ⊤ ≡ 0 that we have just proved and initial condition X ⊤ 0 X

Original: =
Revised: =

Original: I , we have X ⊤ X ≡ 0 .
Revised: I , we have X ⊤ X ≡ 0 .

Original: For Eq.
Revised: For Eq.

Original: (10), we have the exact solution of ϕ 2 is given by X ( t )
Revised: (10), we have the exact solution of ϕ 2 is given by X ( t )

Original: =
Revised: =

Original: X (0) , Y ( t )
Revised: X (0) , Y ( t )

Original: =
Revised: =

Original: Y (0) , and with
Revised: Y (0) , and with

Original: M := γI m × m − 3 a − 2 2 Y ( t ) ,
Revised: M := γI m × m − 3 a − 2 2 Y ( t ) ,

Original: [Equation]
Revised: [Equation]

Original: For Eq.
Revised: For Eq.

Original: (11), we need to check ddt (cid:0) X ⊤ ( t ) X ( t ) (cid:1) = 0 and ddt ( X ⊤ ( t ) V ( t ))
Revised: (11), we need to check ddt (cid:0) X ⊤ ( t ) X ( t ) (cid:1) = 0 and ddt ( X ⊤ ( t ) V ( t ))

Original: = 0 .
Revised: = 0 .

Original: By the uniqueness of solution of ODE, we know that the following ODE for ( X ⊤ X, X ⊤ V, V ⊤ V ) , derived from Eq.
Revised: By the uniqueness of solution of ODE, we know that the following ODE for ( X ⊤ X, X ⊤ V, V ⊤ V ) , derived from Eq.

Original: (11), has a unique solution
Revised: (11), has a unique solution

Original: [Equation]
Revised: [Equation]

Original: We can see that given initial condition ( X ⊤ 0 X 0 , X ⊤ 0 V 0 , V ⊤ 0 V 0 )
Revised: We can see that given initial condition ( X ⊤ 0 X 0 , X ⊤ 0 V 0 , V ⊤ 0 V 0 )

Original: =
Revised: =

Original: ( I, 0 , V ⊤ 0 V 0 ) , the unique solution is ( X ⊤ X, X ⊤ V, V ⊤ V ) ≡ ( I, 0 , V ⊤ 0 V 0 ) .
Revised: ( I, 0 , V ⊤ 0 V 0 ) , the unique solution is ( X ⊤ X, X ⊤ V, V ⊤ V ) ≡ ( I, 0 , V ⊤ 0 V 0 ) .

Original: So we know that the constraint X ⊤ X
Revised: So we know that the constraint X ⊤ X

Original: =
Revised: =

Original: I , X ⊤ V = 0 are preserved by continuous dynamics Eq.
Revised: I , X ⊤ V = 0 are preserved by continuous dynamics Eq.

Original: K P ROOF OF T HEOREM
Revised: K P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: Eq.
Revised: Eq.

Original: (12) is structure preserving.
Revised: (12) is structure preserving.

Original: We use the idea from Tao and Ohsawa (2020).
Revised: We use the idea from Tao and Ohsawa (2020).

Original: Assume the initial value X 0 , Y 0 and V 0 satisfies the constraint, i.e., X ⊤ 0 X 

Revised: Assume the initial value X 0 , Y 0 and V 0 satisfies the constraint, i.e., X ⊤ 0 X

Original: =
Revised: =

Original: I , Y 0 + Y ⊤ 

Revised: I , Y 0 + Y ⊤

Original: = 0 and X ⊤ 0 V 0 = 0 .
Revised: = 0 and X ⊤ 0 V 0 = 0 .

Original: For the first step updating Y , due to the special form of the derivative of Y that is always skewsymmetric, we can tell that Y h is also skew-symmetric and Y h
Revised: For the first step updating Y , due to the special form of the derivative of Y that is always skewsymmetric, we can tell that Y h is also skew-symmetric and Y h

Original: + Y ⊤ h
Revised: + Y ⊤ h

Original: = 0 .
Revised: = 0 .

Original: And the skew-symmetricity of Y h gives us that
Revised: And the skew-symmetricity of Y h gives us that

Original: [Equation]
Revised: [Equation]

Original: so all 3 conditions in Eq.
Revised: so all 3 conditions in Eq.

Original: (28) are satisfied, meaning Eq.
Revised: (28) are satisfied, meaning Eq.

Original: (12) is structure preserving.
Revised: (12) is structure preserving.

Original: Eq.
Revised: Eq.

Original: (14) is structure preserving.
Revised: (14) is structure preserving.

Original: Given initial condition X ⊤ 0 X 

Revised: Given initial condition X ⊤ 0 X

Original: =
Revised: =

Original: I and X ⊤ 0 V 0 = 0 , we check constraint X ⊤ h V h = 0 .
Revised: I and X ⊤ 0 V 0 = 0 , we check constraint X ⊤ h V h = 0 .

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (13) is structure preserving.
Revised: (13) is structure preserving.

Original: First, we show that the discretization is a one order approximation of the exact solution.
Revised: First, we show that the discretization is a one order approximation of the exact solution.

Original: We get V h and X † by forward Euler.
Revised: We get V h and X † by forward Euler.

Original: And since X 0 and V 0 satisfies the constraint 28, we have
Revised: And since X 0 and V 0 satisfies the constraint 28, we have

Original: [Equation]
Revised: [Equation]

Original: which means X h
Revised: which means X h

Original: = X † + O ( h 2 ) =
Revised: = X † + O ( h 2 ) =

Original: X 0 + hV 

Revised: X 0 + hV

Original: + O ( h 2 ) , indeed a one order approximation of exact solution.
Revised: + O ( h 2 ) , indeed a one order approximation of exact solution.

Original: Next we show the numerical discretization is structure preserving.
Revised: Next we show the numerical discretization is structure preserving.

Original: For constraint X ⊤ V = 0 , we can check that
Revised: For constraint X ⊤ V = 0 , we can check that

Original: [Equation]
Revised: [Equation]

Original: which gives us X ⊤ h V h = 0 .
Revised: which gives us X ⊤ h V h = 0 .

Original: For restriction X ⊤ X
Revised: For restriction X ⊤ X

Original: =
Revised: =

Original: I , we have
Revised: I , we have

Original: [Equation]
Revised: [Equation]

Original: L P ROOF OF T HEOREM
Revised: L P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: The composition of structure-preserving maps is structure preserving.
Revised: The composition of structure-preserving maps is structure preserving.

Original: ϕ 1 ◦ ϕ 2 ◦ ϕ 3 is a 1st-order integrator due to operator splitting theory McLachlan and Quispel (2002), and its convergence order is kept after any ϕ
Revised: ϕ 1 ◦ ϕ 2 ◦ ϕ 3 is a 1st-order integrator due to operator splitting theory McLachlan and Quispel (2002), and its convergence order is kept after any ϕ

Original: j gets replaced by ¯ ϕ j as long as the difference between ¯ ϕ j and ϕ j is higher-order
Revised: j gets replaced by ¯ ϕ j as long as the difference between ¯ ϕ j and ϕ j is higher-order

Original: Tao (2016).
Revised: Tao (2016).

Original: M D ISCUSSIONS ON OUR NUMERICAL INTEGRATOR M. B ENEFITS OF THE STEP X † ( X ⊤† X † ) −
Revised: M D ISCUSSIONS ON OUR NUMERICAL INTEGRATOR M. B ENEFITS OF THE STEP X † ( X ⊤† X † ) −

Original: To remain the position X on the manifold, certain techniques that pull point back to manifold are used.
Revised: To remain the position X on the manifold, certain techniques that pull point back to manifold are used.

Original: For example, in Li et al.
Revised: For example, in Li et al.

Original: Lin et al.
Revised: Lin et al.

Original: (2020) for an n -by- m matrix X
Revised: (2020) for an n -by- m matrix X

Original: that is not on the manifold, they perform QR decomposition X = QR and update the new X to be the first m columns of Q .
Revised: that is not on the manifold, they perform QR decomposition X = QR and update the new X to be the first m columns of Q .

Original: However, the QR decomposition is not unique and such step does not have a closed form expression.
Revised: However, the QR decomposition is not unique and such step does not have a closed form expression.

Original: More importantly, it cannot help design a structure-preserving scheme.
Revised: More importantly, it cannot help design a structure-preserving scheme.

Original: Instead, our algorithm has a similarly functioned step via the square root of the inverse matrix
Revised: Instead, our algorithm has a similarly functioned step via the square root of the inverse matrix

Original: X h
Revised: X h

Original: = X † ( X ⊤† X † ) − 1 / 2 .
Revised: = X † ( X ⊤† X † ) − 1 / 2 .

Original: In this case, it follows that X ⊤ h X h
Revised: In this case, it follows that X ⊤ h X h

Original: = I .
Revised: = I .

Original: This step is carefully designed from discretization of ODE ensuring the structure of X , Y , V is preserved at the same time.
Revised: This step is carefully designed from discretization of ODE ensuring the structure of X , Y , V is preserved at the same time.

Original: Through this update, we are able to obtain a structure-preserving method with a closed form expression (see the proof of Theorem 4 in Apdx.
Revised: Through this update, we are able to obtain a structure-preserving method with a closed form expression (see the proof of Theorem 4 in Apdx.

Original: K).
Revised: K).

Original: Note that the square root of the inverse matrix can be iteratively solved with low computational cost (see Algo. 3; Higham (1997)).
Revised: Note that the square root of the inverse matrix can be iteratively solved with low computational cost (see Algo. 3; Higham (1997)).

Original: It was proved to be quadratically convergent, which means only a couple of iterations are needed to reach machine precision, and only matrix multiplications are needed.
Revised: It was proved to be quadratically convergent, which means only a couple of iterations are needed to reach machine precision, and only matrix multiplications are needed.

Original: This lead to our inner loop is more efficient in both cost per iteration and convergence speed comparing to Li et al.
Revised: This lead to our inner loop is more efficient in both cost per iteration and convergence speed comparing to Li et al.

Original: See Sec.
Revised: See Sec.

Original: Q for more details.
Revised: Q for more details.

Original: Input: Symmetric m -by- m matrix A , tol
Revised: Input: Symmetric m -by- m matrix A , tol

Original: Initialization : Y 

Revised: Initialization : Y

Original: =
Revised: =

Original: A , Z 

Revised: A , Z

Original: =
Revised: =

Original: I m × m , k =

Revised: I m × m , k =

Original: while ∥ Y 2 k − A ∥ ≥ tol do
Revised: while ∥ Y 2 k − A ∥ ≥ tol do

Original: Y k +1 =
Revised: Y k +1 =

Original: 12 Y k (3 I − Z k Y k )
Revised: 12 Y k (3 I − Z k Y k )

Original: Z k +1 =
Revised: Z k +1 =

Original: (3 I − Z k Y k )
Revised: (3 I − Z k Y k )

Original: Z k k ← k +5 end 6 return
Revised: Z k k ← k +5 end 6 return

Original: Y k
Revised: Y k

Original: ≈ A 12 and Z k
Revised: ≈ A 12 and Z k

Original: ≈ A −
Revised: ≈ A −

Original: Another benefit of this step is that it is more stable to the truncation error produced by finite machine precision (which can become a significant issue if the model is trained in single accuracy on consumer graphic cards or even quantized).
Revised: Another benefit of this step is that it is more stable to the truncation error produced by finite machine precision (which can become a significant issue if the model is trained in single accuracy on consumer graphic cards or even quantized).

Original: Although X ⊤ † X † follows the expression
Revised: Although X ⊤ † X † follows the expression

Original: [Equation]
Revised: [Equation]

Original: the stability of X † ( X ⊤† X † ) − 1 / 2 is much better than X † ( I + h 2 V ⊤ 0 V 0 ) − 1 / 2 .
Revised: the stability of X † ( X ⊤† X † ) − 1 / 2 is much better than X † ( I + h 2 V ⊤ 0 V 0 ) − 1 / 2 .

Original: Particularly, when n = m , it is an identical map if there is no machine error.
Revised: Particularly, when n = m , it is an identical map if there is no machine error.

Original: M. S IMPLIFICATION OF MATRIX EXPONENTIAL
Revised: M. S IMPLIFICATION OF MATRIX EXPONENTIAL

Original: There are two steps in the discretizations (30) and (12) that use the matrix exponential of m × m matrices.
Revised: There are two steps in the discretizations (30) and (12) that use the matrix exponential of m × m matrices.

Original: Similar matrix exponential is also shown in Li et al.
Revised: Similar matrix exponential is also shown in Li et al.

Original: ; Wen and Yin (2013) but with matrices of much larger size n × n (note n > m ).
Revised: ; Wen and Yin (2013) but with matrices of much larger size n × n (note n > m ).

Original: We introduce two ways of simplifying the computation of the matrix exponential.
Revised: We introduce two ways of simplifying the computation of the matrix exponential.

Original: First, note that Cayley transform is a 2nd-order structurepreserving approximation of the matrix exponential, defined as follows
Revised: First, note that Cayley transform is a 2nd-order structurepreserving approximation of the matrix exponential, defined as follows

Original: [Equation]
Revised: [Equation]

Original: By applying the Cayley transform, the computation is reduced to an inversion of matrix and matrix multiplication while it still keeps the variable on the manifold.
Revised: By applying the Cayley transform, the computation is reduced to an inversion of matrix and matrix multiplication while it still keeps the variable on the manifold.

Original: Additionally, we can use the first-order forward Euler to discretize the ODEs for the two steps, which is also the first-order truncation of the matrix exponential.
Revised: Additionally, we can use the first-order forward Euler to discretize the ODEs for the two steps, which is also the first-order truncation of the matrix exponential.

Original: In this case, the X update in scheme (12) is just X h
Revised: In this case, the X update in scheme (12) is just X h

Original: = X 0 + hY h and is no longer structure-preserving but structure-preserving property of the overall algorithm is not affected (see Thm. 6).
Revised: = X 0 + hY h and is no longer structure-preserving but structure-preserving property of the overall algorithm is not affected (see Thm. 6).

Original: For scheme (30), its Cayley map approximation is structure preserving and defined as follows
Revised: For scheme (30), its Cayley map approximation is structure preserving and defined as follows

Original: [Equation]
Revised: [Equation]

Original: Note the variable V will still satisfy the constraint (28) even if we change the update of V to forward
Revised: Note the variable V will still satisfy the constraint (28) even if we change the update of V to forward

Original: Euler
Revised: Euler

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: In fact, the above two ways and the original matrix exponential share the same complexity per iteration.
Revised: In fact, the above two ways and the original matrix exponential share the same complexity per iteration.

Original: Also, we do not find any significant difference in the convergence in numerical experiments (Apdx. Q).
Revised: Also, we do not find any significant difference in the convergence in numerical experiments (Apdx. Q).

Original: Hence we will always use forward Euler as default, which has the lowest computational cost.
Revised: Hence we will always use forward Euler as default, which has the lowest computational cost.

Original: M. P ROOF OF T HEOREM
Revised: M. P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: The composition of ¯ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 in any order will give a structure preserving scheme, because in Sec.
Revised: The composition of ¯ ϕ 1 , ¯ ϕ 2 , ¯ ϕ 3 in any order will give a structure preserving scheme, because in Sec.

Original: K we proved that each of them is structure preserving.
Revised: K we proved that each of them is structure preserving.

Original: When we apply it in specific order ¯ ϕ 3 ◦ ¯ ϕ 1 ◦ ¯ ϕ 2 , we prove when X ⊤ X
Revised: When we apply it in specific order ¯ ϕ 3 ◦ ¯ ϕ 1 ◦ ¯ ϕ 2 , we prove when X ⊤ X

Original: = I is no longer preserved by ¯ ϕ 2 , the assembled scheme is still structure preserving.
Revised: = I is no longer preserved by ¯ ϕ 2 , the assembled scheme is still structure preserving.

Original: We only need to prove that [ X h , Y h , V h ]
Revised: We only need to prove that [ X h , Y h , V h ]

Original: = ¯ ϕ 3 ( x 0 , Y 0 , V 0 ) satisfies the following: when initial condition satisfies Y ⊤ 

Revised: = ¯ ϕ 3 ( x 0 , Y 0 , V 0 ) satisfies the following: when initial condition satisfies Y ⊤

Original: + Y 

Revised: + Y

Original: = 0 and X ⊤ 0 V 0 = 0 , we have X ⊤ h X h
Revised: = 0 and X ⊤ 0 V 0 = 0 , we have X ⊤ h X h

Original: =
Revised: =

Original: I , Y ⊤ h + Y h = 0 and X ⊤ h V h
Revised: I , Y ⊤ h + Y h = 0 and X ⊤ h V h

Original: = 0 .
Revised: = 0 .

Original: Since we have the step X h = X † ( X ⊤† X † ) − 1 / 2 , we have X ⊤ h X h
Revised: Since we have the step X h = X † ( X ⊤† X † ) − 1 / 2 , we have X ⊤ h X h

Original: = I .
Revised: = I .

Original: Y is not updating so
Revised: Y is not updating so

Original: Y h + Y ⊤ h
Revised: Y h + Y ⊤ h

Original: = 0 .
Revised: = 0 .

Original: [Equation]
Revised: [Equation]

Original: As a result, X
Revised: As a result, X

Original: [Equation]
Revised: [Equation]

Original: M.
Revised: M.

Original: T HE
Revised: T HE

Original: [Table] [Table]
Revised: [Table] [Table]

Original: A
Revised: A

Original: ˜ ϕ
Revised: ˜ ϕ

Original: [Equation]
Revised: [Equation]

Original: 
Revised: 

Original: [Equation]
Revised: [Equation]

Original: ‘gradient’ and 2-order momentum :
Revised: ‘gradient’ and 2-order momentum :

Original: [Equation]
Revised: [Equation]

Original: ˆ
Revised: ˆ

Original: [Equation]
Revised: [Equation]

Original: 
Revised: 

Original: [Equation]
Revised: [Equation]

Original: N P ROOF OF T HEOREM
Revised: N P ROOF OF T HEOREM

Original: Proof.
Revised: Proof.

Original: We will use the same notation as in algorithm 2.
Revised: We will use the same notation as in algorithm 2.

Original: Assume X ⊤
Revised: Assume X ⊤

Original: i
Revised: i

Original: X i
Revised: X i

Original: =
Revised: =

Original: I , Y ⊤ i
Revised: I , Y ⊤ i

Original: + Y i = 0 ,
Revised: + Y i = 0 ,

Original: X ⊤ i V i = 0 , p i = p ⊤ i .
Revised: X ⊤ i V i = 0 , p i = p ⊤ i .

Original: Our initialization satisfies these conditions.
Revised: Our initialization satisfies these conditions.

Original: We prove that these conditions are satisfied after mapped by ˆ ϕ 3 ◦ ˆ ϕ 1 ◦ ˆ ϕ 2 in the following.
Revised: We prove that these conditions are satisfied after mapped by ˆ ϕ 3 ◦ ˆ ϕ 1 ◦ ˆ ϕ 2 in the following.

Original: Step 1 : Due to the skew-symmetricity of f i , we have p
Revised: Step 1 : Due to the skew-symmetricity of f i , we have p

Original: i +1 is symmetric.
Revised: i +1 is symmetric.

Original: Step 2 : X i , Z i , U i + 12 :=
Revised: Step 2 : X i , Z i , U i + 12 :=

Original: ˆ
Revised: ˆ

Original: ϕ 2 ( X i , Z i , U i ) , where U i + 12 = β 1 U
Revised: ϕ 2 ( X i , Z i , U i ) , where U i + 12 = β 1 U

Original: i − 3 a − 2 2 ηU i Z i − (1 − β 1 ) g i .
Revised: i − 3 a − 2 2 ηU i Z i − (1 − β 1 ) g i .

Original: Since
Revised: Since

Original: X ⊤ i
Revised: X ⊤ i

Original: U
Revised: U

Original: i
Revised: i

Original: = 

Revised: =

Original: and X ⊤ i g i
Revised: and X ⊤ i g i

Original: =
Revised: =

Original: 0 , we can check
Revised: 0 , we can check

Original: [Equation]
Revised: [Equation]

Original: Step
Revised: Step

Original: 3 : X i + 12 , Z i +1 , U i + 12 := ˆ ϕ 1 ( X i , Z i , U i + 12 )
Revised: 3 : X i + 12 , Z i +1 , U i + 12 := ˆ ϕ 1 ( X i , Z i , U i + 12 )

Original: where
Revised: where

Original: Z i
Revised: Z i

Original: +
Revised: +

Original: = β
Revised: = β

Original: Z i
Revised: Z i

Original: − (
Revised: − (

Original: − β 1 ) f
Revised: − β 1 ) f

Original: i
Revised: i

Original: and
Revised: and

Original: X i
Revised: X i

Original: + 

Revised: +

Original: =
Revised: =

Original: X i
Revised: X i

Original: +
Revised: +

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: i (cid:16) Z i +1 ⊘ ( p ◦ 12 i +1 + ϵ ) (cid:17) .
Revised: i (cid:16) Z i +1 ⊘ ( p ◦ 12 i +1 + ϵ ) (cid:17) .

Original: Since we have X ⊤
Revised: Since we have X ⊤

Original: i U i + 12 = 

Revised: i U i + 12 =

Original: We can have directly X ⊤ i + 1 2 U i + 12 = 0 .
Revised: We can have directly X ⊤ i + 1 2 U i + 12 = 0 .

Original: And due to f i is symmetric,
Revised: And due to f i is symmetric,

Original: Z i +1 is still skew-symmetric.
Revised: Z i +1 is still skew-symmetric.

Original: Note that X ⊤ i
Revised: Note that X ⊤ i

Original: +
Revised: +

Original: 12 X
Revised: 12 X

Original: i + 

Revised: i +

Original: = I no longer stands now but it satisfied until the algorithm finishes.
Revised: = I no longer stands now but it satisfied until the algorithm finishes.

Original: Step 4 X
Revised: Step 4 X

Original: i +
Revised: i +

Original: , Z
Revised: , Z

Original: i +
Revised: i +

Original: , U i +1 :=
Revised: , U i +1 :=

Original: ˆ ϕ
Revised: ˆ ϕ

Original: X i + 12 , Z i +1 , U i + 12 ) .
Revised: X i + 12 , Z i +1 , U i + 12 ) .

Original: First we can see from X i +

Revised: First we can see from X i +

Original: = X † ( X ⊤† X † ) −that X ⊤
Revised: = X † ( X ⊤† X † ) −that X ⊤

Original: i
Revised: i

Original: +

Revised: +

Original: X i +

Revised: X i +

Original: =
Revised: =

Original: I as long as X † =
Revised: I as long as X † =

Original: X i + 12 + η ˜
Revised: X i + 12 + η ˜

Original: UX ⊤ i +
Revised: UX ⊤ i +

Original: 12 X
Revised: 12 X

Original: i + 12 is full rank, which is always true since
Revised: i + 12 is full rank, which is always true since

Original: X i is full rank and η is small.
Revised: X i is full rank and η is small.

Original: Since we have X ⊤ i + 12 ˜ U = 0 by the special form of matrix I − X i +
Revised: Since we have X ⊤ i + 12 ˜ U = 0 by the special form of matrix I − X i +

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: i +
Revised: i +

Original: Since X ⊤† X † = I ,
Revised: Since X ⊤† X † = I ,

Original: [Equation]
Revised: [Equation]

Original: So we proved all the constrain are satisfied again, which means X ⊤ i +

Revised: So we proved all the constrain are satisfied again, which means X ⊤ i +

Original: X i +

Revised: X i +

Original: =
Revised: =

Original: I , Y ⊤ i +

Revised: I , Y ⊤ i +

Original: + Y
Revised: + Y

Original: i +

Revised: i +

Original: = 0 ,
Revised: = 0 ,

Original: X ⊤ i +1 V
Revised: X ⊤ i +1 V

Original: i +

Revised: i +

Original: = 0 ,
Revised: = 0 ,

Original: p i +1 =
Revised: p i +1 =

Original: p ⊤ i +
Revised: p ⊤ i +

Original: O O PTIMIZATION ON SO ( n ) AS A SPECIAL CASE FOR n = m ON St ( n, m )
Revised: O O PTIMIZATION ON SO ( n ) AS A SPECIAL CASE FOR n = m ON St ( n, m )

Original: Our method for Stiefel optimization can naturally be applied on the special orthogonal group SO ( n ) (Apdx. B.3) and is defined as follows
Revised: Our method for Stiefel optimization can naturally be applied on the special orthogonal group SO ( n ) (Apdx. B.3) and is defined as follows

Original: [Equation]
Revised: [Equation]

Original: Tao and Ohsawa (2020) proposed an efficient algorithm based on Lie group structure for the optimization on SO ( n ) although it cannot be generalized to Stiefel case.
Revised: Tao and Ohsawa (2020) proposed an efficient algorithm based on Lie group structure for the optimization on SO ( n ) although it cannot be generalized to Stiefel case.

Original: Our method restores the same integrator in Tao and Ohsawa (2020) while using different approach applied to a family of metrics.
Revised: Our method restores the same integrator in Tao and Ohsawa (2020) while using different approach applied to a family of metrics.

Original: In greater detail, for SO ( n ) , the canonical-type metric (2) degenerates to Euclidean metric up to constant scaling, i.e., g ( [Equation]
Revised: In greater detail, for SO ( n ) , the canonical-type metric (2) degenerates to Euclidean metric up to constant scaling, i.e., g (∆ 1 , ∆

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: Tr (∆ ⊤ 1 ∆ 2 ) .
Revised: Tr (∆ ⊤ 1 ∆ 2 ) .

Original: Then the corresponding Lagrangian is
Revised: Then the corresponding Lagrangian is

Original: [Equation]
Revised: [Equation]

Original: 36) which results in the following position-momentum ( X, Q ) dynamics
Revised: 36) which results in the following position-momentum ( X, Q ) dynamics

Original: [Equation]
Revised: [Equation]

Original: (37) with the same initialization as (7).
Revised: (37) with the same initialization as (7).

Original: Next, apply the same Y, V decomposition of Q .
Revised: Next, apply the same Y, V decomposition of Q .

Original: Since I − XX ⊤ = 0 , we have V = 0 , i.e., the momentum Q is purely in Y -direction.
Revised: Since I − XX ⊤ = 0 , we have V = 0 , i.e., the momentum Q is purely in Y -direction.

Original: Hence the equivalent dynamics to (37) is the following
Revised: Hence the equivalent dynamics to (37) is the following

Original: [Equation]
Revised: [Equation]

Original: Note when we take b = − 1 (namely, a = 1 / 2 ), (38) is identical to the continuous dynamics in Tao and Ohsawa (2020), although Tao and Ohsawa (2020) uses an intrinsic form which is coordinate-free but not explicit enough.
Revised: Note when we take b = − 1 (namely, a = 1 / 2 ), (38) is identical to the continuous dynamics in Tao and Ohsawa (2020), although Tao and Ohsawa (2020) uses an intrinsic form which is coordinate-free but not explicit enough.

Original: The numerical integrator corresponding to (S)GD is defined the same as Section M.4 and is summarized in Alg.4.
Revised: The numerical integrator corresponding to (S)GD is defined the same as Section M.4 and is summarized in Alg.4.

Original: Note Step 5 is optional in the n = m case.
Revised: Note Step 5 is optional in the n = m case.

Original: It gives the identity map if no arithmetic error due to machine precision is incurred.
Revised: It gives the identity map if no arithmetic error due to machine precision is incurred.

Original: However, this step is beneficial under low precision (see
Revised: However, this step is beneficial under low precision (see

Original: Apdx. M.1).
Revised: Apdx. M.1).

Original: In addition, we also have an adaptive version which was absent in Tao and Ohsawa (2020):

Original: Algorithm 4: Momentum SGD on SO ( n ) (generalization to Tao and Ohsawa (2020))
Revised: Algorithm 4: Momentum SGD on SO ( n ) (generalization to Tao and Ohsawa (2020))

Original: Hyperparameter : η ∈ (0 , + ∞ ) , γ ∈ [0 , + ∞ ) , a < 1 , maximum number of iteration N
Revised: Hyperparameter : η ∈ (0 , + ∞ ) , γ ∈ [0 , + ∞ ) , a < 1 , maximum number of iteration N

Original: Initialization : {+X 0 , V 0 , Y 0 s.t. X ⊤ 
+}
Revised: Initialization : {+X 0 , V 0 , Y 0 s.t. X ⊤ 
+}

Original: [Equation]{+X 
 = I , X ⊤ 0 V 0 = 0 , Y 
 + Y ⊤ 
 = 0 , b = aa −+}
Revised: {+X 
 = I , X ⊤ 0 V 0 = 0 , Y 
 + Y ⊤ 
 = 0 , b = aa −+}

Original: Output: Local minimum X of f 1 for i = 0 , ..., N − 1 do 
Revised: Output: Local minimum X of f 1 for i = 0 , ..., N − 1 do f

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: X

Original: [Equation]
Revised: [Equation]

Revised: − ∂f∂X ( X i )

Revised: ⊤ X i (cid:17)

Original: ηY
Revised: Y i

Original: [Equation]
Revised: [Equation]

Revised: = µY i − f i

Original: X 
Revised: X † [Equation] ηY i [Equation]

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: 6 end 7 return X N
Revised: 6 end 7 return X N

Revised: In addition, we also have an adaptive version which was absent in Tao and Ohsawa (2020):

Original: O. A DAM ON SO ( n )
Revised: O. A DAM ON SO ( n )

Original: In this section, we extend the above optimizer to an Adam version.
Revised: In this section, we extend the above optimizer to an Adam version.

Original: Note it can be derived from either the special case of the Stiefel optimizer, or the structure of Lie algebra so ( n ) , i.e., left-trivialization
Revised: Note it can be derived from either the special case of the Stiefel optimizer, or the structure of Lie algebra so ( n ) , i.e., left-trivialization

Original: (see Tao and Ohsawa (2020))
Revised: (see Tao and Ohsawa (2020))

Original: such that the momentum Y can be recognized as an element in
Revised: such that the momentum Y can be recognized as an element in

Original: T I SO ( n ) , the tangent space at the identity.
Revised: T I SO ( n ) , the tangent space at the identity.

Original: In the latter interpretation, the momentum will always stay in the same tangent space T I SO ( n ) which passes on a message that trivialization almost reproduce the convenience in Euclidean space for SO ( n ) .
Revised: In the latter interpretation, the momentum will always stay in the same tangent space T I SO ( n ) which passes on a message that trivialization almost reproduce the convenience in Euclidean space for SO ( n ) .

Original: Hence unlike the general Adam-Stiefel optimizer (Algo. 2), no extra technique, for example projection, is needed for Adam- SO ( n ) .
Revised: Hence unlike the general Adam-Stiefel optimizer (Algo. 2), no extra technique, for example projection, is needed for Adam- SO ( n ) .

Original: Detailed iterations are shown in Algo.
Revised: Detailed iterations are shown in Algo.

Original: Note this is also a structure preserving scheme (see Thm. 7).
Revised: Note this is also a structure preserving scheme (see Thm. 7).

Original: Algorithm 5: Adam on SO
Revised: Algorithm 5: Adam on SO

Original: Hyperparameter : η ∈ (0 , + ∞ ) , β 1 ∈ [0 , 1) , β 2 ∈ [0 , 1) , 0 < ϵ ≪ 1 , maximum number of iteration N
Revised: Hyperparameter : η ∈ (0 , + ∞ ) , β 1 ∈ [0 , 1) , β 2 ∈ [0 , 1) , 0 < ϵ ≪ 1 , maximum number of iteration N

Original: Initialization : X 0 , V 0 , Y 0 s.t.
Revised: Initialization : X 0 , V 0 , Y 0 s.t.

Original: X ⊤ 

Revised: X ⊤

Original: X 

Revised: X

Original: = I , X ⊤ 0 V 0 = 0 , Y 

Revised: = I , X ⊤ 0 V 0 = 0 , Y

Original: + Y ⊤ 

Revised: + Y ⊤

Original: = 0 , p 0 = 0 , q 0 =1 for i = 0 , ..., N − 1 do f i = 1 − b 2 (cid:16) X ⊤ i ∂f∂X ( X i ) − ∂f∂X
Revised: = 0 , p 0 = 0 , q 0 =1 for i = 0 , ..., N − 1 do f i = 1 − b 2 (cid:16) X ⊤ i ∂f∂X ( X i ) − ∂f∂X

Original: ( X i )
Revised: ( X i )

Original: ⊤ X i (cid:17) p i +1 =
Revised: ⊤ X i (cid:17) p i +1 =

Original: β 2 p
Revised: β 2 p

Original: i + (1 − β 2 ) f ◦ 2 i
Revised: i + (1 − β 2 ) f ◦ 2 i

Original: Y i +1 =
Revised: Y i +1 =

Original: β 1 Y i − (1 − β 1 )
Revised: β 1 Y i − (1 − β 1 )

Original: f
Revised: f

Original: i
Revised: i

Original: X †
Revised: X †

Original: =
Revised: =

Original: X i expm (cid:18) η (cid:113) 1 − β
Revised: X i expm (cid:18) η (cid:113) 1 − β

Original: i +12 Y
Revised: i +12 Y

Original: i +1 ⊘ ( p ◦− 12 i +1 + ϵ ) (cid:19)
Revised: i +1 ⊘ ( p ◦− 12 i +1 + ϵ ) (cid:19)

Original: X i +

Revised: X i +

Original: = X † ( X ⊤† X † ) − 12 .
Revised: = X † ( X ⊤† X † ) − 12 .

Original: 7 end 8 return X N
Revised: 7 end 8 return X N

Original: P E XPERIMENTAL DETAILS
Revised: P E XPERIMENTAL DETAILS

Original: Note: codes are provided in supplementary materials.
Revised: Note: codes are provided in supplementary materials.

Original: Experiments are conducted on a high-performance computing cluster whose name shall be revealed post anonymous period.
Revised: Experiments are conducted on a high-performance computing cluster whose name shall be revealed post anonymous period.

Original: Single v100 GPU was used.
Revised: Single v100 GPU was used.

Original: P. G ENERAL DISCUSSION
Revised: P. G ENERAL DISCUSSION

Original: Initialization For variable X with the initial X 0 , in fact, any initialization with full rank is valid, including the non-orthogonal X 0 such that X 0 does not start on the manifold.
Revised: Initialization For variable X with the initial X 0 , in fact, any initialization with full rank is valid, including the non-orthogonal X 0 such that X 0 does not start on the manifold.

Original: This is due to the structure-preserving property of our algorithm (both SGD and Adam versions, see Apdx. M).
Revised: This is due to the structure-preserving property of our algorithm (both SGD and Adam versions, see Apdx. M).

Original: After just one iteration, X will stay on the manifold.
Revised: After just one iteration, X will stay on the manifold.

Original: Here we suggest one way of obtaining orthogonal
Revised: Here we suggest one way of obtaining orthogonal

Original: X 0 which is performed in Saxe et al.
Revised: X 0 which is performed in Saxe et al.

Original: After randomly generating an entry-wise i.i.d.
Revised: After randomly generating an entry-wise i.i.d.

Original: normal distributed n -by- m matrix, we can perform the QR decomposition.
Revised: normal distributed n -by- m matrix, we can perform the QR decomposition.

Original: 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 .
Revised: 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 .

Original: Then we can take X 0 = ˜ Q ˜ R and thus have X ⊤ 0 X 

Revised: Then we can take X 0 = ˜ Q ˜ R and thus have X ⊤ 0 X

Original: = I .
Revised: = I .

Original: For the variables U , Z in Stiefel SGD and U , Z , p , q in Stiefel Adam, we would use zero matrices as their initialization since zero matrices always satisfy these constraints (28).
Revised: For the variables U , Z in Stiefel SGD and U , Z , p , q in Stiefel Adam, we would use zero matrices as their initialization since zero matrices always satisfy these constraints (28).

Original: 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. 1.1) requires separate adjustments of learning rates for the two types of parameters.
Revised: 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. 1.1) requires separate adjustments of learning rates for the two types of parameters.

Original: As an illustration, in projected SGD and Adam on Stiefel manifold Li et al.
Revised: As an illustration, in projected SGD and Adam on Stiefel manifold Li et al.

Original: (2020), learning rate may vary 10s of times for constrained and unconstrained parameters.
Revised: (2020), learning rate may vary 10s of times for constrained and unconstrained parameters.

Original: This can be intuitively understood as the result of applying different optimization methods for different parameters.
Revised: This can be intuitively understood as the result of applying different optimization methods for different parameters.

Original: However, using various learning rates may lead to the divergence of the algorithms.
Revised: However, using various learning rates may lead to the divergence of the algorithms.

Original: In contrast, our Stiefel method can be established for both the constrained and unconstrained parameters with the same learning rate.
Revised: In contrast, our Stiefel method can be established for both the constrained and unconstrained parameters with the same learning rate.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: For Stiefel Adam, we recommend to use 10 − 3 to be the learning rate and ( β 1 , β 2 ) =
Revised: For Stiefel Adam, we recommend to use 10 − 3 to be the learning rate and ( β 1 , β 2 ) =

Original: [Table]
Revised: [Table]

Original: How to choose between Adam-Stiefel and momentum
Revised: How to choose between Adam-Stiefel and momentum

Original: 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.
Revised: 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.

Original: P. D ETAILS OF PRW
Revised: P. D ETAILS OF PRW

Original: For MNIST experiment, we use a pretrained CNN to extract 128-dim.
Revised: For MNIST experiment, we use a pretrained CNN to extract 128-dim.

Original: features of figures in MNIST.
Revised: features of figures in MNIST.

Original: For Shakespeare’s plays experiment, we compute the PRW distances between all pairs of items in a corpus of eight Shakespeare’s operas embedded into 200-dim. space using word2vec.
Revised: For Shakespeare’s plays experiment, we compute the PRW distances between all pairs of items in a corpus of eight Shakespeare’s operas embedded into 200-dim. space using word2vec.

Original: For each digit or play pair, we use the extracted features as point sets { x i } and { y i } .
Revised: For each digit or play pair, we use the extracted features as point sets { x i } and { y i } .

Original: We choose the dimension of the target space of the projection to be k = 2 in both experiments.
Revised: We choose the dimension of the target space of the projection to be k = 2 in both experiments.

Original: The mean optimal transport values are taken among all digits or movie pairs at the specific iteration.
Revised: The mean optimal transport values are taken among all digits or movie pairs at the specific iteration.

Original: All the setting are same as the original paper Lin et al.
Revised: All the setting are same as the original paper Lin et al.

Original: (2020) except that early termination is removed.
Revised: (2020) except that early termination is removed.

Original:  The hyperparameters for each optimizer are listed in Tab. 3.
Revised: 6 The hyperparameters for each optimizer are listed in Tab. 3.

Original: P. D ETAILS OF V I T EXPERIMENTS
Revised: P. D ETAILS OF V I T EXPERIMENTS

Original: The structure of ViT and training process is exactly same as https://github.com/ omihub777/ViT-CIFAR but we use a different ViT implementation from https://github. com/lucidrains/vit-pytorch .
Revised: The structure of ViT and training process is exactly same as https://github.com/ omihub777/ViT-CIFAR but we use a different ViT implementation from https://github. com/lucidrains/vit-pytorch .

Original: The hyperparameters are carefully adjusted for each optimizer to obtain the best performance.
Revised: The hyperparameters are carefully adjusted for each optimizer to obtain the best performance.

Original: For projected Stiefel SGD/Adam Li et al.
Revised: For projected Stiefel SGD/Adam Li et al.

Original: , we tune the learning rates for constrained and unconstrained parameters separately.
Revised: , we tune the learning rates for constrained and unconstrained parameters separately.

Original: For regularizer SGD/Adam, the regularizer scaling parameter is also carefully chosen.
Revised: For regularizer SGD/Adam, the regularizer scaling parameter is also carefully chosen.

Original: The largest learning rate for momentumless
Revised: The largest learning rate for momentumless

Original: Stiefel SGD Wen and Yin (2013) is applied but momentum methods still admit faster convergence.
Revised: Stiefel SGD Wen and Yin (2013) is applied but momentum methods still admit faster convergence.

Original: Additional experimental results can be seen in Fig.
Revised: Additional experimental results can be seen in Fig.

Original: Model structure
Revised: Model structure

Original: The images are cut into 4 × 4 patches.
Revised: The images are cut into 4 × 4 patches.

Original: d model =
Revised: d model =

Original: d feedforward = 384 , n head = 12 .
Revised: d feedforward = 384 , n head = 12 .

Original: d q = d k = d v = 32 (which also equals d model /n head , so ‘orthogonality across heads’ constraint can be applied; see Section 3.2).
Revised: d q = d k = d v = 32 (which also equals d model /n head , so ‘orthogonality across heads’ constraint can be applied; see Section 3.2).

Original: The ‘classification token’ in Dosovitskiy et al.
Revised: The ‘classification token’ in Dosovitskiy et al.

Original: (2020) is used.
Revised: (2020) is used.

Original: Total number of layer of our ViT model is 7.
Revised: Total number of layer of our ViT model is 7.

Original: Data augmentation Following https://github.com/omihub777/ViT-CIFAR , we use the implementation from https://github.com/DeepVoltaire/AutoAugment for Cubuk et al.
Revised: Data augmentation Following https://github.com/omihub777/ViT-CIFAR , we use the implementation of (Cubuk et al., 2018) from https://github.com/DeepVoltaire/ AutoAugment .

Original: Training All the training uses the same scheduler, which is to let the learning rate linearly increase from 0 to the target learning rate in 5 epochs, followed by cosine annealing learning rate Loshchilov and Hutter (2016) with no restart, and minimum learning rate as 0.01 times max learning rate.
Revised: Training All the training uses the same scheduler, which is to let the learning rate linearly increase from 0 to the target learning rate in 5 epochs, followed by cosine annealing learning rate Loshchilov and Hutter (2016) with no restart, and minimum learning rate as 0.01 times max learning rate.

Original: Label smoothing with parameter 0.1 is used.
Revised: Label smoothing with parameter 0.1 is used.

Original: weight decay = 5 e − 5 unless specified.
Revised: weight decay = 5 e − 5 unless specified.

Original: All the hyperparameters used for optimizers are listed in Tab. 4.
Revised: All the hyperparameters used for optimizers are listed in Tab. 4.

Original: [Table] [Table]
Revised: [Table] [Table]

Original: Q A DDITIONAL EXPERIMENT : SYSTEMATIC TESTS ON THE LEADING EIGENVALUE PROBLEM
Revised: Q A DDITIONAL EXPERIMENT : SYSTEMATIC TESTS ON THE LEADING EIGENVALUE PROBLEM

Original: In this section, we consider the leading eigenvalue problem: given an n -by- n matrix A , find the m largest eigenvalues of A .
Revised: In this section, we consider the leading eigenvalue problem: given an n -by- n matrix A , find the m largest eigenvalues of A .

Original: This problem is at the core of many data sciences tasks, where n can be very large but m remains small.
Revised: This problem is at the core of many data sciences tasks, where n can be very large but m remains small.

Original: 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.
Revised: 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.

Original: Particularly, this problem can be formulated as an optimization problem on St ( n, m ) as follows
Revised: Particularly, this problem can be formulated as an optimization problem on St ( n, m ) as follows

Original: [Equation]
Revised: [Equation]

Revised: The idea is one seeks an optimal m -dimensional subspace, represented by X via m orthonormal bases in R n corresponding to its m columns, so that eigenvalues projected onto this subspace sums up to a maximum value.

Original: Such formulation is not unique.
Revised: Such formulation is not unique.

Original: For example, Tao and Ohsawa (2020) considers applying their momentum-accelerated SO ( n ) optimizer to arg min R ∈ SO ( n )
Revised: For example, Tao and Ohsawa (2020) consider applying their momentum-accelerated SO ( n ) optimizer to arg min R ∈ SO ( n )

Original: Tr ( E ⊤ R ⊤ ARE ) , where the constant padding matrix E :=
Revised: Tr ( E ⊤ R ⊤ ARE ) , where the constant padding matrix E :=

Original: [ I m × m ; 0 ( n − m ) × m ] .
Revised: [ I m × m ; 0 ( n − m ) × m ] .

Original: However, their setting is computationally less efficient than our Stiefel simplification.
Revised: However, their setting is computationally less efficient than our Stiefel simplification.

Revised: In fact, our formulation will be particularly suitable to cases when m is small but n is very large.

Original: Our test uses a deterministic matrix A generated by A
Revised: Our test uses a deterministic matrix A generated by A

Original: =
Revised: =

Original: (Ξ + Ξ ⊤ ) / 2 / √ n , where Ξ is an instance of n -by- n matrix with i.i.d.
Revised: (Ξ + Ξ ⊤ ) / 2 / √ n , where Ξ is an instance of n -by- n matrix with i.i.d.

Original: random normal elements.
Revised: random normal elements.

Original: The exact solution is obtainable from matrix decomposition so that the error can be quantified.
Revised: The exact solution is obtainable from matrix decomposition so that the error can be quantified.

Original: As is mentioned above, several aspects of our algorithm are examined, including convergence speed, manifold perseverance, and computational complexity in Fig.
Revised: As is mentioned above, several aspects of our algorithm are examined, including convergence speed, manifold perseverance, and computational complexity in Fig.

Original: 3, as well as other exact manifold-preserving methods.
Revised: 3, as well as other exact manifold-preserving methods.

Original: We can conclude that our (S)GD-Stiefel method has the fastest convergence rate and exact manifold perseverance with the lowest computational complexity.
Revised: We can conclude that our (S)GD-Stiefel method has the fastest convergence rate and exact manifold perseverance with the lowest computational complexity.

Original: Additionally, in Fig.
Revised: Additionally, in Fig.

Original: 4, we test the influence of different canonical-type metric, i.e., different a in
Revised: 4, we test the influence of different canonical-type metric, i.e., different a in

Original: Def. 1, and the three different ways to ‘compute’ matrix exponential in Apdx.
Revised: Def. 1, and the three different ways to ‘compute’ matrix exponential in Apdx.

Original: M.2, i.e., the matrix exponential itself, Cayley map, and forward Euler.
Revised: M.2, i.e., the matrix exponential itself, Cayley map, and forward Euler.

Original: No significant difference is found in leading eigenvalue test and as a result, the most commonly used canonical metric and the cheapest forward
Revised: No significant difference is found in leading eigenvalue test and as a result, the most commonly used canonical metric and the cheapest forward

Original: Euler are chosen as default in all the other experiments.
Revised: Euler are chosen as default in all the other experiments.

Original: The hyperparameter used for each algorithm are listed in Tab.
Revised: The hyperparameter used for each algorithm are listed in Tab.

Original: We also experimentally prove that better manifold preservation leads to better convergence, for both our algorithm 1 and Projected Stiefel SGD (Li et al., 2020) in Fig 5.
Revised: We also experimentally prove that better manifold preservation leads to better convergence, for both our algorithm 1 and Projected Stiefel SGD (Li et al., 2020) in Fig 5.

Original: What’s more, this experiment shows our (inner) iterative solver is quadratically convergent (as opposed to the linearly convergent
Revised: What’s more, this experiment shows our (inner) iterative solver is quadratically convergent (as opposed to the linearly convergent

Original: Cayley approximation in Li et al.
Revised: Cayley approximation in Li et al.

Original: (2020)), and
Revised: (2020)), and

Original: therefore a smaller number of iteration is needed to reach machine precision.
Revised: therefore a smaller number of iteration is needed to reach machine precision.

Original:  Though our special retraction on tangent bundle (polar retraction for position and no extra handling for momentum) is cheap, it must be used with our specially designed algorithm and cannot be applied to other algorithms even if the position will still stay on Stiefel manifold.
Revised: [Figure] Though our special retraction on tangent bundle (polar retraction for position and no extra handling for momentum) is cheap, it must be used with our specially designed algorithm and cannot be applied to other algorithms even if the position will still stay on Stiefel manifold.

Original: The reason is that the loss of structure of momentum will lead to slow convergence.
Revised: The reason is that the loss of structure of momentum will lead to slow convergence.

Original: Fig 6 shows that projected Stiefel SGD[Figure] with our retraction for tangent bundle convergences slower than not only our algorithm but also the original projected Stiefel SGD.
Revised: 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.

Original: [Figure]
Revised: [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Figure] [Figure]
Revised: [Figure] [Figure]