Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MPAX: Mathematical Programming in JAX

Haihao Lu, Zedong Peng, Jinwen Yang

TL;DR

MPAX introduces a JAX-native first-order solver for large-scale LPs and convex QPs that leverages XLA compilation, batching, and autodiff to integrate optimization directly into modern ML workflows. It implements two PDHG-based schemes—the restarted Reflected Halpern PDHG for LP ($r^2$HPDHG) and the Restarted Accelerated PDHG for QP (rAPDHG)—with diagonal preconditioning, adaptive restarts, feasibility polishing, and primal-weight updates. The framework supports across-hardware execution (CPU/GPU/TPU), true batched solving, distributed optimization via SPMD data sharding, and differentiable optimization through surrogate gradient techniques or unrolled differentiation. Empirical results show substantial GPU speedups, near-linear multi-GPU scaling for dense LPs, strong batched-solve performance, and competitive benchmark results relative to GPU-optimized solvers, all while enabling end-to-end differentiable decision-making in ML pipelines. MPAX is publicly available at https://github.com/MIT-Lu-Lab/MPAX.

Abstract

We present MPAX (Mathematical Programming in JAX), an open-source first-order solver for large-scale linear programming (LP) and convex quadratic programming (QP) built natively in JAX. The primary goal of MPAX is to exploit modern machine learning infrastructure for large-scale mathematical programming, while also providing advanced mathematical programming algorithms that are easy to integrate into machine learning workflows. MPAX implements two PDHG variants, r2HPDHG for LP and rAPDHG for QP, together with diagonal preconditioning, adaptive restarts, adaptive step sizes, primal-weight updates, infeasibility detection, and feasibility polishing. Leveraging JAX's compilation and parallelization ecosystem, MPAX provides across-hardware portability, batched solving, distributed optimization, and automatic differentiation. We evaluate MPAX on CPUs, NVIDIA GPUs, and Google TPUs, observing substantial GPU speedups over CPU baselines and competitive performance relative to GPU-based codebases on standard LP/QP benchmarks. Our numerical experiments further demonstrate MPAX's capabilities in high-throughput batched solving, near-linear multi-GPU scaling for dense LPs, and efficient end-to-end differentiable training. The solver is publicly available at https://github.com/MIT-Lu-Lab/MPAX.

MPAX: Mathematical Programming in JAX

TL;DR

MPAX introduces a JAX-native first-order solver for large-scale LPs and convex QPs that leverages XLA compilation, batching, and autodiff to integrate optimization directly into modern ML workflows. It implements two PDHG-based schemes—the restarted Reflected Halpern PDHG for LP (HPDHG) and the Restarted Accelerated PDHG for QP (rAPDHG)—with diagonal preconditioning, adaptive restarts, feasibility polishing, and primal-weight updates. The framework supports across-hardware execution (CPU/GPU/TPU), true batched solving, distributed optimization via SPMD data sharding, and differentiable optimization through surrogate gradient techniques or unrolled differentiation. Empirical results show substantial GPU speedups, near-linear multi-GPU scaling for dense LPs, strong batched-solve performance, and competitive benchmark results relative to GPU-optimized solvers, all while enabling end-to-end differentiable decision-making in ML pipelines. MPAX is publicly available at https://github.com/MIT-Lu-Lab/MPAX.

Abstract

We present MPAX (Mathematical Programming in JAX), an open-source first-order solver for large-scale linear programming (LP) and convex quadratic programming (QP) built natively in JAX. The primary goal of MPAX is to exploit modern machine learning infrastructure for large-scale mathematical programming, while also providing advanced mathematical programming algorithms that are easy to integrate into machine learning workflows. MPAX implements two PDHG variants, r2HPDHG for LP and rAPDHG for QP, together with diagonal preconditioning, adaptive restarts, adaptive step sizes, primal-weight updates, infeasibility detection, and feasibility polishing. Leveraging JAX's compilation and parallelization ecosystem, MPAX provides across-hardware portability, batched solving, distributed optimization, and automatic differentiation. We evaluate MPAX on CPUs, NVIDIA GPUs, and Google TPUs, observing substantial GPU speedups over CPU baselines and competitive performance relative to GPU-based codebases on standard LP/QP benchmarks. Our numerical experiments further demonstrate MPAX's capabilities in high-throughput batched solving, near-linear multi-GPU scaling for dense LPs, and efficient end-to-end differentiable training. The solver is publicly available at https://github.com/MIT-Lu-Lab/MPAX.

Paper Structure

This paper contains 28 sections, 19 equations, 5 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related works
  3. Preliminary: Primal-Dual First-Order Algorithms for LP and QP
  4. Restarted Reflected Halpern PDHG for LP
  5. Restarted Accelerated PDHG for QP
  6. MPAX: Math Programming in JAX
  7. Across-hardware portability and hardware acceleration
  8. Batched solving
  9. Distributed optimization
  10. Auto-differentiation
  11. Miscellaneous
  12. Experiments
  13. Across-hardware portability and hardware acceleration
  14. Batched solving
  15. Distributed Optimization
  16. ...and 13 more sections

Figures (5)

  • Figure 1: Design of MPAX
  • Figure 2: Solve time of MPAX and Gurobi for different batch sizes on the storm instance.
  • Figure 3: MPAX's distributed optimization performance on a multi-dimensional knapsack problem.
  • Figure 4: Warcraft shortest path dataset: Warcraft terrain map (left), vertex traversal cost (middle) and shortest path (right)
  • Figure 5: The loss and normalized regret curve of the Warcraft shortest path problem