MathOptAI.jl: Embed trained machine learning predictors into JuMP models

2025-07-03

AI Summary

This work addresses the challenge of efficiently integrating machine learning models with mathematical optimization frameworks—specifically JuMP. We propose a gray-box modeling approach that embeds pre-trained models (e.g., neural networks, decision trees, Gaussian processes) from PyTorch—and compatibly supports Flux.jl and Lux.jl—directly into JuMP. Leveraging Julia–Python interoperability, function evaluations, gradients, and Hessians are computed on GPU-accelerated Python backends, while remaining nonlinear operations execute natively on Julia’s CPU for accuracy and efficiency. The method enables end-to-end differentiable optimization with rigorous support for mixed-integer nonlinear programming (MINLP). Empirical results demonstrate significant improvements in MINLP solver performance for data-driven optimization tasks. An open-source implementation under the BSD-3 license is publicly available on GitHub, providing a scalable, high-performance modeling infrastructure for integrating learned components into optimization workflows.

Technology Category

• Machine Learning: Optimization
• Constraint Satisfaction and Optimization: Mixed Discrete/Continuous Optimization
• Search and Optimization: Mixed Discrete/Continuous Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

Abstract

We present exttt{MathOptAI.jl}, an open-source Julia library for embedding trained machine learning predictors into a JuMP model. exttt{MathOptAI.jl} can embed a wide variety of neural networks, decision trees, and Gaussian Processes into a larger mathematical optimization model. In addition to interfacing a range of Julia-based machine learning libraries such as exttt{Lux.jl} and exttt{Flux.jl}, exttt{MathOptAI.jl} uses Julia's Python interface to provide support for PyTorch models. When the PyTorch support is combined with exttt{MathOptAI.jl}'s gray-box formulation, the function, Jacobian, and Hessian evaluations associated with the PyTorch model are offloaded to the GPU in Python, while the rest of the nonlinear oracles are evaluated on the CPU in Julia. MathOptAI is available at https://github.com/lanl-ansi/MathOptAI.jl under a BSD-3 license.

Problem

• Embed trained ML predictors into JuMP optimization models
• Support diverse ML models like neural networks and trees
• Enable GPU acceleration for PyTorch models in Julia

Innovation

• Embed ML predictors into JuMP models
• Support neural networks, trees, Gaussian Processes
• Offload PyTorch evaluations to GPU