Optimality-Informed Neural Networks for Solving Parametric Optimization Problems

TL;DR

This work tackles the computational burden of solving parametric nonlinear constrained optimization by learning a parameter-to-solution map. It introduces Optimality-Informed Neural Networks (OptINNs), which embed the Karush-Kuhn-Tucker (KKT) conditions into both the loss function and architecture to predict primal and dual solutions, yielding feasible and near-optimal surrogates with reduced data needs. The approach is demonstrated across linear and nonlinear problems, showing competitive primal accuracy and, crucially, improved constraint satisfaction and dual-variable estimation compared to quadratic-penalty baselines, especially for larger or more complex problems. The results suggest OptINNs provide accurate, data-efficient surrogates suitable for real-time optimization tasks and model-predictive control, with promising directions for dynamic architectures and sensitivity-enabled training.

Abstract

Many engineering tasks require solving families of nonlinear constrained optimization problems, parametrized in setting-specific variables. This is computationally demanding, particularly, if solutions have to be computed across strongly varying parameter values, e.g., in real-time control or for model-based design. Thus, we propose to learn the mapping from parameters to the primal optimal solutions and to their corresponding duals using neural networks, giving a dense estimation in contrast to gridded approaches. Our approach, Optimality-informed Neural Networks (OptINNs), combines (i) a KKT-residual loss that penalizes violations of the first-order optimality conditions under standard constraint qualifications assumptions, and (ii) problem-specific output activations that enforce simple inequality constraints (e.g., box-type/positivity) by construction. This design reduces data requirements, allows the prediction of dual variables, and improves feasibility and closeness to optimality compared to penalty-only training. Taking quadratic penalties as a baseline, since this approach has been previously proposed for the considered problem class in literature, our method simplifies hyperparameter tuning and attains tighter adherence to optimality conditions. We evaluate OptINNs on different nonlinear optimization problems ranging from low to high dimensions. On small problems, OptINNs match a quadratic-penalty baseline in primal accuracy while additionally predicting dual variables with low error. On larger problems, OptINNs achieve lower constraint violations and lower primal error compared to neural networks based on the quadratic-penalty method. These results suggest that embedding feasibility and optimality into the network architecture and loss can make learning-based surrogates more accurate, feasible, and data-efficient for parametric optimization.

Figures (10)

• Figure 1: The architecture and training of a 4-layer OptINN. The problem parameters are passed through the MLP for the given parameter vector ${\boldsymbol{p}}$. The trivialization layer includes all bound and box constraints on the output variables, giving the variables marked with $\overline{\underline{(\cdot)}}$. The output $\overline{\underline{{\boldsymbol{y}}}}\xspace\in\mathbb{R}\xspace^{n_y}$ corresponds to the concatenation of the primal variables ${\boldsymbol{x}}$ and dual variables ${\boldsymbol{\mu}}$ and ${\boldsymbol{\lambda}}$
• Figure 2: Cosine scheduling for $\alpha$ with $\underline{\alpha}=0.1$, $\overline{\alpha}=0.9$, and initialization and finalization phases of lengths $d^\mathrm{init} = d^\mathrm{final}=50$ over 200 epochs
• Figure 3: Minima of $-x + \gamma_g\mathop{\mathrm{ReLU}}\limits(x-1)^2$ under changing $\gamma_g$ for the optimization problem $\min_{x\in\mathbb{R}\xspace} -x$ s.t. $x \le 1$. The dashed and dotted lines mark the corresponding minimizer and minimum, respectively. Increasing $\gamma_g$ shifts the minimizer closer to the the true optimum $x=1$, but also increases the minimum. All suggested minimizer lie in the infeasible region $x>1$
• Figure 4: Comparison of the training results for a quadratic-penalty-method-based NN (PMNN) and two OptINN using four and no data points. Displayed are the mean (line) and the area between the minium and maximum (shaded) for all three methods over five trainings. We observe for all methods a good tracking of the primal-variables. While the OptINNs exhibits a larger spread along jumps in the dual variables, they are able to give reasonable estimations
• Figure 5: Training curves for both, OptINNs trained with 4 data points (top row) and without data points (bottom row). The left column shows the data-based validation loss (256 validation data points), the right column the KKT-loss. Each color corresponds to a training with different random seed. The validation data were not used for any decision making during training, we included them purely for evaluation. The 4 of the 5 training runs terminate since no improvement was made for 20000 epochs. For both trainings, we do not observe any overfitting, i.e. there is not any increase in the validation data loss

Theorems & Definitions (9)

• Theorem 2.1
• Definition 2.1
• Lemma 2.1
• Theorem 2.2
• Definition 3.1: Unimodal Penalty Function
• Lemma 3.1: Zero KKT-loss for KKT-points
• proof
• Remark 3.1
• Example 4.1