# Momentum Stiefel Optimizer, with Applications to Suitably-Orthogonal Attention, and Optimal Transport

arXiv: https://arxiv.org/abs/2205.14173 · HTML: https://mlsanity.com/paper/2205.14173/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 84 references that are in the map.
- First posted on arXiv 27 May 2022.

## Builds on (most cited first, 30 of 44)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Training data-efficient image transformers & distillation through attention](https://mlsanity.com/paper/2012.12877/) (2020)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions](https://mlsanity.com/paper/2102.12122/) (2021)
- [Exact solutions to the nonlinear dynamics of learning in deep linear neural networks](https://mlsanity.com/paper/1312.6120/) (2013)
- [AutoAugment: Learning Augmentation Policies from Data](https://mlsanity.com/paper/1805.09501/) (2018)
- [Unitary Evolution Recurrent Neural Networks](https://mlsanity.com/paper/1511.06464/) (2015)
- [Parseval Networks: Improving Robustness to Adversarial Examples](https://mlsanity.com/paper/1704.08847/) (2017)
- [A Differential Equation for Modeling Nesterov's Accelerated Gradient Method: Theory and Insights](https://mlsanity.com/paper/1503.01243/) (2015)
- [Riemannian Adaptive Optimization Methods](https://mlsanity.com/paper/1810.00760/) (2018)
- [Dynamical Isometry and a Mean Field Theory of CNNs: How to Train 10,000-Layer Vanilla Convolutional Neural Networks](https://mlsanity.com/paper/1806.05393/) (2018)
- [Cheap Orthogonal Constraints in Neural Networks: A Simple Parametrization of the Orthogonal and Unitary Group](https://mlsanity.com/paper/1901.08428/) (2019)
- [Orthogonal Convolutional Neural Networks](https://mlsanity.com/paper/1911.12207/) (2019)
- [A Variational Perspective on Accelerated Methods in Optimization](https://mlsanity.com/paper/1603.04245/) (2016)
- [Nested Hierarchical Transformer: Towards Accurate, Data-Efficient and Interpretable Visual Understanding](https://mlsanity.com/paper/2105.12723/) (2021)
- [Full-Capacity Unitary Recurrent Neural Networks](https://mlsanity.com/paper/1611.00035/) (2016)
- [First-order Methods for Geodesically Convex Optimization](https://mlsanity.com/paper/1602.06053/) (2016)
- [Subspace Robust Wasserstein Distances](https://mlsanity.com/paper/1901.08949/) (2019)
- [On orthogonality and learning recurrent networks with long term dependencies](https://mlsanity.com/paper/1702.00071/) (2017)
- [Orthogonal Recurrent Neural Networks with Scaled Cayley Transform](https://mlsanity.com/paper/1707.09520/) (2017)
- [Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform](https://mlsanity.com/paper/2002.01113/) (2020)
- [Orthogonal Deep Neural Networks](https://mlsanity.com/paper/1905.05929/) (2019)
- [Riemannian SVRG: Fast Stochastic Optimization on Riemannian Manifolds](https://mlsanity.com/paper/1605.07147/) (2016)
- [Global rates of convergence for nonconvex optimization on manifolds](https://mlsanity.com/paper/1605.08101/) (2016)
- [Projection Robust Wasserstein Distance and Riemannian Optimization](https://mlsanity.com/paper/2006.07458/) (2020)
- [Regularizing CNNs with Locally Constrained Decorrelations](https://mlsanity.com/paper/1611.01967/) (2016)
- [A Brief Introduction to Manifold Optimization](https://mlsanity.com/paper/1906.05450/) (2019)

## Built on by (most cited first, 11 of 11)

- [Convergence of Kinetic Langevin Monte Carlo on Lie groups](https://mlsanity.com/paper/2403.12012/) (2024)
- [Symplectic Autoencoders for Model Reduction of Hamiltonian Systems](https://mlsanity.com/paper/2312.10004/) (2023)
- [Spectrum-Aware Parameter Efficient Fine-Tuning for Diffusion Models](https://mlsanity.com/paper/2405.21050/) (2024)
- [Stiefel Flow Matching for Moment-Constrained Structure Elucidation](https://mlsanity.com/paper/2412.12540/) (2024)
- [Generalizing Adam to Manifolds for Efficiently Training Transformers](https://mlsanity.com/paper/2305.16901/) (2023)
- [Efficient Optimization with Orthogonality Constraint: a Randomized Riemannian Submanifold Method](https://mlsanity.com/paper/2505.12378/) (2025)
- [Extragradient Type Methods for Riemannian Variational Inequality Problems](https://mlsanity.com/paper/2309.14155/) (2023)
- [Quantitative Convergences of Lie Group Momentum Optimizers](https://mlsanity.com/paper/2405.20390/) (2024)
- [Spectral-factorized Positive-definite Curvature Learning for NN Training](https://mlsanity.com/paper/2502.06268/) (2025)
- [Retraction-Free Optimization over the Stiefel Manifold for the LoRA Fine-Tuning](https://mlsanity.com/paper/2607.25299/) (2026)
- [ORCA: The Annealed Spectral Conditioning Optimizer for Faster, Better LLM Training](https://mlsanity.com/paper/2610.06116/) (2026)

## Abstract (arXiv)

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