[Paper] Learning to Control: The iUzawa-Net for Nonsmooth Optimal Control of Linear PDEs

Source: arXiv - 2602.12273v1

Overview

The paper introduces iUzawa‑Net, a deep‑learning‑augmented solver that can produce near‑real‑time solutions for a challenging class of nonsmooth optimal‑control problems governed by linear PDEs. By “unrolling” an inexact Uzawa iteration and swapping out traditional preconditioners and PDE solvers for trainable neural modules, the authors bridge classic numerical optimization with modern data‑driven inference.

Key Contributions

• iUzawa‑Net architecture: a novel neural network that mimics the steps of an inexact Uzawa method for saddle‑point systems, with learnable components replacing costly linear solvers.
• Theoretical guarantees: proofs of universal approximation (the network can represent any admissible solution map) and asymptotic ε‑optimality (the learned solver approaches the true optimal control as training data grows).
• Real‑time performance: empirical results show orders‑of‑magnitude speed‑ups on nonsmooth elliptic and parabolic optimal‑control benchmarks compared with conventional iterative solvers.
• Generalizable framework: the unrolling‑and‑learning recipe is presented as a template for other PDE‑constrained optimization tasks (e.g., inverse problems, model‑predictive control).

Methodology

1. Problem formulation – The authors consider optimal‑control problems where the control enters a linear PDE and the objective includes a nonsmooth regularizer (e.g., ℓ₁‑type sparsity). This leads to a saddle‑point KKT system that is traditionally solved by Uzawa‑type iterations.

2. Inexact Uzawa unrolling – Each iteration of the classical method consists of three steps: (a) update the primal state, (b) update the dual multiplier, and (c) apply a preconditioner to solve a linear subproblem. The authors “unroll” a fixed number T of these steps into a feed‑forward network.

3. Learnable replacements –

   • State update: a convolutional encoder‑decoder that approximates the PDE solve.
   • Multiplier update: a shallow MLP that implements the dual ascent with a learned step size.
   • Preconditioner: a small residual network that mimics a multigrid‑style smoother.

4. Training – Synthetic optimal‑control pairs are generated offline by a high‑accuracy solver. The network is trained end‑to‑end to minimize the discrepancy between its output and the ground‑truth optimal control, using standard Adam optimization and a mixed loss (data fidelity + regularizer).

5. Inference – Once trained, a forward pass through iUzawa‑Net yields an approximate optimal control in milliseconds, independent of mesh size or PDE discretization.

Results & Findings

• Speed: iUzawa‑Net is ~200–300× faster than the fully converged Uzawa solver while staying within a few percent of the optimal objective value.
• Scalability: Because the network operates on the discretized fields directly, inference time does not grow with problem size, making it suitable for high‑resolution simulations.
• Robustness: The learned preconditioner adapts to different PDE coefficients and boundary conditions without re‑tuning.

Practical Implications

• Real‑time model‑predictive control (MPC) – Engineers can embed iUzawa‑Net in embedded controllers for fluid flow, heat diffusion, or structural vibration, where solving an optimal‑control problem on the fly was previously infeasible.
• Rapid design iteration – In computational design pipelines (e.g., topology optimization with sparsity constraints), the network can provide instant feedback on control updates, dramatically shortening the design loop.
• Edge deployment – The lightweight convolutional modules fit on modern GPUs or even specialized AI accelerators, opening the door to on‑device PDE‑constrained decision making (e.g., smart HVAC, autonomous underwater vehicles).
• Hybrid solvers – iUzawa‑Net can serve as a high‑quality initializer for traditional solvers, cutting down the number of required iterations and thus reducing overall compute cost.

Limitations & Future Work

• Training data dependence – The network’s performance hinges on the representativeness of the synthetic training set; out‑of‑distribution PDE parameters may degrade accuracy.
• Theoretical gap for highly nonlinear PDEs – The current analysis assumes linear PDE operators; extending the approach to fully nonlinear dynamics (e.g., Navier‑Stokes) remains open.
• Interpretability of learned preconditioners – While empirically effective, the internal workings of the neural preconditioner are not yet understood in classical numerical‑analysis terms.

Future research directions suggested by the authors include: (1) adaptive online fine‑tuning to handle parameter drift, (2) extending the unrolling framework to stochastic or multi‑objective control settings, and (3) integrating physics‑informed regularization to improve robustness across broader PDE families.

Authors

• Yongcun Song
• Xiaoming Yuan
• Hangrui Yue
• Tianyou Zeng

Paper Information

• arXiv ID: 2602.12273v1
• Categories: math.OC, cs.LG, math.NA
• Published: February 12, 2026
• PDF: Download PDF