[Paper Review] Self-Adaptive Physics-Informed Neural Networks using a Soft Attention Mechanism

Levi D. McClenny, Ulisses Braga-Neto|arXiv (Cornell University)|Sep 7, 2020

Model Reduction and Neural Networks26 references42 citations

TL;DR

The paper introduces Self-Adaptive PINNs (SA-PINNs) that learn trainable per-point weights to create a soft attention mask, maximizing emphasis on difficult regions while minimizing overall loss to solve stiff PDEs more accurately with fewer epochs.

ABSTRACT

Physics-Informed Neural Networks (PINNs) have emerged recently as a promising application of deep neural networks to the numerical solution of nonlinear partial differential equations (PDEs). However, it has been recognized that adaptive procedures are needed to force the neural network to fit accurately the stubborn spots in the solution of "stiff" PDEs. In this paper, we propose a fundamentally new way to train PINNs adaptively, where the adaptation weights are fully trainable and applied to each training point individually, so the neural network learns autonomously which regions of the solution are difficult and is forced to focus on them. The self-adaptation weights specify a soft multiplicative soft attention mask, which is reminiscent of similar mechanisms used in computer vision. The basic idea behind these SA-PINNs is to make the weights increase as the corresponding losses increase, which is accomplished by training the network to simultaneously minimize the losses and maximize the weights. In addition, we show how to build a continuous map of self-adaptive weights using Gaussian Process regression, which allows the use of stochastic gradient descent in problems where conventional gradient descent is not enough to produce accurate solutions. Finally, we derive the Neural Tangent Kernel matrix for SA-PINNs and use it to obtain a heuristic understanding of the effect of the self-adaptive weights on the dynamics of training in the limiting case of infinitely-wide PINNs, which suggests that SA-PINNs work by producing a smooth equalization of the eigenvalues of the NTK matrix corresponding to the different loss terms. In numerical experiments with several linear and nonlinear benchmark problems, the SA-PINN outperformed other state-of-the-art PINN algorithm in L2 error, while using a smaller number of training epochs.

Motivation & Objective

• Motivate and address convergence and accuracy issues of baseline PINNs on stiff PDEs.
• Propose fully trainable, per-point adaptation weights that emphasize hard regions during training.
• Develop a continuous self-adaptive masking strategy and connect it to PDE-constrained optimization theory.
• Provide a practical training framework and theoretical insights into the dynamics of SA-PINN training.

Proposed method

• Define loss with per-point adaptive weights for initial, boundary, and residual points.
• Use a nonnegative, differentiable mask m(λ) to scale each point’s loss as a function of trainable λ.
• Optimize by min_w and max_{λ} to realize a saddle point, effectively a penalty method.
• Derive gradients for λ showing they increase with the corresponding unmasked loss and describe monotone weight growth.
• Map self-adaptive weights to a Gaussian Process for continuous weighting in SGD-friendly training.
• Discuss Neural Tangent Kernel (NTK) implications to understand training dynamics under SA-PINNs.

Experimental results

Research questions

• RQ1 Can per-point trainable weights improve PINN training and convergence on stiff PDEs?
• RQ2 How do self-adaptive weights affect loss component balancing and training dynamics?
• RQ3 What is the impact of SA-PINNs on benchmark stiff PDEs compared to baseline PINNs and prior weighting schemes?
• RQ4 Can a GP-based map of self-adaptive weights enable effective SGD training for SA-PINNs?

Key findings

• SA-PINNs achieve significantly lower L2 error than baselines on the Allen-Cahn equation (2.1% ± 1.21% vs 96.15% ± 6.45% with baseline and 49.61% ± 2.50% for nonadaptive weighting).
• In Burgers, SA-PINNs reach L2 error 4.803e-04 ± 1.01e-4, outperforming the baseline with fewer training epochs.
• For Helmholtz, SA-PINN attains relative L2 error 3.2e-3 ± 2.2e-4, approaching the accuracy of advanced schemes with fewer iterations.
• SA-PINNs learn interpretable weight maps that emphasize difficult regions (e.g., discontinuities in Burgers) and early times in Allen-Cahn.
• The NTK analysis suggests SA-PINNs equalize and smooth the eigenvalue distribution across loss components, aiding training dynamics.
• The authors provide an open-source implementation and demonstrate consistent improvements across multiple benchmarks.