Nonlinear Optimization with GPU-Accelerated Neural Network Constraints

TL;DR

This work introduces a reduced-space formulation for nonlinear optimization with pretrained neural networks embedded as constraints, leveraging GPU-accelerated evaluation to scale to networks with hundreds of millions of parameters. By encoding the NN as a single nonlinear constraint and differentiating through it, the method achieves significantly faster solves and fewer iterations than a full-space approach, particularly when paired with GPU-based Hessian computations. Empirical results on MNIST adversarial generation and a transient-feasibility SCOPF surrogate demonstrate large-scale feasibility and robustness advantages, with substantial runtime reductions compared to CPU-only full-space solves. The findings highlight the practical potential of GPU-accelerated, reduced-space NLP for design and control problems that rely on deep learning surrogates, while acknowledging limitations for global optimization and data-transfer overhead considerations.

Abstract

We propose a reduced-space formulation for optimizing over trained neural networks where the network's outputs and derivatives are evaluated on a GPU. To do this, we treat the neural network as a "gray box" where intermediate variables and constraints are not exposed to the optimization solver. Compared to the full-space formulation, in which intermediate variables and constraints are exposed to the optimization solver, the reduced-space formulation leads to faster solves and fewer iterations in an interior point method. We demonstrate the benefits of this method on two optimization problems: Adversarial generation for a classifier trained on MNIST images and security-constrained optimal power flow with transient feasibility enforced using a neural network surrogate.