1 Introduction↩︎

A recent trend in the mathematical optimization literature is to embed trained machine learning predictors into a larger optimization model. The most common application is for a practitioner to train a machine learning predictor as a surrogate for a more complicated subsystem that cannot be directly embedded into an optimization model, for example, because it does not have an algebraic form or because it is non-differentiable. [1] provide a review of the field.

We present MathOptAI.jl, an open-source Julia [2] library for embedding trained machine learning predictors into a JuMP [3] model. MathOptAI.jl can embed a wide variety of neural networks, decision trees, and Gaussian processes into a larger mathematical optimization model. MathOptAI.jl supports a range of mixed-integer linear, smooth nonlinear, and non-smooth nonlinear reformulations. The design of MathOptAI.jl is inspired by previous work such as gurobi-machinelearning[4] and OMLT[5].

More formally, the purpose of MathOptAI.jl is to build optimization problems in the form given by Problem 1 : \[\begin{array}{cl} \displaystyle \min_{\boldsymbol{x}\in\mathbb{R}^N,\;\boldsymbol{y}\in\mathbb{R}^P} & f_0(\boldsymbol{x},\boldsymbol{y}) \\ \text{s.t.} & f_i(\boldsymbol{x},\boldsymbol{y}) \in S_i,\quad i = 1,\ldots,M\\ & \boldsymbol{y} = F(\boldsymbol{x}), \\ \end{array} \label{eqn:nlopt}\tag{1}\] where \(f_i(\boldsymbol{x}, \boldsymbol{y}): \mathbb{R}^{N+P} \rightarrow \mathbb{R}^{D_i}\) is a function that must belong to the set \(S_i\subset \mathbb{R}^{D_i}\). Note that, here and throughout, bold-faced symbols are vectors, and \(x_i\) is the \(i\)-th component of the vector \(\boldsymbol{x}\). The functions \(f_i\) and sets \(S_i\) are defined by the MathOptInterface standard form [6] and abstract over the various problem classes of mathematical optimization, from mixed-integer linear to nonlinear. The relationship \(\boldsymbol{y} = F(\boldsymbol{x})\), where \(F(\boldsymbol{x}): \mathbb{R}^N \rightarrow \mathbb{R}^P\), is the machine learning predictor (or concatenation of multiple predictors) that we want to embed in the optimization model. We delay explaining how we embed the relationship \(\boldsymbol{y} = F(\boldsymbol{x})\) until Section 2.

1.1 Examples of machine learning predictors↩︎

Examples of predictors \(F(\boldsymbol{x})\) supported by MathOptAI.jl include:

• \(\textsf{Affine}_{\{A,\boldsymbol{b}\}}(\boldsymbol{x}) = A\boldsymbol{x} + \boldsymbol{b}\)

• \(\textsf{ReLU}(\boldsymbol{x}) = \max(\boldsymbol{0}, \boldsymbol{x})\)

• \(\textsf{Sigmoid}(\boldsymbol{x}) = \frac{1}{1 + e^{-\boldsymbol{x}}}\)

• \(\textsf{SoftMax}(\boldsymbol{x}) = \frac{e^\boldsymbol{x}}{||e^\boldsymbol{x}||_1}\)

• \(\textsf{Tanh}(\boldsymbol{x}) = \tanh(\boldsymbol{x})\).

By convention, all predictors take vector inputs and produce vector outputs. Some predictors are naturally vector-valued, for example, \(\textsf{Affine}_{\{A,\boldsymbol{b}\}}(\boldsymbol{x}): \mathbb{R}^m \rightarrow \mathbb{R}^n\), where \(m\) and \(n\) are derived from the shape of \(A\) and \(\boldsymbol{b}\). Other predictors involve operators (like \(e\), \(\max\), and \(\tanh\)) that produce scalar outputs. When scalar-valued operators are applied to a vector input, the convention is that the operator is applied element-wise.

MathOptAI.jl also supports predictors which are not simple algebraic equations. One example is a binary decision tree (\(\mathbb{R}^m \rightarrow \mathbb{R}^{1}\)), which can be represented by the set of constraints: \[\begin{align} \textsf{BinaryDecisionTree}_{\{\mathcal{P},f,y\}}(\boldsymbol{x}) = & \sum\limits_{p\in \mathcal{P}} \delta_p y_p \\ \text{such that:\;\;\;}& \delta_p \implies f_n(\boldsymbol{x}) \le 0 & \forall p\in\mathcal{P}, n \in p \\ & \sum\limits_{p\in \mathcal{P}} \delta_p=1 \\ & \delta_p \in \{0, 1\} & \forall p \in \mathcal{P}, \end{align}\] where \(\mathcal{P}\) is a set of paths through the tree, \(y_p\) is the leaf value associated with path \(p\), and \(f_n\) describes the relation that must hold at each node \(n\) if path \(p\) is taken [4].

Some predictors are defined by the composition of simpler predictors. For example, a logistic regression predictor is: \[\boldsymbol{y} = F(\boldsymbol{x}) = \frac{1}{1 + e^{-(A\boldsymbol{x} + \boldsymbol{b})}},\] which is the composition: \[\boldsymbol{y} = F(\boldsymbol{x}) = (\textsf{Sigmoid} \circ \textsf{Affine}_{\{A,\boldsymbol{b}\}})(\boldsymbol{x}).\]

Similarly, a feed-forward neural network predictor is defined by the repeated application of an affine transformation and a nonlinear activation function \(\sigma\) over \(L\) layers: \[\begin{align} \boldsymbol{y}_l &= \sigma_l (W_l \boldsymbol{y}_{l-1} + \boldsymbol{b}_l), \quad l =1,\dots,L, \end{align}\] where \(\boldsymbol{y}_0 = \boldsymbol{x}\) and \(\boldsymbol{y} = \boldsymbol{y}_L\). Examples of activation functions include \(\sigma(x) = \max(0, x)\), and \(\sigma(x) = \tanh(x)\). Instead of considering a layer comprised of both an affine transform and an activation function, MathOptAI.jl models a neural network as the sequential composition of predictors. For example: \[\boldsymbol{y} = F(\boldsymbol{x}) = (\textsf{SoftMax} \circ \textsf{Affine}_{\{W_L,\boldsymbol{b}_L\}} \circ \cdots \circ \textsf{ReLU} \circ \textsf{Affine}_{\{W_1,\boldsymbol{b}_1\}})(\boldsymbol{x}).\]

In addition to the predictors listed above, MathOptAI.jl also supports more complicated predictors, such as random forests, gradient boosted trees, and a predictor that returns the quantile of a distribution or Gaussian Process. We omit these formulations in the interest of brevity. Moreover, due to the on-going open-source development of the package, we also omit a complete list of the supported predictors since the list will inevitably become outdated. Readers should consult the associated code archive [7] for the full list of supported predictors.

1.2 Contribution and Outline↩︎

The main contribution of this paper is the Julia library MathOptAI.jl, an open-source Julia [2] library for embedding trained machine learning predictors into a JuMP [3] model. The MathOptAI.jl library is archived at [7].

A novel feature of MathOptAI.jl is its support for the gray-box formulation of a neural network (discussed in Section 2.3). Most usefully, the gray-box formulation allows us to embed large-scale neural networks into a Julia-based JuMP model whilst evaluating the associated zeroth-, first-, and second-order derivatives of the network’s loss function on a GPU from Python. Although we are not the first to develop the technique (see [8] and [9]), we are the first to present it in a general purpose library. As secondary contributions, we provide a review of software alternatives to MathOptAI.jl, and we explain our design principles. These contributions may be useful for authors looking to implement a similar package in other programming languages.

Readers should note that the purpose of this paper is not to provide computational experiments of the relative merits of embedding machine learning predictors into an optimization model. We direct interested readers to our related work [10], and to the ever growing literature on the subject like [9].

The rest of this paper is organized as follows. In Section 2, we explain how we embed a predictor into an optimization model. In Section 3 we provide a code example of using MathOptAI.jl. Section 4 is a comparison to related work, and we conclude in Section 5 with our guiding design principles.

2 Algebraic representation↩︎

This section explains the three ways in which MathOptAI.jl embeds the \(\boldsymbol{y} = F(\boldsymbol{x})\) equation into an optimization model. The three approaches are denominated full-space, reduced-space, and gray-box.

2.1 Full-space↩︎

In the full-space formulation, we add a vector-valued decision variable \(\boldsymbol{y}\) to represent the output of each predictor, and we add constraints and additional variables as necessary to ensure that \(\boldsymbol{y} = F(\boldsymbol{x})\) holds in a feasible solution. For example, an \(\textsf{Affine}_{\{A,\boldsymbol{b}\}}\) predictor is encoded as the vector-valued linear equality constraint: \[\boldsymbol{y} = A\boldsymbol{x} + \boldsymbol{b},\] and a \(\textsf{Sigmoid}\) predictor is encoded as the set of nonlinear equality constraints: \[y_i = \frac{1}{1 + e^{-x_i}},\quad i = 1,\ldots,P,\] where \(P\) is the output dimension of the predictor.

Predictors formed by the composition of other predictors re-use the intermediate \(y_i\) variables at each step of the composition. For example, the logistic regression example \(\textsf{Sigmoid} \circ \textsf{Affine}_{\{A,\boldsymbol{b}\}}\) can be formulated as: \[\begin{align} \boldsymbol{z} = A\boldsymbol{x} + \boldsymbol{b} \\ y_{i} = \frac{1}{1 + e^{-z_i}}&,\quad i = 1,\ldots,P_1, \end{align}\] where \(P_1\) is the output dimension of the affine predictor.

Some formulations may introduce additional variables and constraints. For example, the complementarity or quadratic formulation of \(\textsf{ReLU}\) is: \[\begin{align} y_i = x_i + z_i,\quad i = 1,\ldots,P \\ y_i \cdot z_i \le 0,\quad i = 1,\ldots,P \\ y_i, z_i \ge 0,\quad i = 1,\ldots,P. \end{align}\]

The benefits of the full-space approach are that solvers and modeling languages can efficiently exploit the affine constraints, and that each nonlinear constraint is sparse (it uses only a small number of the total decision variables) and involves few nonlinear terms. The downside of the full-space approach is that we add intermediate variables and constraints. For deep (or wide) neural networks in particular, this can result in a large number of additional variables and constraints in the optimization model. Because of its simplicity, the full-space formulation is supported by all software packages that we compare in Section 4.

2.2 Reduced-space↩︎

In the reduced-space formulation we represent each predictor as an expression, and we delay adding intermediate variables until necessary. For example, the logistic regression example is equivalent to the expression: \[y_{i} \coloneq \frac{1}{1 + e^{-\boldsymbol{e}_i^\top(A\boldsymbol{x} + \boldsymbol{b})}},\quad i = 1,\ldots,P_1.\] If \(y_i\) appears in other parts of the optimization model, then the nonlinear expression is used instead of an intermediate decision variable \(y_i\).

Note that some predictors do not have a reduced-space formulation. One example is the complementarity \(\textsf{ReLU}\) formulation. For these predictors we always use the full-space formulation.

The benefit of the reduced-space approach is that we add fewer variables and constraints to the problem. The downside is that the nonlinear expressions become complicated with a very large number of terms. Of the packages that we compare in Section 4, the reduced-space formulation is supported only by MathOptAI.jl and OMLT. One reason is that other software packages like gurobi-machinelearning have limited support for large nonlinear expressions.

2.3 Gray-box↩︎

In the gray-box formulation, we do not attempt to encode the predictor \(F\) algebraically. Instead, we exploit the fact that nonlinear local solvers such as Ipopt [11] require only callback oracles to evaluate the function \(F(\boldsymbol{x}): {\mathbb{R}}^N \rightarrow \mathbb{R}^P\), the Jacobian \(\nabla F(\boldsymbol{x}): {\mathbb{R}}^N \rightarrow \mathbb{R}^{P\times N}\) and, optionally, the Hessian-of-the-Lagrangian \(\nabla^2\mathcal{L}(\boldsymbol{x}, \lambda): {\mathbb{R}}^{N+P} \rightarrow \mathbb{R}^{N\times N}\), which is defined as: \[\nabla^2\mathcal{L}(\boldsymbol{x}, \lambda) = \sum\limits_{i=1}^P \lambda_i \nabla^2 F_i(\boldsymbol{x}),\] where \(\lambda \in \mathbb{R}^P\) is a vector of Lagrange multipliers that varies between iterations of the solution algorithm.

Using JuMP’s support for user-defined nonlinear operators, we implement the function evaluation as a nonlinear operator \(F(\boldsymbol{x})\), and we use automatic differentiation to compute the Jacobian, \(\nabla F(\boldsymbol{x})\), and Hessian-of-the-Lagrangian, \(\nabla^2\mathcal{L}(\boldsymbol{x}, \lambda)\). Then, we add \(F\) using the full-space formulation: \[F(\boldsymbol{x}) - \boldsymbol{y} = \boldsymbol{0},\] or the reduced-space formulation \(\boldsymbol{y} \coloneq F(\boldsymbol{x})\).

There are three main benefits to the gray-box approach. First, the number of variables and constraints in the optimization model scales with the dimension of the input vector \(\boldsymbol{x}\) and output vector \(\boldsymbol{y}\), and not with the size or complexity of the intermediate predictors because these are not represented explicitly in the optimization model. As shown in our related work [10], this enables the gray-box formulation to scale and embed machine learning predictors with hundreds of millions of parameters. Second, the external evaluation of the oracles means we can trivially use customized tools, such as PyTorch’s GPU acceleration, to improve the evaluation performance. We again direct readers to [10] for more details. Third, the gray-box formulation enables models which are too complicated (or tedious) to express algebraically in the full-space or reduced-space formulations. This includes, for example, user-defined custom layers in PyTorch.

The downside of the gray-box approach is that the explicit representation of the constraints are not exposed to the solver, which means that the gray-box formulation cannot be used by nonlinear global optimization solvers. Moreover, some predictors do not have a gray-box formulation because they are not amenable to oracle-based nonlinear optimization. One example is the BinaryDecisionTree predictor, which is discrete and non-differentiable.

Of the packages that we compare in Section 4, the gray-box formulation is unique to MathOptAI.jl. The gray-box formulation was previously suggested by [8] and [9], but they developed their work as a custom implementation and not as part of a wider general purpose library.

2.3.1 Hessian-of-the-Lagrangian↩︎

Naively computing the Hessian-of-the-Lagrangian for the gray-box formulation requires computing \(P\) many \(N\times N\) Hessian matrices \(\nabla^2 F_i(\boldsymbol{x},\lambda)\), and then computing a sum-product over the \(\lambda\) vector (in \(\mathbb{R}^P\)). We improve performance by creating the Lagrangian function as a new predictor: \[G(\boldsymbol{x}, \lambda) = (\textsf{Affine}_{\{\lambda,\boldsymbol{0}\}} \circ F)(\boldsymbol{x}),\] where \(\lambda\) is treated as a \(1\times P\) matrix, then differentiating \(G\) directly using the automatic differentiation frameworks of PyTorch [12] or Flux.jl [13]. Integrating with machine learning frameworks to evaluate the Hessian-of-the-Lagrangian allows us to avoid the sum-product and memory requirement to represent the explicit Hessians, and to exploit GPU acceleration, if available, when computing \(\nabla^2 G(\boldsymbol{x},\lambda)\).

2.4 Relative performance↩︎

The relative performance of each formulation depends on the problem instance, solution algorithm, and machine hardware. For example, the full-space formulation may perform better on some instances due to the advantages of “lifting” discussed by [14]. In contrast, we have observed that reduced-space formulations sometimes require significantly fewer iterations to converge [10]. While gray-box formulations can exploit GPU acceleration, we note that similar speedups may be possible for full-space formulations by exploiting the structure of the machine learning model in the solver’s linear algebra operations. We leave a numerical comparison of the different approaches to future work.

2.5 Variable bounds↩︎

In addition to adding the variables and constraints required by each predictor, we also add bounds to the new variables by propagating bounds from the input \(\boldsymbol{x}\) variables (when they exist) through the predictors by direct substitution. Although theoretically redundant, adding these bounds often improves the performance of the optimization solver.

3 Code example↩︎

As a trivial example of the usage of MathOptAI.jl, we consider embedding the following neural network into a JuMP model. It has 10 input variables, two output variables, and a \(\textsf{ReLU}\) activation function.

First, we build and train (code not shown) the model in PyTorch [12] . Then we save the trained model to disk using torch.save:

#!/usr/bin/python3
import torch
from torch import nn
model = nn.Sequential(nn.Linear(10, 16), nn.ReLU(), nn.Linear(16, 2))
torch.save(model, "model.pt")

We can then embed the network from PyTorch into a JuMP model:

#!/usr/bin/julia
using JuMP, MathOptAI, PythonCall, Ipopt
model = Model(Ipopt.Optimizer)
@variable(model, 0 <= x[1:10] <= 1)
predictor = MathOptAI.PytorchModel("model.pt")
y, formulation = MathOptAI.add_predictor(model, predictor, x)

These lines: load the relevant Julia packages; create a JuMP model with Ipopt [11] as the optimizer; define a vector of 10 decision variables; reference the PyTorch file on disk; and add the predictor using the full-space formulation. The return y is the output vector of decision variables. The formulation object contains the variables and constraints added by the predictor. To add the predictor using the reduced-space formulation, we can do:

y, _ = MathOptAI.add_predictor(model, predictor, x; reduced_space = true)

To add the predictor using the gray-box formulation, we can do:

y, _ = MathOptAI.add_predictor(model, predictor, x; gray_box = true)

The gray-box formulation has additional options. For example, here we evaluate the oracles on the GPU (device = "cuda") and we do not provide Hessians to Ipopt (hessian = false):

y, _ = MathOptAI.add_predictor(model, predictor, x; gray_box = true,
                               device = "cuda", hessian = false)

By default, MathOptAI.jl uses the non-smooth nonlinear formulation of \(\textsf{ReLU}\). To support a mixed-integer linear solver like HiGHS [15], we can specify a different reformulation:

#!/usr/bin/julia
using JuMP, MathOptAI, PythonCall, HiGHS
model = Model(HiGHS.Optimizer)
@variable(model, 0 <= x[1:10] <= 1)
predictor = MathOptAI.PytorchModel("model.pt")
config = Dict(:ReLU => MathOptAI.ReLUSOS1)
y, _ = MathOptAI.add_predictor(model, predictor, x; config)

Here config is a dictionary which maps the name of a PyTorch layer to an AbstractPredictor constructor for MathOptAI.jl to use when formulating the predictor. The ReLUSOS1 formulation uses the SOS1 formulation of [16]. Note that HiGHS does not support SOS1 constraints natively, but, when the input variables are bounded, JuMP automatically reformulates SOS1 constraints into a MIP equivalent.

MathOptAI.jl supports a variety of machine learning predictors that are not neural networks. For example, to train and embed a decision tree from DecisionTree.jl [17], we can do:

#!/usr/bin/julia
using JuMP, MathOptAI, DecisionTree, HiGHS
truth(x::Vector) = x[1] <= 0.5 ? -2 : (x[2] <= 0.3 ? 3 : 4)
features = abs.(sin.((1:10) .* (3:4)'));
labels = truth.(Vector.(eachrow(features)));
predictor = DecisionTree.build_tree(labels, features)
model = Model(HiGHS.Optimizer);
@variable(model, 0 <= x[1:2] <= 1);
y, _ = MathOptAI.add_predictor(model, predictor, x);

4 Comparison to related work↩︎

Table 1 compares the features of MathOptAI.jl to OMLT[5], gurobi-machinelearning[4], PySCIPOpt-ML[16], and GAMSPy[18]. We choose these packages because they are the most prominent and well developed in the ecosystem. Other alternatives exist—see [1] for a comparison—but they are significantly less developed or feature-diverse than these tools. Note that the table is not intended to be a comprehensive comparison of features, for example, it does not cover alternative formulations for each predictor, supported file types for input and output, or the supported higher level interfaces to packages such as pytorch, scikit-learn, and XGBoost.

All packages are open-source. MathOptAI.jl is implemented in Julia, while OMLT, gurobi-machinelearning, PySCIPOpt-ML, and GAMSPy are implemented in Python. Note that, while gurobi-machinelearning is open-source, it requires a commercial license for the Gurobi optimizer. Similarly, while GAMSPy is MIT licensed, it requires a commercial license for the GAMS modeling language. MathOptAI.jl, OMLT, and GAMSPy target the solver-independent modeling languages JuMP [3], Pyomo [19], and GAMS, while gurobi-machinelearning and PySCIPOpt-ML are solver-specific implementations for Gurobi [20] and SCIP [21] respectively.

GAMSPy and OMLT support a wide variety of features and machine learning predictors. For example, GAMSPy and OMLT support support convolution and pooling predictors, as well as a wide range of activation functions. OMLT uniquely supports graph neural networks (GNNs). Because gurobi-machinelearning is specific to the Gurobi optimizer, it has a smaller set of supported predictors (it is particularly missing smooth nonlinear predictors), and it implements only the full-space formulation. gurobi-machinelearning does, however, have strong support for gradient boosted trees and random forests, which coincide with Gurobi’s strength as a mixed-integer optimizer. PySCIPOpt-ML is similar to gurobi-machinelearning, except that it uses SCIP’s support for global nonlinear optimization to support nonlinear neural network layers.

MathOptAI.jl uniquely has support for embedding a predictor via the gray-box formulation. As a corollary, MathOptAI.jl is the only package to support GPU acceleration of the predictors via this formulation. In theory, the gray-box feature could be added to OMLT, since the work of [9] shows how to embed a PyTorch model into Pyomo. We argue that the results of [10] and [9] demonstrate that the gray-box formulation is a worthwhile feature addition for OMLT. The gray-box feature cannot be added to gurobi-machinelearning or PySCIPOpt-ML because they require the algebraic form of the constraints for their global solvers.

Even though OMLT can be used to build JuMP models, the user experience of OMLT+JuMP has more friction than MathOptAI.jl+JuMP because the former forces the user to install and manage both Julia and Python, and there are some nuances for when and how the OMLT block and JuMP model interact. However, OMLT’s JuMP interface is a good option for users looking to embed predictors that are supported by OMLT and not by MathOptAI.jl, such as a SAGEConv layer in a graph neural network.

5 Design principles↩︎

When designing MathOptAI.jl, we had the benefit of being able to observe existing implementations such as OMLT[5]. From this analysis, and our prior experience developing other Julia packages, we developed the following design principles.

5.1 Leverage Python’s strengths↩︎

Julia has a strong optimization ecosystem, particularly, for our interest, around JuMP. However, Julia’s machine learning ecosystem—while feature rich and well developed—is significantly less popular than Python’s machine learning ecosystem. Thus, we support both Julia-based machine learning libraries such as Lux.jl [22] and Flux.jl [13], and Python-based machine learning libraries such as PyTorch [12]. Our experience has shown [10] that the Julia-Python interface was both simple to engineer and performant to use in practice. OMLT’s similarly positive experience integrating Julia and Python libraries suggests that the future may see more integration between Python and Julia, with libraries in each language able to leverage the strengths of the other.

5.2 Leverage Julia’s strengths↩︎

The design of MathOptAI.jl leverages Julia’s strengths by relying heavily on multiple dispatch. Inside MathOptAI.jl, we define new predictors as a subtype of AbstractPredictor. The subtypes implement a common MathOptAI.add_predictor function. For example, the basic code to implement a non-smooth nonlinear \(\textsf{ReLU}\) predictor is:

struct ReLU <: MathOptAI.AbstractPredictor end
function MathOptAI.add_predictor(model::JuMP.Model, predictor::ReLU, x::Vector)
    y = JuMP.@variable(model, [1:length(x)])
    cons = JuMP.@constraint(model, y .== max.(0, x))
    return y, MathOptAI.Formulation(predictor, y, cons)
end
MathOptAI.add_predictor(::JuMP.Model, p::MathOptAI.ReducedSpace{ReLU}, x::Vector) =
    (max.(0, x), MathOptAI.Formulation(p))

A second strength of Julia is its package extension mechanism, which allows code to be dynamically loaded if two other packages are loaded. We use this feature to implement package extensions between MathOptAI.jl and third-party Julia packages. As one example, the code to implement the interface between the linear modeling package GLM.jl [23] and MathOptAI.jl is:

import MathOptAI, GLM
MathOptAI.add_predictor(model::JuMP.Model, p::GLM.LinearModel, x::Vector) =
    MathOptAI.add_predictor(model, MathOptAI.Affine(GLM.coef(p)), x)

The result is a clean and simple implementation, where we have a set of predictors that subtype from AbstractPredictor and implement a common API to convert their data structures into a set of JuMP variables and constraints, and a set of package extensions that convert third-party machine learning models into the AbstractPredictors supported by MathOptAI.jl.

5.3 Composition of predictors↩︎

Packages such as OMLT, gurobi-machinelearning, and PySCIPOpt-ML approach neural networks with the assumption that basic building block is a layer: \[\boldsymbol{y} = \sigma(W\boldsymbol{x} + \boldsymbol{b}),\] which combines both an affine transform and an activation function.

MathOptAI.jl follows the PyTorch [12] convention in which everything is a predictor \(\boldsymbol{y} = F(\boldsymbol{x}): \mathbb{R}^N \rightarrow \mathbb{R}^P\), including linear layers and nonlinear activation functions. A layer in MathOptAI.jl is therefore the composition of two predictors: \[\boldsymbol{y} = (\sigma \circ \textsf{Affine}_{(W, \boldsymbol{b})})(\boldsymbol{x}).\] This approach significantly simplifies the implementation of the library because we can reuse and recombine the various predictors into different pipelines.

The complete code to implement predictor composition is the following:

struct Pipeline <: MathOptAI.AbstractPredictor
    layers::Vector{MathOptAI.AbstractPredictor}
end
function MathOptAI.add_predictor(model::JuMP.Model, predictor::Pipeline, x::Vector)
    formulation = MathOptAI.PipelineFormulation(predictor, Any[])
    for layer in predictor.layers
        x, inner_formulation = MathOptAI.add_predictor(model, layer, x)
        push!(formulation.layers, inner_formulation)
    end
    return x, formulation
end

which enables syntax like:

A, b = rand(2, 3), rand(2)
logistic = MathOptAI.Pipeline([MathOptAI.Affine(A, b), MathOptAI.Sigmoid()])

Note how Pipeline is itself an AbstractPredictor, which enables nested pipelines.

5.4 Inputs and outputs are vectors↩︎

A key feature of MathOptAI.jl are the types of inputs and outputs that are supported. One approach would be to allow the input and output types to depend on the predictor. Thus, an \(\textsf{Affine}\) predictor might input and output n-dimensional arrays, while a binary classifier predictor might output a scalar Boolean. Unfortunately, such a design makes the predictor-composition principle difficult to implement. Instead, MathOptAI.jl requires that all predictors accept only a vector x::Vector as input and return a vector y::Vector as output, even if the input and output is a scalar.

The biggest limitation of this approach is that MathOptAI.jl does not natively support \(n\)-dimensional array inputs or outputs. One common example is an image input, in which the natural input type is a three-dimensional array for the row, column, and color channel. Our solution is to force users to reshape their input and output data into a one-dimensional vector. For example, a \((M, N, C)\) dimensional image can be reshaped into a vector of length \(M\times N \times C\). The corresponding input size information of \((M, N, C)\) is stored as an attribute of the predictor so that we can recover the original array when needed. This decision is a significant trade-off, but it greatly simplified the implementation of MathOptAI.jl and we have not found it to be a burden in practice.

6 Conclusions↩︎

This paper introduced MathOptAI.jl, an open-source Julia library for embedding trained machine learning predictors into a JuMP model. A novel feature of MathOptAI.jl is its support for the gray-box formulation. Readers are directed to [7] for the archived source code and to https://github.com/lanl-ansi/MathOptAI.jl for the latest version.

The roadmap for future development of MathOptAI.jl includes a way to write decomposition algorithms and custom linear system solvers that leverage a priori knowledge about the structure of the machine learning predictors. We are also interested in adding support for new package extensions and predictors; however, the space of predictors that we could include is much larger than the set of predictors we can afford to maintain and test in an open-source package, so care is needed to avoid feature bloat.

References↩︎