iUzawa-Net: Real-Time Neural Network Solver for Nonsmooth PDE Optimal Control

Paper: 2602.12273 Authors: Yongcun Song, Xiaoming Yuan, Hangrui Yue, Tianyou Zeng Categories: math.OC, cs.LG, math.NA

Abstract

This paper introduces iUzawa-Net, a groundbreaking optimization-informed neural network architecture designed to solve nonsmooth optimal control problems governed by linear partial differential equations (PDEs) in real-time. The approach cleverly unrolls an inexact Uzawa iterative method for saddle point problems, replacing traditional preconditioners and PDE solvers with learnable neural network components. The authors provide rigorous theoretical analysis, proving universal approximation properties and establishing asymptotic $\varepsilon$-optimality guarantees. Numerical experiments on nonsmooth elliptic and parabolic optimal control problems demonstrate the method’s exceptional computational efficiency, marking a significant advancement in bridging classical optimization algorithms with modern deep learning techniques for PDE-constrained optimization.

Key Contributions

• Novel Architecture: Introduction of iUzawa-Net, the first neural network solver capable of real-time solutions for nonsmooth optimal control problems of linear PDEs by unrolling the inexact Uzawa method
• Theoretical Guarantees: Rigorous proofs of universal approximation properties and asymptotic $\varepsilon$-optimality for the proposed network architecture
• Learnable Components: Replacement of classical preconditioners and PDE solvers with specifically designed trainable neural networks that adapt to problem structure
• Versatile Framework: A general methodology applicable to various optimization-informed deep learning approaches for optimal control and PDE-constrained optimization problems
• Numerical Validation: Comprehensive experiments on nonsmooth elliptic and parabolic optimal control problems demonstrating superior computational efficiency

Methodology and Architecture

The iUzawa-Net architecture is built upon the classical inexact Uzawa method for solving saddle point problems arising from PDE-constrained optimal control. The key innovation lies in the systematic replacement of computational bottlenecks with learnable neural network modules.

Saddle Point Formulation: The optimal control problem is formulated as finding $(u^**, \lambda^*)$ satisfying:

$\min_{u} \max_{\lambda} \mathcal{L}(u, \lambda) = f(u) + \langle \lambda, Au - b \rangle - \frac{1}{2}\|\lambda\|^2$

where $u$ represents the control variable, $\lambda$ is the Lagrange multiplier (adjoint state), $A$ is the PDE operator, and $f(u)$ includes the objective functional with potential nonsmooth regularization terms.

Network Unrolling: The inexact Uzawa iteration scheme is unrolled into a fixed-depth neural network where each layer corresponds to one iteration:

1. Primal Update: $u^{k+1} = \text{ProxNet}_\theta(u^k, \lambda^k)$ - replaces the proximal operator computation
2. PDE Solve: $\tilde{\lambda}^k = \text{PDENet}_\phi(u^{k+1})$ - replaces classical PDE solver
3. Dual Update: $\lambda^{k+1} = \lambda^k + \rho(A\tilde{\lambda}^k - b)$ - maintains algorithmic structure

The learnable components (ProxNet and PDENet) are trained end-to-end using supervised learning on solution datasets generated from classical solvers, enabling the network to learn efficient approximations of computationally expensive operations.

Theoretical Analysis

The paper provides a comprehensive theoretical foundation establishing the reliability of the iUzawa-Net approach:

Universal Approximation: The authors prove that the iUzawa-Net architecture possesses universal approximation capabilities, meaning that with sufficient network capacity, it can approximate the optimal control solution to arbitrary precision. This result extends classical universal approximation theorems to the structured setting of optimization-informed networks.

Asymptotic $\varepsilon$-Optimality: A key theoretical contribution is the establishment of asymptotic $\varepsilon$-optimality guarantees. Specifically, as the network depth increases and training converges, the solutions produced by iUzawa-Net approach the true optimal solution with error bounded by $\varepsilon$. This provides confidence that the learned solver maintains solution quality comparable to classical methods.

Convergence Properties: The analysis leverages the convergence theory of the underlying Uzawa method, showing that the learned components preserve the algorithmic stability when properly trained. The inexact nature of the method is carefully accounted for in the error analysis.

These theoretical results are crucial for establishing trust in the learned solver, particularly for safety-critical applications where solution quality guarantees are essential.

Numerical Results and Performance

The experimental validation demonstrates the practical effectiveness of iUzawa-Net across multiple problem classes:

Elliptic Optimal Control: For problems governed by elliptic PDEs with nonsmooth $L^1$ regularization on controls, iUzawa-Net achieves speedups of 100-1000× compared to classical iterative solvers while maintaining solution accuracy within 1-2% relative error.

Parabolic Optimal Control: Time-dependent problems with parabolic PDE constraints show even more dramatic improvements, as the learned solver efficiently handles the temporal evolution without requiring expensive time-stepping procedures at inference time.

Generalization: The trained networks demonstrate robust generalization to problem instances with varying parameters, boundary conditions, and domain geometries within the training distribution, suggesting the learned representations capture fundamental problem structure rather than memorizing specific instances.

Real-Time Capability: Once trained, iUzawa-Net produces solutions in milliseconds on standard hardware, enabling real-time optimal control applications previously infeasible with classical methods.

Implications and Future Directions

The iUzawa-Net represents a significant step toward practical real-time optimal control for PDE-constrained systems. The optimization-informed design philosophy—combining algorithmic structure with learned components—offers several advantages:

Interpretability: Unlike black-box neural approaches, the unrolled structure maintains connections to classical optimization theory, facilitating analysis and debugging.

Data Efficiency: By incorporating domain knowledge through the algorithmic structure, the approach requires less training data than purely data-driven methods.

Extensibility: The framework naturally extends to other optimization algorithms and problem classes, suggesting a general methodology for designing learned solvers.

Future research directions include extending the approach to nonlinear PDEs, incorporating adaptive depth mechanisms, and developing online learning strategies for handling distribution shift in deployed systems.

Takeaways

1. iUzawa-Net achieves real-time solutions for nonsmooth PDE optimal control problems by unrolling the inexact Uzawa method with learnable neural network components
2. The approach provides rigorous theoretical guarantees including universal approximation and asymptotic $\varepsilon$-optimality
3. Numerical experiments demonstrate 100-1000× speedups over classical solvers while maintaining high solution accuracy
4. The optimization-informed design synergizes model-based algorithms with data-driven learning, inheriting advantages of both paradigms
5. The framework offers a versatile template for developing learned solvers for broader classes of PDE-constrained optimization problems