An adaptive augmented Lagrangian method for training physics and equality constrained artificial neural networks

TL;DR

The paper tackles constrained training of physics-informed neural networks (PECANN) by replacing a single global penalty in the augmented Lagrangian with adaptive, per-constraint penalties. It introduces a RMSprop-inspired dual update scheme and an expected-constraint formulation to enable mini-batch training, improving stability and convergence for diverse PDE constraints. The approach is validated on forward and inverse problems, including heat diffusion in composite media, a 1D wave equation, and Navier–Stokes flow up to $Re=1000$, showing improved accuracy and robustness over standard PINN methods. This yields a scalable, plug-and-play framework for PDE-constrained learning with reduced hyperparameter tuning and better handling of large constraint sets in practical simulations.

Abstract

Physics and equality constrained artificial neural networks (PECANN) are grounded in methods of constrained optimization to properly constrain the solution of partial differential equations (PDEs) with their boundary and initial conditions and any high-fidelity data that may be available. To this end, adoption of the augmented Lagrangian method within the PECANN framework is paramount for learning the solution of PDEs without manually balancing the individual loss terms in the objective function used for determining the parameters of the neural network. Generally speaking, ALM combines the merits of the penalty and Lagrange multiplier methods while avoiding the ill conditioning and convergence issues associated singly with these methods . In the present work, we apply our PECANN framework to solve forward and inverse problems that have an expanded and diverse set of constraints. We show that ALM with its conventional formulation to update its penalty parameter and Lagrange multipliers stalls for such challenging problems. To address this issue, we propose an adaptive ALM in which each constraint is assigned a unique penalty parameter that evolve adaptively according to a rule inspired by the adaptive subgradient method. Additionally, we revise our PECANN formulation for improved computational efficiency and savings which allows for mini-batch training. We demonstrate the efficacy of our proposed approach by solving several forward and PDE-constrained inverse problems with noisy data, including simulation of incompressible fluid flows with a primitive-variables formulation of the Navier-Stokes equations up to a Reynolds number of 1000.

Figures (19)

• Figure 1: Effects of augmented Lagrangian penalty update strategies on the heat transfer in a composite medium problem using the point-wise constraint formulation (i.e. section \ref{['sec:constrained_optimization_formulation_pointwise_constraint']}). (a) predicted temperature at t = 1, (b) predicted heat flux at t = 1
• Figure 2: Distribution of Lagrange multipliers arising from using the formulation with point-wise constraints (i.e. section \ref{['sec:constrained_optimization_formulation_pointwise_constraint']}). (a) ALM with monotonically increasing penalty update (Algorithm 1). (b) ALM with conditional penalty update (Algorithm 2). (c) ALM with adaptive penalty update (Algorithm 3, proposed).
• Figure 3: Heat transfer in a composite medium using the adaptive penalty update strategy and the formulation based expectation of constraints (i.e. section \ref{['sec:proposed_formulation']}). (a) predicted temperature (b) absolute point-wise error in predicted temperature , (c) predicted heat flux, (d) absolute point-wise error in heat flux prediction.
• Figure 4: Error levels during training using ALM with different penalty update strategy. Top row: Temperature (a) relative error (b) absolute error. Bottom row: Heat flux (c) relative error, (d) absolute error.
• Figure 5: Evolution of penalty parameters with different penalty update strategy: (a) evolution of the penalty parameter $\mu_{\mathcal{B}}$ for enforcing boundary conditions, $\mu_{\mathcal{I}}$ for enforcing initial conditions, and $\mu_{\mathcal{Q}}$ for enforcing flux constraints using ALM with adaptive penalty update strategy . (b) evolution of penalty parameters for enforcing boundary condition, initial condition and flux equation using ALM with MPU and CPU. (c) evolution of constraints during training using ALM with APU