Unrolled Graph Neural Networks for Constrained Optimization

TL;DR

This work addresses constrained optimization by learning to emulate the dual ascent dynamics through two coupled GNNs that unroll primal and dual updates. A nested training framework enforces layerwise descent in the primal and ascent in the dual, enabling the models to capture DA dynamics while remaining robust to distribution shifts. Experiments on mixed-integer quadratic programs demonstrate near-optimal, near-feasible solutions and strong out-of-distribution generalization, with fast inference via a single forward pass. Overall, the approach yields scalable, robust learned solvers for constrained problems.

Abstract

In this paper, we unroll the dynamics of the dual ascent (DA) algorithm in two coupled graph neural networks (GNNs) to solve constrained optimization problems. The two networks interact with each other at the layer level to find a saddle point of the Lagrangian. The primal GNN finds a stationary point for a given dual multiplier, while the dual network iteratively refines its estimates to reach an optimal solution. We force the primal and dual networks to mirror the dynamics of the DA algorithm by imposing descent and ascent constraints. We propose a joint training scheme that alternates between updating the primal and dual networks. Our numerical experiments demonstrate that our approach yields near-optimal near-feasible solutions and generalizes well to out-of-distribution (OOD) problems.

Figures (3)

• Figure 1: The structure of the unrolled primal GNN (left) and dual GNN (right). The forward pass alternates between: i) each dual layer $\Phi_{\text{D}}^l$ outputs $\boldsymbol{\lambda}_l$, which is sent to the primal network $\Phi_{\text{P}}$ to query for the Lagrangian minimizer ${\mathbf x}^*(\boldsymbol{\lambda}_l)$, and ii) the primal network performs its internal forward pass and returns ${\mathbf x}_l \approx {\mathbf x}^*(\boldsymbol{\lambda}_l)$.
• Figure 2: Descent Guarantees. (Left) The gradient norm of the Lagrangian across the primal layers. (Middle) The constraint violation across the dual layers. (Right) The complementary slackness $\boldsymbol{\lambda}_L^\top {\mathbf f}({\mathbf x}_l)$ across the unrolled dual layers. The constrained model exhibits a consistent decrease in all three quantities across layers, whereas the unconstrained model shows a more oscillatory pattern.
• Figure 3: 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 other learning-based methods in optimality (top) and feasibility (bottom) across all OOD scenarios.