An Embarrassingly Simple Way to Optimize Orthogonal Matrices at Scale

Orthogonality constraints are ubiquitous in robust and probabilistic machine learning. Unfortunately, current optimizers are computationally expensive and do not scale to problems with hundreds or thousands of constraints. One notable exception is the Landing algorithm (Ablin et al., 2024) which, however comes at the expense of temporarily relaxing orthogonality. In this work, we revisit and improve on the ideas behind Landing, enabling the inclusion of modern adaptive optimizers while ensuring that orthogonal constraints are effectively met. Remarkably, these improvements come at little to no cost, and reduce the number of required hyperparemeters. Our algorithm POGO is fast and GPU-friendly, consisting of only 5 matrix products, and in practice maintains orthogonality at all times. On several challenging benchmarks, POGO greatly outperforms recent optimizers and shows it can optimize problems with thousands of orthogonal matrices in minutes while alternatives would take hours. As such, POGO sets a milestone to finally exploit orthogonality constraints in ML at scale. A public PyTorch implementation of POGO is available at https://github.com/adrianjav/pogo.

When building AI systems, engineers often want certain internal components to behave like a pure rotation---transforming data without stretching or distorting it. This turns out to make models more reliable and more stable to train. The catch is that existing methods for enforcing this property are far too slow for large models, or quietly abandon the requirement mid-training, defeating the purpose. We designed POGO, a new method that always respects this property, works naturally alongside the standard tools researchers already use to train AI, and requires less manual tuning than alternatives. On a modern graphics card, it solves problems in minutes that would otherwise take hours. Practically, this means that a whole class of AI models where genuinely maintaining this rotational structure matters becomes feasible to train at real-world scale for the first time. On top of that, we made our code freely available for everyone to use it.