Unrolled Neural Networks for Constrained Optimization

TL;DR

This work introduces constrained dual unrolling (CDU), a framework that jointly trains two unrolled neural networks to solve constrained optimization by approximating the saddle point of the Lagrangian. The primal network learns a stationary point for a given multiplier, while the dual network traces multiplier trajectories toward optimality, with descent and ascent constraints enforcing DA-like dynamics. Through an alternating training scheme and problem-agnostic GNN architectures, CDU delivers near-optimal, near-feasible solutions with strong out-of-distribution generalization on MIQP relaxations and wireless power allocation, while achieving substantial online speedups (few forward passes) after offline training. The approach provides a principled, scalable pathway to learning-based solvers that preserve optimization structure and robustness in diverse settings.

Abstract

In this paper, we develop unrolled neural networks to solve constrained optimization problems, offering accelerated, learnable counterparts to dual ascent (DA) algorithms. Our framework, termed constrained dual unrolling (CDU), comprises two coupled neural networks that jointly approximate the saddle point of the Lagrangian. The primal network emulates an iterative optimizer that finds a stationary point of the Lagrangian for a given dual multiplier, sampled from an unknown distribution. The dual network generates trajectories towards the optimal multipliers across its layers while querying the primal network at each layer. Departing from standard unrolling, we induce DA dynamics by imposing primal-descent and dual-ascent constraints through constrained learning. We formulate training the two networks as a nested optimization problem and propose an alternating procedure that updates the primal and dual networks in turn, mitigating uncertainty in the multiplier distribution required for primal network training. We numerically evaluate the framework on mixed-integer quadratic programs (MIQPs) and power allocation in wireless networks. In both cases, our approach yields near-optimal near-feasible solutions and exhibits strong out-of-distribution (OOD) generalization.

Figures (9)

• Figure 1: Trajectories generated by (left) DA algorithm, (middle) constrained dual unrolling and (right) its unconstrained counterpart for a QP instance: (Top) primal trajectories toward the stationary point of the Lagrangian ${\mathcal{L}}(\cdot, \boldsymbol{\lambda}; {\mathbf z})$, and (bottom) dual trajectories maximizing the dual function $g(\cdot \, ; {\mathbf x}^*(\boldsymbol{\lambda}),{\mathbf z})$. The constrained networks generate descent trajectories emulating the DA algorithm, while the unconstrained models generate random trajectories and fail to hit the optimum of the dual function (bottom right).
• Figure 2: Constrained-optimization unrolling in \ref{['eq:primal_unrolling']}--\ref{['eq:recoverability']}. (Left) The primal network with $K=3$ unrolled layers approximates the stationary point ${\mathbf x}_l$ for a given $\boldsymbol{\lambda}_l$. (Middle) The dual network with $L=3$ layers (red) calls the primal network (blue) at each layer, steering the iterates toward $\boldsymbol{\lambda}_3 \approx \boldsymbol{\lambda}^*$. (Right) The solution ${\mathbf x}^*$ is recovered by feeding $\boldsymbol{\lambda}_3$ into the primal network.
• Figure 3: Performance of constrained dual unrolling across $14$ layers vs an iterative DA algorithm. (Left) The distance to the primal optimum ${\mathbf x}^*$, (middle) the distance to the dual optimum $\boldsymbol{\lambda}^*$, and (right) the objective function (a measure of optimality). The 14-layer outputs of our method are evenly distributed across the 600 iterations for clearer visual comparison.
• Figure 4: Descent Guarantees. (Left) The gradient norm of the Lagrangian across the primal layers, evaluated over a test dataset of problem instances and multiplier samples. (Middle) The constraint violation across the dual layers. (Right) The complementary slackness $\boldsymbol{\lambda}_L^\top \max\{\mathbf{0}, {\mathbf f}({\mathbf x}_l)\}$ across the unrolled dual layers (a measure of feasibility). The constrained model exhibits a consistent decrease in all three quantities across layers, whereas the unconstrained model shows a more oscillatory pattern.
• Figure 5: Robustness under OOD problems, varying (left) the number of optimization variables $n$, (middle) the number of linear constraints $m$, and (right) the number of integer-valued variables $r$. The red dotted line represents the in-distribution scenario: $n=80, m=45,$ and $r=10$. Our constrained dual unrolling outperforms the learning-based benchmarks in optimality (bottom row) and feasibility (middle row) across all OOD scenarios.