GPU-Enabled Large-Scale Optimization Using Randomized Linear Algebra

Pratik Rathore; Zachary Frangella; Parth Nobel; Xuning Hu; Madeleine Udell

Source-linked AI summary

GPU-Enabled Large-Scale Optimization Using Randomized Linear Algebra

Pratik Rathore, Zachary Frangella, Parth Nobel, Xuning Hu, Madeleine Udell

arXiv:2609.08136v1Sep 8, 2026cs.LGmath.OCstat.ML

TL;DR

RandNLA has strong algorithmic foundations, but few practitioner-facing implementations combine GPU acceleration with simple optimization modeling. rlaopt provides GPU-enabled RandNLA solvers and a natural-syntax modeling interface, with experiments showing substantial GPU speedups and regime-dependent benefits from randomized preconditioning.

Problem
Few implementations combine RandNLA algorithms, GPU acceleration, and a simple interface for specifying large-scale optimization problems.
Method
rlaopt combines GPU-enabled NyströmPCG, NysADMM, and SAPPHIRE solvers with a modeling language that checks compatibility, performs decomposition, and supports differentiation through iterations.
Results
The experiments show substantial GPU speedups and that randomized preconditioning benefits depend on spectral decay and regularization, including a 7.23× GPU speedup over CG for one ridge-regression setting.
Takeaways & Limitations
rlaopt provides a unified PyTorch framework for specifying and solving large-scale, ill-conditioned optimization problems with RandNLA methods and GPU execution.
Takeaways & Limitations
The experiments use accuracy checks that may be more stringent than needed for low-to-moderate precision, potentially favoring interior-point methods over first-order methods.

Abstract

from arXiv · show

This paper introduces rlaopt, a PyTorch-based package for large-scale optimization and scientific computing using randomized numerical linear algebra (RandNLA). Despite substantial progress in RandNLA-based algorithms, few implementations combine GPU acceleration with a simple interface for specifying optimization problems. rlaopt addresses this gap by providing GPU-enabled solvers for positive-definite linear systems and convex empirical risk minimization with constraints and regularizers. These solvers use RandNLA to accelerate conjugate gradient (NystromPCG), operator splitting (NysADMM), and stochastic gradient methods (SAPPHIRE). Moreover, rlaopt includes a modeling language that lets users specify problems using natural mathematical syntax. rlaopt automatically checks compatibility with the selected solver and performs the required problem decomposition. The solvers also support differentiation through their iterations, enabling applications such as hyperparameter tuning. Experiments on ridge regression, bounded multinomial logistic regression, and bounded elastic net identify when randomized preconditioning improves performance and demonstrate substantial speedups from GPU execution. The package is open-source under an Apache license, with source code at https://github.com/udellgroup/rlaopt and version 0.1.0 available on PyPI.

1 Introduction

rlaopt addresses the gap between RandNLA algorithms and practitioner-facing software with GPU-enabled solvers, a natural mathematical modeling language, and differentiable iterations. Its evaluation examines when randomized preconditioning and GPU execution improve large-scale optimization.

Motivation: RandNLA-based implementations remain limited despite strong algorithmic foundations and the scale demands of modern machine learning and scientific computing.Large datasets create optimization problems requiring efficient algorithms and hardware-aware implementations.
Package: rlaopt provides GPU-enabled NyströmPCG, NysADMM, and SAPPHIRE solvers for linear systems, constrained convex optimization, and empirical risk minimization.The package targets positive-definite systems, operator splitting, and stochastic optimization.
Modeling interface: Its modeling language lets users compose losses, regularizers, and constraints with natural mathematical syntax while checking solver compatibility and introducing ADMM splitting variables automatically.The interface is inspired by disciplined convex programming and CVXPY.
Evaluation: The experiments compare rlaopt with state-of-the-art methods on ridge regression, bounded multinomial logistic regression, and bounded elastic net regression.They characterize problem regimes favoring RandNLA solvers and GPU execution.
Differentiation: rlaopt supports differentiating through solver iterations for applications including hyperparameter tuning and end-to-end learning.This capability connects optimization to modern machine learning pipelines requiring gradient computation.

2 Problem Classes

rlaopt targets positive-definite linear systems and composite convex optimization with constraints and regularizers. These formulations cover applications such as kernel ridge regression, logistic regression, elastic net, and constrained multinomial regression.

Positive-Definite Linear Systems: The first problem class is solving symmetric positive-definite systems Ax = b, including regularized systems (K + λI)x = b.Here A is symmetric positive-definite, while K is symmetric positive-semidefinite and λ > 0.
Positive-Definite Linear Systems: These linear systems arise in kernel ridge regression and Gaussian process inference, where dense kernel matrices must be handled.Kernel ridge regression optimality conditions produce a system involving K + λI.
Positive-Definite Linear Systems: Direct methods require O(n^3) time and O(n^2) storage, while conjugate gradient has O(n^2) per-iteration cost but slows on ill-conditioned systems.Ill-conditioning is common for kernel matrices with rapidly decaying spectra.
Composite Convex Optimization: The second class is composite convex optimization with a smooth convex term, proxable nonsmooth terms, and linear mappings that encode constraints and regularizers.The formulation supports empirical risk minimization with terms such as ℓ1 penalties, box constraints, and elastic net regularization.
Solvers: rlaopt applies NysADMM when full gradients are available and SAPPHIRE when full-gradient computation is too expensive for large-scale empirical risk minimization.NysADMM accelerates ADMM linear-system subproblems, while SAPPHIRE uses randomized curvature estimates in stochastic variance-reduced methods.

3 Related Work

Related work spans direct and iterative linear-system solvers, kernel methods, interior-point and operator-splitting methods, proximal and stochastic optimization, and differentiable optimization frameworks. rlaopt is positioned as a GPU-enabled, RandNLA-based alternative for large, ill-conditioned problems across these areas.

Large-Scale Linear Systems: Direct Cholesky methods offer high accuracy but their O(n^3) cost and O(n^2) storage limit them, motivating iterative solvers for larger problems.Conjugate gradient improves scaling but requires effective preconditioning on ill-conditioned systems.
Large-Scale Linear Systems: Nyström preconditioning constructs low-rank approximations with sketching and is GPU-amenable because its dominant cost is matrix-matrix multiplication.rlaopt provides a GPU-enabled NyströmPCG implementation for ill-conditioned positive-definite systems.
Large-Scale Linear Systems: rlaopt differs from application-specific kernel tools by targeting general positive-definite systems and providing a modeling language for problem specification.This distinguishes it from tools such as FALKON and GPyTorch.
Optimization Solvers: Interior-point methods provide high-accuracy solutions but rely on matrix factorizations whose cubic cost limits their use on large-scale machine learning problems.rlaopt instead combines first-order and operator-splitting methods with RandNLA.
Optimization Solvers: NysADMM accelerates inexact ADMM by using NyströmPCG for ill-conditioned primal subproblems, including subproblems formed from second-order approximations.This extends operator splitting with randomized preconditioning.
Differentiable Optimization: rlaopt’s differentiable RandNLA solvers in PyTorch complement frameworks such as cvxpylayers, JAXopt, and MPAX by targeting different ecosystems or problem classes.The comparison includes implicit and unrolled differentiation approaches.

4 Modeling Language

rlaopt provides a composable modeling language for expressing optimization problems with natural mathematical syntax, then solving them through unified interfaces. It also supports differentiating through solver iterations for applications such as hyperparameter tuning.

Core abstractions: The modeling language uses variables, atoms, and solvers as its three core abstractions.
Variables: Variables specify optimized quantities, including their shape, optional name, and CPU or GPU device.
Atoms: Atoms provide mathematically characterized building blocks for losses, regularizers, and constraints, including smoothness and proximal-operator properties.
Objective composition: Users compose objectives with standard operators such as + and *, mirroring mathematical formulations including bounded elastic net.The example combines squared loss, elastic-net regularization, and box constraints.
Solvers: The unified solver interface supports stepped iteration and direct one-call solving, with both modes returning variable values.
Differentiation: rlaopt supports automatic differentiation through solver trajectories, yielding validation-loss derivatives with respect to parameters such as a lasso regularization value.
Model preparation: Unlike conic reformulation, rlaopt preserves composite structure and works directly with gradient and proximal oracles.The paper identifies this as advantageous for machine-learning problems such as logistic regression.

5 Automatic Detection of Problem Structure

rlaopt automatically analyzes composite objectives, validates whether their structure matches the selected solver, and performs the required splitting. Proximal gradient uses directly proxable, disjoint variable terms, while ADMM handles affine nonsmooth expressions through auxiliary variables and linear constraints.

Overview: rlaopt detects problem structure and decomposes user-specified objectives into forms compatible with proximal gradient or ADMM.
Objective partitioning: Atoms declare smoothness and proxability, allowing rlaopt to partition smooth terms into f and nonsmooth terms into g1, ..., gk without user input.
Proximal gradient: Proximal-gradient splitting requires each nonsmooth atom to act directly on a raw variable and nonsmooth atoms to use disjoint variable sets.
Compatibility checks: Violating proximal-gradient conditions triggers an informative compatibility error that can suggest using ADMM instead.
ADMM splitting: ADMM handles nonsmooth atoms applied to affine expressions by introducing auxiliary variables and consensus constraints.
ADMM splitting: Each ADMM decomposition extracts an auxiliary variable, a proxable atom, and the affine expression’s linear operator and offset.
Scalable operators: Linear operators are represented implicitly through matrix-vector and adjoint products, avoiding explicit formation of potentially very large data matrices.
ADMM execution: In ADMM, rlaopt accelerates the x-subproblem’s linear solve with Nyström preconditioning, including for affine ℓ1-regularized least squares.

6 Experiments

Across ridge regression, bounded multinomial regression, and bounded elastic net, rlaopt’s randomized solvers show problem-dependent advantages, especially for large, dense, ill-conditioned workloads. GPU execution provides substantial speedups, while solver effectiveness varies by dataset structure and baseline.

Ridge regression: 7.23× faster than CG, NyströmPCG was most advantageous for rapidly decaying spectra and weak regularization.At α = 2 and λ = 10−6 on GPU, the preconditioned solver achieved this ratio; at α = 0.5 and λ = 10−2, the ratio was 0.81.
Ridge regression: 57×–128× GPU speedups over CPU were observed for NyströmPCG across the synthetic ridge experiments.The measured solve time includes preconditioner construction.
Bounded multinomial logistic regression: SAPPHIRE succeeded on three GPU datasets but was slower than APG on all three and slower than JIT-enabled JAXopt baselines on its successful instances.GPU SAPPHIRE succeeded on CIFAR-10, News20, and RCV1, taking approximately 576, 613, and 490 seconds, respectively.
Bounded elastic net: NysADMM solved two large dense random-feature problems where no competing solver succeeded, whereas conic solvers were superior on sparse problems.GPU NysADMM took approximately 3093 and 811 seconds on acsincome-rf and yearpredictionmsd-rf; SCS was faster on yolanda-rf, realsim, and e2006.
GPU speedups: NysADMM’s GPU speedups exceeded 1.16–6.22× because CPU runs timed out on all five elastic-net datasets.The reported NysADMM speedups are lower bounds; successful GPU solves occurred on four datasets.
Differentiating through the solver: Differentiating through proximal-gradient iterations enabled gradient-based tuning of the lasso regularization parameter, with validation loss decreasing during optimization.The inner lasso solve and validation-loss gradient computation were repeated at each outer iteration.

7 Conclusion

rlaopt unifies GPU-enabled RandNLA solvers with a CVXPY-inspired modeling language and automatic problem decomposition. The conclusion reports substantial GPU speedups, problem-dependent randomized-preconditioning benefits, and differentiation through solver iterations, while documenting configuration and stopping behavior across the supported methods.

Contributions: rlaopt provides GPU-enabled NyströmPCG, NysADMM, and SAPPHIRE within a unified PyTorch-based framework for large-scale optimization.Its modeling language uses natural mathematical syntax and automatically detects problem structure for solver-specific decomposition.
Experimental conclusions: GPU NysADMM solves two dense bounded elastic-net instances where no competing conic solver succeeds, while conic solvers are faster on sparse instances.The comparison shows that solver performance depends on matrix structure and problem sparsity.
Experimental conclusions: Randomized preconditioning benefits depend on spectral decay and regularization, and SAPPHIRE benefits from GPU execution although JAXopt is faster on the multinomial instances it solves.These conclusions identify regimes where RandNLA acceleration helps and where competitor methods remain faster.
Differentiation: Differentiation through solver iterations enables gradient-based hyperparameter tuning.The package therefore supports optimization workflows that require gradients through the numerical solver.
Future work and availability: Future work includes more atoms and solvers, additional RandNLA scientific-computing applications, improved sparse-matrix performance, and a possible JAX backend with JIT compilation.The package is available as open-source software, with documentation and version 0.1.0 on PyPI.
Solver configuration: NyströmPCG, NysADMM, proximal gradient, and SAPPHIRE expose configuration and termination controls, including residual, gradient-mapping, and iteration-limit criteria.The appendix specifies solver defaults, preconditioner settings, and stopping behavior for the supported methods.

B Experimental Details

The experimental appendix documents the settings and data summaries used for the paper’s Section 6 experiments. It provides reference tables for the Nyström preconditioner and training dimensions, preprocessing, and class counts.

Experimental setup: The appendix provides details and settings for the experiments reported in Section 6.It serves as the experimental reference for the subsequent data and configuration tables.
Configuration: Table 10 summarizes the parameters controlling the randomized Nyström preconditioner.The table is the configuration reference for the preconditioner used by the solvers.
Datasets and preprocessing: Table 11 reports source and preprocessed feature dimensions, standardization and normalization procedures, and the number of classes for multinomial logistic regression.The column d denotes source feature dimension, p denotes preprocessed feature dimension, and K denotes the class count.

B.1 Datasets and Preprocessing

The experiments combine LIBSVM and OpenML datasets with task-specific preprocessing and explicit accuracy checks. Bounded problems use stationarity and feasibility criteria, while ridge regression uses a relative-residual check.

Datasets: Training data come from LIBSVM for seven named datasets and from OpenML for acsincome, yolanda, and fashion-mnist.Fashion-mnist uses the first 60000 of 70000 OpenML examples for training, and class labels are remapped to consecutive integers.
Preprocessing: acsincome and yolanda targets are standardized by centering and population standard deviation, while other targets retain their source values.realsim source labels are used as regression targets in the bounded elastic-net experiment.
Preprocessing: Sparse features remain sparse when supported, but rlaopt forces them to be dense.This preprocessing choice affects how sparse datasets are represented for the package’s experiments.
Random features: Random features use Gaussian or ReLU constructions, with Gaussian weights scaled by 1/√m and cosine features using σ = 1.The cosine map also uses independent phase offsets θ_j drawn uniformly from [0, 2π].
Accuracy checks: The bounded multinomial and elastic-net checks use task-specific gradients while separately enforcing feasibility.At a lower bound, a nonnegative gradient is stationary; at an upper bound, a nonpositive gradient is stationary.

B.3 Calibration and Production Refinement

The experiments calibrate solver-native tolerances against common accuracy checks and rerun failures with stricter tolerances. This procedure affects which runtimes and outcomes appear in the final figures.

Calibration: Calibration searches for the loosest native tolerance whose solutions pass common accuracy checks on every calibration instance for a solver and backend.The goal is to reduce runs on large datasets while retaining common checks as the production acceptance criterion.
Production refinement: A solver may satisfy its calibrated stopping tolerance yet fail the common checks, so experiments rerun it with a stricter tolerance.The final runtime is taken from the first qualifying attempt, including a successful strict-tolerance rerun.
Ridge outcomes: 13 ridge runs with residuals from 1.05×10^-6 to 3.09×10^-6 passed after reducing the native tolerance from 10^-9 to 10^-10.The affected runs included six CPU LSQR and seven GPU cuML LSMR runs.
Elastic net and multinomial outcomes: Direct GPU SCS on yearpredictionmsd-rf failed the common checks at 10^-7 and remained nonqualifying at 10^-8 and 10^-9 within the 3600-second budgets.It reached tolerance after 350 iterations in 365.56 seconds, with stationarity 1.178×10^-2 and feasibility violation 5.002×10^-6.

B.4 Hardware, Timing, and Solver Settings

The experiments use specified CPU and GPU hardware, distinct timing protocols for ridge and constrained optimization, and solver examples that support iterative execution and differentiable optimization.

Hardware: GPU jobs use NVIDIA H200 NVL GPUs, while each task receives 64 physical CPU cores and 128 GiB host memory.
Timing: Ridge timing excludes data generation, device placement, and warmup, while including synchronized GPU work and enforcing 900-second solve and startup limits.Measured ridge solves use three seeds, a 2p iteration ceiling, and separate accounting for warmup failures.
Timing: Bounded elastic net and multinomial logistic regression use cold solves with 3600-second solve limits, 1800-second startup limits, and at most 100000 iterations.Preprocessing and feature construction occur outside the timer, while solver-internal setup and JIT compilation are included.
Problem settings: The bounded regression formulation constrains coefficients through r = Xw + b1 −y with 0 ≤w ≤1.
Differentiable optimization: The differentiable optimization experiment uses 29 training observations, 27 validation observations, and 26 Gaussian features with 24 nonzero coefficients.Each inner solve uses 100 proximal-gradient steps, followed by 40 outer steps starting from µ = 0.2.
Solver settings: The examples apply operators without materializing X^T X, initialize solver state, and run fixed iteration counts for ridge, logistic, and elastic-net problems.Software versions are recorded in Table 13.

C Additional Experimental Results

The additional experiments extend ridge-regression sweeps across regularization levels, sample and feature dimensions, and matrix shapes, while also summarizing preconditioning and JIT-enabled multinomial logistic-regression comparisons. All figures use common accuracy checks and resource limits, with solve times reported as medians over three seeds and min–max error bars where specified.

Ridge problems: Figures 7–9 sweep square ridge problems across n = p at λ = 10^-2, 10^-4, and 10^-6.Each point reports median solve time over three seeds, with min–max error bars.
Ridge problems: Figures 10–12 fix p = 214 while varying n across λ = 10^-2, 10^-4, and 10^-6.The plots use median solve times over three seeds and min–max error bars.
Ridge problems: Figures 13–14 fix n = 216 while varying p at λ = 10^-2 and 10^-4.The plotted points summarize median solve time over three seeds, with min–max error bars.
Preconditioning and logistic regression: Figure 15 compares CG solve time with NyströmPCG solve time across eight matrix shapes, spectral decay rates α, and regularization levels λ.Inequalities indicate lower bounds when CG times out.
Preconditioning and logistic regression: Figure 16 compares SAPPHIRE with JAXopt baselines whose JIT compilation is included in solve time; the baselines are at least one order of magnitude faster.These supplementary figures use the same common accuracy checks and resource limits as the main-paper figures.