MPAX用JAX实现高效线性与二次规划求解,支持自动微分和多硬件加速。
MPAX: Mathematical Programming in JAX
- 基于JAX构建,集成对偶混合梯度法与自适应优化技术
- 在GPU上比CPU快数倍,多卡扩展接近线性增速
- 适合需要可微训练的机器学习场景,如神经网络联合优化
我们提出MPAX(Mathematical Programming in JAX),一个原生基于JAX的大型线性规划(LP)与凸二次规划(QP)一阶求解器。目标是利用现代机器学习基础设施实现大规模数学规划,同时提供易于集成到机器学习流程中的先进算法。MPAX实现了两种PDHG变体:r2HPDHG用于LP,rAPDHG用于QP,包含对角预处理、自适应重启、自适应步长、原始权重更新、不可行性检测与可行性精化。借助JAX的编译与并行化生态,MPAX支持跨硬件移植、批量求解、分布式优化与自动微分。我们在CPU、NVIDIA GPU和Google TPU上进行了评估,观察到相比CPU基线有显著的GPU加速,且在标准LP/QP基准测试中表现具有竞争力。数值实验进一步证明其在高吞吐批量求解、密集LP的近线性多GPU扩展以及高效的端到端可微训练中的能力。求解器开源地址:https://github.com/MIT-Lu-Lab/MPAX。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。