Formulations and scalability of neural network surrogates in nonlinear optimization problems

We study nonlinear optimization problems in the form given by Equation (1):

We consider interior point methods, such as IPOPT [22], for solving (1). These methods iteratively compute search directions $d$ by solving the linear system (2):

where $\alpha$ and $\beta$ are additional terms that are not shown for simplicity. The matrix on the left-hand side is referred to as the Karush-Kuhn Tucker, or KKT, matrix. To construct this system, solvers rely on callbacks that provide the Jacobian $\nabla g$ and, optionally, the Hessian of the Lagrangian function, $\nabla^{2}\mathcal{L}$. These are typically provided by the automatic differentiation system of an algebraic modeling environment. If the Hessian $\nabla^{2}\mathcal{L}$ is not available, a limited-memory quasi-Newton approximation is used [17].

A neural network predictor is a function denoted $y={\rm NN}(x)$. We consider neural networks defined by repeated application of an affine transformation and a nonlinear activation function $\sigma$ over $L$ layers:

where $y_{0}=x$ and $y=y_{L}$. Weights $W_{l}$ and biases $b_{l}$ are parameters that are optimized to minimize error on a set of training data representing desired inputs and outputs of the neural network. To a nonlinear optimization solver using a trained neural network, $W_{l}$ and $b_{l}$ are constant. To fit assumptions made by these solvers, we consider only smooth activation functions, e.g., sigmoid and hyperbolic tangent functions.

In this section, we explain the three ways in which we encode pre-trained neural network predictors into the constraints of a nonlinear optimization model of the form (1). The three approaches are denominated full-space, reduced-space, and gray-box.

Encoding a trained neural network into an optimization model using the forms described in Section 2.3 is tedious and error-prone. To simplify experimentation, we developed a new Julia package, MathOptAI.jl, which is a JuMP extension for embedding a range of machine learning models into a JuMP model. In addition to supporting neural networks trained using PyTorch, which are the focus of this paper, MathOptAI.jl also supports Julia-based deep-learning libraries, as well as other machine learning models such as decision trees and Gaussian Processes. MathOptAI.jl is provided as an open-source package at https://github.com/lanl-ansi/MathOptAI.jl under a BSD-3 license.

We model the SCOPF problem using PowerModels [9], PowerModelsSecurityConstrained [8], and JuMP [15]. Neural networks are modeled using PyTorch [18] and embedded into the optimization problem using MathOptAI.jl. Optimization problems are solved using the IPOPT nonlinear optimization solver [22] with MA27 [10] as the linear solver. The full-space and reduced-space models support evaluation on a CPU but not on a GPU. Because gray-box models use PyTorch, they can be evaluated on a CPU or GPU. We run our experiments on the Venado supercomputer. CPU-only experiments use compute nodes with two 3.4 GHz NVIDIA Grace CPUs and 240 GB of RAM, while CPU+GPU experiments use nodes with a Grace CPU with 120 GB of RAM and an NVIDIA H100 GPU.

Table 1 shows the numbers of variables, constraints, and nonzero entries of the derivative matrices for the optimization problem with different neural networks and formulations. We note that the reduced-space and gray-box formulations have approximately the same numbers of constraints and variables as the original problem, but more nonzero entries in the Jacobian and Hessian matrices due to the dense, nonlinear stability constraints. With these formulations, the structure of the optimization problem does not change as the neural network surrogate adds more interior layers. By contrast, the full-space formulation grows in numbers of variables, constraints, and nonzeros as the neural network gets larger. These problem structures suggest that the full-space formulation will lead to expensive KKT matrix factorizations, while this will not be an issue for reduced-space and gray-box formulations.

Runtimes for the different formulations with neural network surrogates of increasing size are given in Table 2. For gray-box formulations, we compare optimization solves using exact and approximate Hessian evaluations and different hardware platforms. The results immediately show that full-space and reduced-space formulations are not scalable to neural networks with more than a few million trained parameters. The full-space formulation fails with a segmentation fault—likely due to memory issues in MA27—while the reduced-space formulation exceeds time and memory limits building the constraint expressions in JuMP. A breakdown of solve times, given in Table 3, confirms the bottlenecks in these formulations. The full-space formulation spends almost all of its solve time in the IPOPT algorithm, which we assume is dominated by KKT matrix factorization, while the reduced-space formulation spends most of its solve time evaluating the Hessian.

By contrast, the gray-box formulation is capable of solving the optimization problem with the largest neural network surrogates tested. While a CPU-only solve with exact Hessian matrices takes an unacceptably-long 748 s, a GPU-accelerated solve with approximate Hessian matrices solves in only one second. This is slower than the original SCOPF problem (with no stability constraint) by a factor of 2.5, which may be acceptable for some applications. In all cases, the solve time with the gray-box formulation is dominated by function and Hessian evaluations, which explains the large speed-ups obtained with the GPU (50$\times$ and 17$\times$ for the 592M-parameter network with exact and approximate Hessians).

Approximating the Hessian matrix also speeds up the solve significantly. Hessian approximation is not a common approach when exact Hessians are available because it can lead to slow and unreliable convergence. In the 592M-parameter case, approximating the Hessian increases the iteration count by 14, but it makes up for this by decreasing the time per iteration by a factor of 60. These results suggest that this is an appropriate trade-off for optimization problems constrained by large neural networks.