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 wiedely 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, the Grassmann manifold and the symplectic Grassmann manifold: all of these are homogeneous spaces and as such admit a global tangent space representation. This global tangent space representation is used to perform all of the steps in the Adam optimizer. The resulting algorithm is then applied to train a transformer for which orthogonality constraints are enforced up to machine precision and we observe significant speed-ups in the training process. Optimization of neural networks where they weights do not lie on a manifold is identified as a special case of the presented framkework. This allows for a flexible implementation in which the learning rate is adapted simultaneously for all parameters, irrespective of whether they are an element of a general manifold or a vector space.

翻译：神经网络成功的主要原因之一是一系列高效新优化器的涌现，其中最重要的可能是Adam优化器。它被广泛用于训练神经网络，但因其难以解释而著称。由于缺乏清晰的物理直觉，Adam难以推广到流形上。已有一些尝试直接将Adam算法中的部分步骤应用于流形，或寻找其底层结构，但完整的推广仍未能实现。本文提出了一种新方法，利用与神经网络优化相关的特殊流形结构，例如Stiefel流形、辛Stiefel流形、Grassmann流形和辛Grassmann流形：所有这些流形都是齐次空间，因此具有全局切空间表示。该全局切空间表示被用于执行Adam优化器中的所有步骤。由此产生的算法随后被应用于训练Transformer，其中正交性约束被强制执行到机器精度，并且我们观察到训练过程实现了显著加速。对于权重不位于流形上的神经网络优化，被识别为所提出框架的一个特例。这使得实现具有灵活性，即学习率能够同时适用于所有参数进行自适应调整，无论这些参数是位于一般流形上还是向量空间中。