提出轻量级算法POGO,高效实现大规模正交矩阵优化
An Embarrassingly Simple Way to Optimize Orthogonal Matrices at Scale
- 基于改进的Landing思想,仅需5次矩阵乘法实现正交约束
- 在千级正交约束下分钟级完成优化,速度远超现有方法
- 兼容自适应优化器,适合需要高精度正交性的大规模模型
正交性约束在鲁棒和概率机器学习中广泛应用,但现有优化器计算成本高,难以扩展到数百甚至数千个约束。尽管Landing算法(Ablin等,2024)是例外,却需临时放松正交性。本文重新审视并改进Landing思路,在保持正交性的同时引入现代自适应优化器,且几乎不增加计算开销与超参数数量。所提算法POGO高效且适配GPU,仅包含5次矩阵乘法,在实践中始终维持正交性。在多个挑战性基准测试中,POGO显著优于近期优化器,可在数分钟内完成千级正交矩阵的优化,而其他方法需数小时。POGO为大规模机器学习中有效利用正交约束树立了里程碑。PyTorch实现已公开于https://github.com/adrianjav/pogo。
原文摘要 · Abstract (English)
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 PyTorch implementation of POGO is publicly available at https://github.com/adrianjav/pogo.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。