MPAX: Mathematical Programming in JAX

要約

このペーパーでは、線形プログラミング（LP）を機械学習ワークフローに統合するための多用途で効率的なツールボックスであるMPAX（JAXの数学プログラミング）を紹介します。
MPAXは、最先端の1次方法を実装し、平均プライマルデュアルハイブリッド勾配を再開し、JAXのLPSを解くために再起動Halpern Primal-Dual Hybrid Gradientを反映しました。
これにより、ハードウェアアクセラ化のネイティブサポートと、バッチ解決、自動分化、デバイスの並列性などの機能が提供されます。
広範な数値実験は、既存のソルバーに対するMPAXの利点を示しています。
ソルバーはhttps://github.com/mit-lu-lab/mpaxで入手できます。

要約(オリジナル)

This paper presents MPAX (Mathematical Programming in JAX), a versatile and efficient toolbox for integrating linear programming (LP) into machine learning workflows. MPAX implemented the state-of-the-art first-order methods, restarted average primal-dual hybrid gradient and reflected restarted Halpern primal-dual hybrid gradient, to solve LPs in JAX. This provides native support for hardware accelerations along with features like batch solving, auto-differentiation, and device parallelism. Extensive numerical experiments demonstrate the advantages of MPAX over existing solvers. The solver is available at https://github.com/MIT-Lu-Lab/MPAX.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google