# Generalizing Adam to Manifolds for Efficiently Training Transformers

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 67 references that are in the map.
- First posted on arXiv 26 May 2023.

## Builds on (most cited first, 14 of 14)

- [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)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Limitations of the Empirical Fisher Approximation for Natural Gradient Descent](https://mlsanity.com/paper/1905.12558/) (2019)
- [Riemannian Adaptive Optimization Methods](https://mlsanity.com/paper/1810.00760/) (2018)
- [Cheap Orthogonal Constraints in Neural Networks: A Simple Parametrization of the Orthogonal and Unitary Group](https://mlsanity.com/paper/1901.08428/) (2019)
- [Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform](https://mlsanity.com/paper/2002.01113/) (2020)
- [A mathematical model for automatic differentiation in machine learning](https://mlsanity.com/paper/2006.02080/) (2020)
- [Autotune: A Derivative-free Optimization Framework for Hyperparameter Tuning](https://mlsanity.com/paper/1804.07824/) (2018)
- [Effective Extensible Programming: Unleashing Julia on GPUs](https://mlsanity.com/paper/1712.03112/) (2017)
- [Momentum Stiefel Optimizer, with Applications to Suitably-Orthogonal Attention, and Optimal Transport](https://mlsanity.com/paper/2205.14173/) (2022)
- [Optimization on the symplectic Stiefel manifold: SR decomposition-based retraction and applications](https://mlsanity.com/paper/2211.09481/) (2022)

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

- [Symplectic Autoencoders for Model Reduction of Hamiltonian Systems](https://mlsanity.com/paper/2312.10004/) (2023)
- [Sparks of Quantum Advantage and Rapid Retraining in Machine Learning](https://mlsanity.com/paper/2407.16020/) (2024)
- [Natural gradient descent with momentum](https://mlsanity.com/paper/2604.15554/) (2026)

## Abstract (arXiv)

One of the primary reasons behind the success of neural networks has been the emergence of an array of new, highly-successful optimizers, perhaps most importantly the Adam optimizer. It is widely used for training neural networks, yet notoriously hard to interpret. Lacking a clear physical intuition, Adam is difficult to generalize to manifolds. Some attempts have been made to directly apply parts of the Adam algorithm to manifolds or to find an underlying structure, but a full generalization has remained elusive. In this work a new approach is presented that leverages the special structure of the manifolds which are relevant for optimization of neural networks, such as the Stiefel manifold, the symplectic Stiefel manifold and the Grassmann manifold: all of these are homogeneous spaces and as such admit a global tangent space representation. This is a common vector space, often called the Lie subspace, that makes the generalization of all steps in the Adam optimizer (as well as other optimizers) possible. It is thus possible to extend the Adam optimizer to manifolds without a projection step, something that was not possible before. The resulting algorithm is then applied to train transformers and a symplectic autoencoder for which orthogonality constraints are enforced up to machine precision and we conclusively demonstrate the advantage of the proposed optimizer over existing methods.