Generalizing Adam to Manifolds for Efficiently Training Transformers

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.