Stochastic Flows and Geometric Optimization on the Orthogonal Group
Krzysztof Choromanski,u00a0David Cheikhi,u00a0Jared Davis,u00a0Valerii Likhosherstov,u00a0Achille Nazaret,u00a0Achraf Bahamou,u00a0Xingyou Song,u00a0Mrugank Akarte,u00a0Jack Parker-Holder,u00a0Jacob Bergquist,u00a0Yuan Gao,u00a0Aldo Pacchiano,u00a0Tamas Sarlos,u00a0Adrian Weller,u00a0Vikas Sindhwani
We present a new class of stochastic, geometrically-driven optimization algorithms on the orthogonal group O(d) and naturally reductive homogeneous manifolds obtained from the action of the rotation group SO(d). 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 Humanoid agent from OpenAI Gym and improving convolutional neural networks.
