# Formulations and scalability of neural network surrogates in nonlinear optimization problems

arXiv: https://arxiv.org/abs/2412.11403 · HTML: https://mlsanity.com/paper/2412.11403/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 7 of its 20 references that are in the map.
- First posted on arXiv 16 Dec 2024.

## Builds on (most cited first, 7 of 7)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [JuMP 1.0: Recent improvements to a modeling language for mathematical optimization](https://mlsanity.com/paper/2206.03866/) (2022)
- [OMLT: Optimization & Machine Learning Toolkit](https://mlsanity.com/paper/2202.02414/) (2022)
- [Global Deterministic Optimization with Artificial Neural Networks Embedded](https://mlsanity.com/paper/1801.07114/) (2018)
- [Recent Developments in Security-Constrained AC Optimal Power Flow: Overview of Challenge 1 in the ARPA-E Grid Optimization Competition](https://mlsanity.com/paper/2206.07843/) (2022)
- [Process flowsheet optimization with surrogate and implicit formulations of a Gibbs reactor](https://mlsanity.com/paper/2310.09307/) (2023)

## Built on by (most cited first, 1 of 1)

- [A Simultaneous Approach for Training Neural Differential-Algebraic Systems of Equations](https://mlsanity.com/paper/2504.04665/) (2025)

## Abstract (arXiv)

We compare full-space, reduced-space, and gray-box formulations for representing trained neural networks in nonlinear constrained optimization problems. We test these formulations on a transient stability-constrained, security-constrained alternating current optimal power flow (SCOPF) problem where the transient stability criteria are represented by a trained neural network surrogate. Optimization problems are implemented in JuMP and trained neural networks are embedded using a new Julia package: MathOptAI.jl. To study the bottlenecks of the three formulations, we use neural networks with up to 590 million trained parameters. The full-space formulation is bottlenecked by the linear solver used by the optimization algorithm, while the reduced-space formulation is bottlenecked by the algebraic modeling environment and derivative computations. The gray-box formulation is the most scalable and is capable of solving with the largest neural networks tested. It is bottlenecked by evaluation of the neural network's outputs and their derivatives, which may be accelerated with a graphics processing unit (GPU). Leveraging the gray-box formulation and GPU acceleration, we solve our test problem with our largest neural network surrogate in 2.5× the time required for a simpler SCOPF problem without the stability constraint.