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

Oscar Dowson; Robert B Parker; Russel Bent

REVIEW 3 major objections 3 minor 2 references

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

T0 review · 3 major / 3 minor · reviewed 2026-08-06 · deepseek-v4-flash

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

Pith's one-line read This paper presents MathOptAI.jl, a Julia library that embeds trained machine-learning predictors into JuMP optimization models, including a gray-box mode that offloads derivative computation to a GPU.

desk verdict A useful, honestly-scoped software paper that deserves peer review, but the gray-box Hessian story is under-specified for the reduced-space path—worth asking the authors to pin down. read the letter →

arxiv 2507.03159 v2 pith:BF47PUQO submitted 2025-07-03 cs.LG math.OC

Oscar Dowson , Robert B Parker , Russel Bent This is my paper

classification cs.LGmath.OC

keywords MathOptAI.jlJuMPJuliamachinelearningpredictorsgray-boxformulationGPUaccelerationPyTorchnonlinearoptimization

verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

MathOptAI.jl is a Julia library whose purpose is to take a trained machine-learning predictor—a neural network, decision tree, random forest, or Gaussian process—and make it part of a mathematical optimization model built in JuMP. The paper's central claim is that this embedding can be done in three ways: a full-space formulation that adds variables and constraints for every piece of the predictor, a reduced-space formulation that inlines the predictor as one nonlinear expression, and a gray-box formulation that leaves the predictor's internals outside the solver and supplies function, Jacobian, and Hessian-of-the-Lagrangian evaluations through callbacks. The gray-box path is the paper's new contribution: among the five packages compared in Table 1, MathOptAI.jl is the only one with this formulation, and it is what allows PyTorch models to be evaluated on a GPU while the rest of the model remains a Julia-side JuMP problem. A sympathetic reader should care because the library removes the need for a closed algebraic form: if a predictor has trained weights and a differentiation engine, it can be optimized over directly.

What carries the argument

The load-bearing object is the AbstractPredictor interface, which forces every supported model to act as a vector-in, vector-out function y = F(x), and the Pipeline wrapper that composes predictors so a neural network is just a sequence of layers. On top of this, the gray-box formulation rests on JuMP's user-defined nonlinear operator support: F(x) is registered as an operator, automatic differentiation supplies ∇F(x), and the Hessian-of-the-Lagrangian is computed by forming the composed predictor G(x, λ) = (Affine{λ,0} ∘ F)(x) and differentiating it once, avoiding the need to build P separate Hessians. These pieces together let the solver treat the predictor as a black-box oracle while still receiving exact derivative information.

What would settle it

Take a small pretrained network with known weights, embed it in JuMP once in full-space and once in gray-box, and solve identical optimization problems with Ipopt; a disagreement between the two solutions beyond solver tolerance, or a mismatch between the gray-box Jacobian and Hessian-of-the-Lagrangian and finite-difference values at a random feasible point, would show that the oracle bridge does not represent y = F(x).

Watch

Extended reading notes

Core claim

The discovery the paper puts forward is an engineering capability: a general-purpose, solver-independent modeling library can embed a trained predictor without requiring the predictor to be written out algebraically, as long as the modeling layer can provide oracle evaluations. The paper shows the full mechanism for JuMP and Ipopt, and it argues that the gray-box formulation is what makes very large neural networks tractable, because the optimization model's size scales with the predictor's inputs and outputs rather than with the number of layers or parameters. On the paper's own terms, MathOptAI.jl establishes that the gray-box technique, previously demonstrated only in custom implementations, can be packaged in a general library and combined with GPU acceleration through Python.

Load-bearing premise

The gray-box formulation is only correct if the Julia–Python bridge returns the function value and the function's first- and second-order derivatives in exactly the order and convention Ipopt expects; if the derivatives are misordered, mis-scaled, or computed incorrectly, the solver is optimizing a different problem than y = F(x).

Editorial extensions

If this is right

Users can embed PyTorch, Flux, Lux, GLM, and DecisionTree models into a JuMP model and solve with any of the many solvers JuMP supports, not just one vendor's optimizer.
The gray-box formulation keeps the number of optimization variables and constraints tied to the predictor's input and output dimensions, so networks with very large parameter counts become embeddable; the authors point to related work demonstrating scales of hundreds of millions of parameters.
GPU-accelerated oracle evaluation is available for PyTorch models, moving function, Jacobian, and Hessian-of-the-Lagrangian computation to the GPU while the rest of the nonlinear oracles remain on the CPU.
Choosing a reformulation is a runtime option: the same embedded network can be solved with a nonlinear solver like Ipopt or, by changing the ReLU reformulation, with a mixed-integer solver like HiGHS.
Because predictors compose and third-party Julia packages plug in through package extensions, adding a new model type is a small amount of code, as demonstrated by the GLM adapter.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

Inference: the gray-box design is portable—any modeling layer with user-defined nonlinear operators and any automatic-differentiation engine with a GPU backend could reproduce it, so the Julia-plus-Python combination in the paper is an implementation choice rather than a requirement.
Inference: the authors leave a numerical comparison of the three formulations to future work, so a natural next test is to benchmark full-, reduced-, and gray-box on the same set of problems to map where each wins in iterations, wall time, and solution quality.
Inference: the vector-only input and output convention, which the paper calls its biggest limitation, points to an obvious extension—native support for n-dimensional tensors such as images, which would remove the reshaping burden the library currently places on users.
Inference: in an iterative training-and-optimization loop, retraining a surrogate would not require rewriting the optimization model, only reloading the predictor, so closed-loop pipelines could become simpler to maintain.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Desk Editor's Note

read the letter

This is a genuinely useful software paper, and the authors are honest about what it is: a general-purpose library for embedding trained ML predictors into JuMP, with the gray-box idea attributed to Casas et al. rather than claimed as new. The design principles are sensible, the comparison table is informative, and the code is archived. If you work on optimization with learned surrogates, this is a practical tool worth knowing about.

What the paper does well: it packages the gray-box formulation in a way that no other compared library does, supports GPU evaluation of the required oracles through PyTorch, and spells out the full-space, reduced-space, and gray-box formulations clearly. The composition abstraction (everything is a predictor F(x) mapping vectors to vectors) is a good design choice and is explained with concrete code. The authors also explicitly point to Parker et al. for numerical scaling results, so they are not overclaiming inside this paper.

The soft spots are real but not fatal. First, the stress-test note about reduced-space gray-box Hessians has force. Section 2.3.1 derives the Hessian contribution as the derivative of G(x, lambda) = lambda^T F(x), where lambda is the multiplier on the constraint F(x) - y = 0. That is correct for the full-space path. But the reduced-space path y := F(x) has no such equality constraint, so the correct weight on the Hessian of F is not an Ipopt multiplier for F(x)-y=0; it is the total derivative of the Lagrangian with respect to the operator's output. The paper does not say how JuMP's user-defined operator interface supplies that weight in the reduced-space case. The implementation might well be correct—the archive is the only evidence, and no numerical check appears in the paper—but the missing explanation leaves a hole in the central claim. A small verification example, even in an appendix, would close it.

Second, the paper has no standalone numerical validation. The authors disclose this and point to a companion paper, which is fair, but for a software contribution a minimal sanity check (e.g., a small problem where the gray-box solution matches the full-space solution) would strengthen the submission at low cost.

Overall: this is a solid engineering artifact, clearly written, with honest attribution and a useful feature comparison. The math that is shown is standard; the gap is in documenting the derivative plumbing for one advertised path. A serious referee should be asked to examine that plumbing, but this should not be desk-rejected—it deserves review.

Recommendation: send to peer review, and ask the authors to clarify or test the reduced-space gray-box Hessian path.

Referee Report

3 major / 3 minor

Summary. MathOptAI.jl is an open-source Julia package for embedding trained machine-learning predictors into JuMP optimization models. The paper describes three embedding formulations, full-space, reduced-space, and gray-box, and presents the gray-box approach as the novel contribution: the predictor is exposed to Ipopt through user-defined nonlinear operator callbacks for function, Jacobian, and Hessian-of-the-Lagrangian evaluations, with PyTorch-based oracle evaluations offloaded to a GPU through Julia's Python interface. The paper also gives a usage example, a feature comparison with OMLT, gurobi-machinelearning, PySCIPOpt-ML, and GAMSPy, and a set of design principles. The authors explicitly state that the paper contains no computational experiments and refer readers to a companion paper for performance evidence.

Significance. If the gray-box implementation is correct, this is a genuinely useful software contribution: it is the only package in the Table 1 comparison with a general-purpose gray-box formulation and GPU acceleration, and it fills a gap between Python machine-learning tooling and the JuMP optimization ecosystem. The algebra in Section 2.3.1 is standard and the library is transparently disclosed, with source code archived and prior work by Casas and coauthors credited. I see no circularity issue, since the formulations rest on standard nonlinear programming constructions rather than fitted parameter claims. However, the significance is conditional on the correctness of the derivative plumbing described in Sections 2.3 and 2.3.1, which is the least secured part of the paper: the reduced-space gray-box Hessian derivation is incomplete, and no in-paper validation is provided. The stress-test concern therefore lands and should be addressed before the central gray-box claim is accepted.

major comments (3)

[Sections 2.3 and 2.3.1] The gray-box Hessian derivation is written for the full-space constraint F(x)-y=0, where λ is the vector of Lagrange multipliers of the P equality constraints. The same section, however, states that the predictor can be added via the reduced-space form y := F(x). In the reduced-space problem there are no constraints y=F(x), so the correct weight multiplying ∇²F_k in the Hessian of the Lagrangian is the gradient of the original Lagrangian with respect to y_k, evaluated at y=F(x), not an Ipopt multiplier of a nonexistent constraint. The manuscript never explains how MathOptAI obtains this weight for the reduced-space path, and if it reuses the full-space multiplier vector, the Hessian is wrong and Ipopt may converge to a non-stationary point. Please either give the explicit reduced-space derivation and the exact operator-callback semantics, or restrict the gray-box claim to the full-space formulation.
[Section 2.3.1] Even for the full-space path, the identity G(x,λ) = (Affine{λ,0}∘F)(x) is correct only under a matching sign and ordering convention. The paper writes the constraint as F(x)-y=0 but does not state whether λ is the multiplier of F(x)-y=0 or y-F(x)=0, and it does not specify the row/column ordering of ∇F and of the Hessian block. Because JuMP and Ipopt also have their own conventions for how user-defined operator Hessians are weighted, an unexplained mismatch can silently produce a wrong search direction. Please state the conventions explicitly and verify them on a small example with a finite-difference or KKT-residual check.
[Sections 1.2, 2.4, and 6] The paper explicitly disclaims computational experiments and directs readers to a companion paper for performance evidence. This disclosure is honest, but for the gray-box claim it leaves the central novel mechanism unverified within the manuscript itself: the correctness of the derivative callbacks, not just their speed, is exactly what the algebraic description cannot establish with certainty. If the archived code contains derivative-consistency tests or correctness tests, the manuscript should say so and summarize them; otherwise, please add a minimal numerical example showing that Ipopt converges to the correct KKT point in both full-space and reduced-space gray-box modes.

minor comments (3)

[Section 1.1] The display for BinaryDecisionTree mixes an output expression with a set of side constraints; please define the path set P and leaf values y_p precisely, and avoid overloading y for both the predictor output and the leaf values.
[Section 3] The default non-smooth nonlinear ReLU formulation is used with Ipopt in the main code example, but Ipopt is a smooth NLP solver; the paper should either note which nonsmooth-capable solvers are expected to work with this default or explain how the formulation is regularized when Ipopt is selected.
[Table 1] For consistency and readability, use explicit 'No' entries instead of blanks in the Formulations and GPU-acceleration rows, so that readers can distinguish 'not supported' from 'not listed'.

Circularity Check

0 steps flagged · score 1.0 of 10

No significant circularity: the gray-box scheme is credited to prior external work and the central library claims are self-contained.

full rationale

The paper's central claim is that MathOptAI.jl provides general-purpose embeddings of trained predictors into JuMP models, not that a new mathematical result is derived from fitted data. There are no fitted parameters renamed as predictions, and no uniqueness theorem is imported from the authors' prior work. The gray-box formulation is explicitly attributed to Casas (2024) and Casas et al. (2025), so the novelty claim is limited to packaging the technique in a general-purpose library. The Hessian-of-the-Lagrangian construction defines G(x, lambda) = (Affine{lambda,0} circ F)(x) and then differentiates G; this is a direct computational identity with sum_i lambda_i grad^2 F_i(x), so it is tautological in execution but not a circular derivation of a result from its inputs. The only self-citations (Parker et al. 2025) concern scaling and performance evidence, which the paper expressly refers elsewhere and does not use to deduce the formulation; Section 2.4 even defers numerical comparison to future work, and the code archive provides an independently inspectable implementation. Concerns about whether the reduced-space gray-box path supplies the correct multiplier weights to Ipopt are correctness risks in derivative plumbing, not circularity: they do not make any claimed result equivalent to its own input by construction.

Assumptions & free parameters 0 free parameters · 3 assumptions · 0 invented entities

The library's claims rest on standard assumptions about the correctness of the user-defined operator interface in JuMP, the correctness of external automatic differentiation, and the smoothness of predictors in gray-box mode. No free parameters are fit, and no new physical or mathematical entities are introduced.

assumptions (3)

domain assumption JuMP's user-defined nonlinear operator interface correctly exposes the external predictor and its derivatives to Ipopt.
Section 2.3 relies on this interface to add F(x)-y=0 and provide Jacobian and Hessian-of-the-Lagrangian callbacks.
domain assumption PyTorch and Flux automatic differentiation return correct Jacobians and Hessians of the predictor.
Section 2.3.1 constructs G(x, lambda) = lambda^T F(x) and differentiates it to obtain the Hessian of the Lagrangian.
domain assumption Gray-box predictors are smooth, and twice differentiable when Hessians are requested.
Section 2.3 states that discrete or non-differentiable predictors such as BinaryDecisionTree do not have a gray-box formulation.

how reviews work

0 comments

Cite this review

Pith. "Pith review of MathOptAI.jl: Embed trained machine learning predictors into JuMP models." pith.science (2026). https://pith.science/paper/BF47PUQO

@misc{pith2026250703159,
  author       = {Pith},
  title        = {Pith review of: MathOptAI.jl: Embed trained machine learning predictors into JuMP models},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/BF47PUQO}},
  note         = {Machine review of arXiv:2507.03159}
}

read the original abstract

We present \texttt{MathOptAI.jl}, an open-source Julia library for embedding trained machine learning predictors into a JuMP model. \texttt{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 \texttt{Lux.jl} and \texttt{Flux.jl}, \texttt{MathOptAI.jl} uses Julia's Python interface to provide support for PyTorch models. When the PyTorch support is combined with \texttt{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.

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

2 extracted references · 1 canonical work pages

[1]

Albersmeyer J, Diehl M (2010) The lifted newton method and its application in optimization.SIAM Journal on Opti- mization20(3):1655–1684. Bates D, Noack A, Kornblith S, Bouchet-Valat M, Borregaard MK, Arslan A, White JM, Kleinschmidt D, Alday P, Lynch G, Dunning I, Mogensen PK, Lendle S, Aluthge D, Dutta M, pdeffebach, Jos´e Bayo´an Santiago Calder´on P, ...

arXiv 2010
[2]

Innes M (2018) Flux: Elegant machine learning with julia.Journal of Open Source Software3(25):602, URLhttp: //dx.doi.org/10.21105/joss.00602

Huangfu Q, Hall JAJ (2018) Parallelizing the dual revised simplex method.Mathematical Programming Computation 10(1):119–142. Innes M (2018) Flux: Elegant machine learning with julia.Journal of Open Source Software3(25):602, URLhttp: //dx.doi.org/10.21105/joss.00602. Legat B, Dowson O, Dias Garcia J, Lubin M (2021) MathOptInterface: a data structure for ma...

work page arXiv 2018

Pith tools

Reviewed August 6, 2026 · model on record in the stance chip above.

[1] [1]

Albersmeyer J, Diehl M (2010) The lifted newton method and its application in optimization.SIAM Journal on Opti- mization20(3):1655–1684. Bates D, Noack A, Kornblith S, Bouchet-Valat M, Borregaard MK, Arslan A, White JM, Kleinschmidt D, Alday P, Lynch G, Dunning I, Mogensen PK, Lendle S, Aluthge D, Dutta M, pdeffebach, Jos´e Bayo´an Santiago Calder´on P, ...

arXiv 2010

[2] [2]

Innes M (2018) Flux: Elegant machine learning with julia.Journal of Open Source Software3(25):602, URLhttp: //dx.doi.org/10.21105/joss.00602

Huangfu Q, Hall JAJ (2018) Parallelizing the dual revised simplex method.Mathematical Programming Computation 10(1):119–142. Innes M (2018) Flux: Elegant machine learning with julia.Journal of Open Source Software3(25):602, URLhttp: //dx.doi.org/10.21105/joss.00602. Legat B, Dowson O, Dias Garcia J, Lubin M (2021) MathOptInterface: a data structure for ma...

work page arXiv 2018