Stochastic Flows and Geometric Optimization on the Orthogonal Group
Krzysztof Choromanski, David Cheikhi, Jared Davis, Valerii Likhosherstov, Achille Nazaret, Achraf Bahamou, Xingyou Song, Mrugank Akarte, Jack Parker-Holder, Jacob Bergquist, Yuan Gao, Aldo Pacchiano
Abstract
We present a new class of stochastic, geometrically-driven optimization algorithms on the orthogonal group and naturally reductive homogeneous manifolds obtained from the action of the rotation group . We theoretically and experimentally demonstrate that our methods can be applied in various fields of machine learning including deep, convolutional and recurrent neural networks, reinforcement learning, normalizing flows and metric learning. We show an intriguing connection between efficient stochastic optimization on the orthogonal group and graph theory (e.g. matching problem, partition functions over graphs, graph-coloring). We leverage the theory of Lie groups and provide theoretical results for the designed class of algorithms. We demonstrate broad applicability of our methods by showing strong performance on the seemingly unrelated tasks of learning world models to obtain stable policies for the most difficult agent from and improving convolutional neural networks.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext cdd6e73b-b11f-4f9b-9416-c35b7dc6c8a9Cited by top-tier papers1
Ask how each one uses itBuilds on1
Related papers
- Coordinate Descent on the Orthogonal Group for Recurrent Neural Network TrainingEstelle M. Massart, Vinayak AbrolAAAI 2022 · 13 citations
- Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley TransformJun Li, Fuxin Li, Sinisa TodorovicICLR 2020 · 139 citations
- A General Framework for Equivariant Neural Networks on Reductive Lie GroupsIlyes Batatia, Mario Geiger, Jose M. Munoz, Tess E. Smidt et al.NeurIPS 2023 · 28 citations
- Riemannian SAM: Sharpness-Aware Minimization on Riemannian ManifoldsJihun Yun, Eunho YangNeurIPS 2023 · 10 citations
- A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix GroupsMarc Finzi, Max Welling, Andrew Gordon WilsonICML 2021 · 226 citations
