Unrolled Neural Networks for Constrained Optimization

The gist In this paper, we develop unrolled neural networks to solve constrained optimization problems, offering accelerated, learnable counterparts to dual ascent (DA) algorithms.

Constrained Dual Unrolling Framework

The framework introduced is termed constrained dual unrolling (CDU), which 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. This approach departs from standard unrolling by imposing primal-descent and dual-ascent constraints through constrained learning.

Training Procedure

The training is formulated as a nested optimization problem where the inner level trains the primal network over a joint distribution of constrained problems and Lagrangian multipliers, and the outer level trains the dual network over the distribution of problem instances. To address uncertainty in the multiplier distribution required for primal network training, an alternating scheme is proposed where the dual network first generates multipliers for training the primal network, whose solutions are then used to update the dual network. This alternating procedure ensures that each network is optimized against the recent output of the other.

Key Innovations and Dynamics

The authors depart from traditional unrolling literature in two key aspects. First, instead of hard-coding the structure of an iterative solver, they use off-the-shelf neural networks to model the primal and dual iterates, allowing for more expressive and flexible architectures. Second, they enforce the optimization dynamics by imposing descent and ascent constraints on the outputs of the unrolled layers during training. The primal network predicts the minimizer of the Lagrangian function with respect to its first argument for a given multiplier. The dual network approximates the maximizer of the dual function g, which depends on λ directly and indirectly through x∗(λ).

Performance and Generalization

The numerical experiments show that imposing these descent constraints improves the final-layer performance and enhances the out-of-distribution (OOD) generalization. In both cases, our approach yields near-optimal near-feasible solutions and exhibits strong out-of-distribution (OOD) generalization. The constrained model achieves performance comparable to that of the classical DA algorithm in just 14 unrolled layers, which is significantly fewer than the 600 iterations required by DA. Furthermore, our constrained model consistently outperforms its unconstrained counterpart in optimality and feasibility under all OOD scenarios.

Applications

The framework is agnostic to the primal and dual architectures and the specific constrained problem being solved. In experiments, we tackle mixed-integer quadratic programs (MIQP) and power allocation in wireless networks, using graph neural networks (GNNs). GNNs are particularly well-suited for these tasks as they naturally capture the relational structure between variables and constraints. The final solution to (1), x∗ ∈ x∗(λ∗), can be recovered by feeding the final dual estimate, λL = ΦD z; θD, θP, to the primal network.

Conclusion

Descent- and ascent-driven unrolling offers a principled path to neural networks that solve optimization problems reliably, even far beyond their training distribution. The imposed constraints actively steer the networks toward solutions that are optimal and feasible. This acceleration, however, comes at the expense of a computationally intensive training phase. Once trained, the unrolled networks can be deployed to solve new instances with a fixed, small number of forward passes. The framework demonstrates that the performance gap widens as the constraint-to-variable ratio increases, demonstrating the effectiveness of the inductive biases introduced by descent constraints. This suggests that jointly learning the dual and primal variables can lead to solutions that match–or even exceed–in-distribution performance. The constrained dual unrolling outperforms both its unconstrained counterpart and the naive GNN across all OOD settings.

How it works

The primal network, denoted by ΦP(λ, z; θP), predicts the minimizer of the Lagrangian function with respect to its first argument for a given multiplier λ and a problem instance z. The network consists of a cascade of K layers whose outputs constitute a trajectory from an initial point xe0 to a minimizer of (5), i.e., xeK ≈ x∗(λ). The dual network, denoted by ΦD(z; θD, θP), approximates the maximizer of the dual function g(λ; x∗(λ), z). The l-th dual layer is defined as λl = Φl D λl−1, ΦP(λl−1, z; θP), z; θl D. A single forward pass of the dual network triggers L forward passes of the primal network.

Training Details

We formulate an unsupervised training problem as a nested optimization problem where the deterministic objectives are replaced with statistical counterparts evaluated over a distribution of problem instances. The training is performed through an alternating scheme: the dual network is updated on the outer-level objective, producing multiplier trajectories across its layers, and these multipliers are then used as training data for the primal network in the inner problem. Descent constraints are imposed by reformulating the inner problem as a constrained problem where E h L xek,λ; z − αkL xek−1,λ; z i ≤ 0, ∀k.

Experimental Setup

In our experiments, we consider a family of MIQP problems with n = 80 decision variables, m = 45 linear inequality constraints and r = 10 binary-valued variables. We use a primal network consisting of K = 14 unrolled primal layers, each composed of TP = 3 graph sub-layers. The dual network mirrors this structure, also comprising L = 14 dual layers with the same number of graph sub-layers, i.e., TD = 3. We alternate between the two networks for 400 iterations by running one epoch of primal network training (Algorithm 2) after every 15 epochs of dual network training (Algorithm 3).

Numerical Results

Figure 3 illustrates that our constrained dual unrolling approach achieves performance comparable to that of the classical DA algorithm, reaching near-optimal solutions in just 14 unrolled layers compared to 600 iterations required by DA. Figure 5 shows that our constrained model attains an MSE of 0.133 in x and a mean constraint violation of 0.049, compared to 0.454 and 0.094 for the unconstrained model. The constrained model preserves a descent behavior across the primal and dual layers, whereas the unconstrained one fails to do so. Our constrained dual unrolling outperforms both its unconstrained counterpart and the naive GNN across all OOD settings. The performance gap widens as the constraint level increases or when the number of agents grows, which leads to higher interference as the average inter-agent distance decreases <ref:2601.

1. Bold header: Constrained Dual Unrolling for Accelerated Optimization

This system can solve constrained optimization problems by jointly training a primal network to find a stationary point of the Lagrangian and a dual network to maximize the dual function, finding the saddle point of (4)–(5).

2. Bold header: Imposition of Optimization Dynamics

The framework enforces optimization dynamics by imposing descent and ascent constraints on the outputs of the unrolled layers during training, which leads to trajectories that are predominantly monotone with only mild oscillations.

3. Bold header: Robustness to Out-of-Distribution (OOD) Generalization

The constrained models exhibit strong out-of-distribution (OOD) generalization and consistently outperform unconstrained models, as the descent constraints steer the networks toward solutions that are optimal and feasible even when problem parameters shift.

4. Bold header: Application in Mixed-Integer Quadratic Programs (MIQP)

The system can solve MIQPs by using Graph Neural Networks (GNNs) to model the problem structure, predicting a saddle point of the Lagrangian function for any instance drawn from the problem family.

5. Bold header: Real-time Power Allocation in Wireless Networks

The system can optimize power allocation by maximizing sum rate subject to minimum rate constraints, achieving comparable performance to that of dual ascent in much fewer steps compared to iterative methods like SA + Dual Ascent.

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. Unlike standard unrolling, which maps the updates of a prescribed iterative solver directly to network layers, CDU preserves the underlying optimization dynamics while learning the update mappings. The desired primal-dual dynamics are therefore induced through constrained learning by enforcing primal-descent and dual-ascent constraints. 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 establish theoretical guarantees on the feasibility and optimality of the coupled networks. We numerically evaluate the framework on quadratic programs (QPs) with linear constraints and power allocation in wireless networks.