GPU-enabled large-scale optimization using randomized linear algebra

Despite the growing importance of optimization, existing solvers suffer from several key drawbacks: (i) they often cannot scale to large-scale optimization problems due to high computational/memory costs (e.g., expensive matrix factorizations), (i) they are unable to take advantage of modern parallel hardware (i.e., limited to CPU), and (iii) they do not provide a simple, user-friendly syntax for modeling modern optimization problems. To address these issues, we develop rlaopt, a PyTorch-based package containing high-quality, GPU-enabled methods for large-scale optimization and scientific computing. rlaopt is easy to use: we create a flexible modeling language that lets users specify problems via natural mathematical syntax and demonstrate its simplicity for formulating and solving large-scale linear systems, empirical risk minimization, and more. Randomized numerical linear algebra is central to our approach, as it mitigates ill-conditioning for linear systems and unlocks scalability in stochastic optimization methods and operator splitting frameworks. Moreover, rlaopt allows for differentiating through the solver, making it suitable for modern machine learning pipelines.