Residual-based attention in physics-informed neural networks

In plain words

A physics-informed neural network learns the solution of an equation by being penalised whenever its output breaks that equation at sample points. Ordinary training averages the penalty over all sample points, so a small region with a large error can be drowned out. The model then looks fine on average while some sharp feature stays wrong.

This paper adds a simple attention mask. Every sample point keeps a running score that fades over time and grows with how large the equation error has been at that point. Points with a large recent error count for more in the penalty, so training is pushed toward the hard spots. The score uses only the errors that training already computes, so it costs almost nothing extra and needs no second network to learn the weights.

The authors also use a few standard tricks that help these networks train, such as building the boundary conditions directly into the network. They test the idea on a stiff equation with sharp transitions, on a Helmholtz equation, and on a fluid-flow reconstruction from measurements.

TL;DR

Residual-based attention (RBA) gives each collocation point a multiplier that is a decaying sum of its normalised PDE residual. The multipliers are never trained by gradient descent, so the method adds no auxiliary network. The authors report a relative L2 error of 5.8e-5 on 1D Allen-Cahn and 8.21e-6 on 2D Helmholtz with Fourier features, both five-seed means, with faster early convergence than vanilla PINNs. Implement from Eq. 12 and Algorithm 1, but use Adam as in the experiments, and check the Paper vs code section: the released code differs from the paper in several numbers and in one loss form.

Problem

A PINN fits a network to a PDE by minimising a weighted sum of loss terms: the PDE residual, the initial and boundary conditions, and any data §2.1 eq. 6. Each term is a mean over its points eq. 10. A plain mean lets a small set of points with large errors go unnoticed while the total loss keeps falling, which hurts multiscale problems and can block the flow of information from the boundary into the domain §2.2.

Earlier weighting schemes address this but carry costs. Self-adaptive weights train an extra multiplier per point by gradient ascent §1. Other methods need an auxiliary network, extra multipliers, or gradients of the loss, and some multipliers grow without bound §1. Exact boundary conditions help but are problem-specific §2.3.3.

Method

Loss with multipliers. Each loss term is the mean of a squared, weighted residual. The multipliers are global (one scalar per term) or local (one per collocation point). RBA uses local multipliers on the PDE residual only §2.1 eq. 10 §3.1 §3.2.

The RBA update. For collocation point i at iteration k, the multiplier decays by a factor gamma and receives the point's residual, divided by the largest residual magnitude in the current set §2.2 eq. 12:

\lambda_i^{k+1} = \gamma\,\lambda_i^{k} + \eta^{*}\,\frac{|e_i|}{\lVert e \rVert_{\infty}} eq. 12

Here eta* is the RBA learning rate, distinct from the network's learning rate eta §2.2. The normalisation keeps each new term in [0, 1], so the sequence is bounded. As k grows the multipliers settle inside

\lambda_i^{k} \in \left(0,\ \frac{\eta^{*}}{1-\gamma}\right] eq. 13

§2.2. The decay factor keeps the multipliers from exploding §2.2 app. A.1.

Initialisation. All multipliers start at zero §2.2. Two options are given. Mode 1 applies eta* from the first step. Mode 2 uses eta* = 1 at the first step only, which makes the first multiplier equal to the normalised residual §2.2 alg. 1. The authors report similar results for both §3.4.

Offset. An optional offset lambda_o raises the lower bound of the multipliers to lambda_o. The paper uses it only in the Navier-Stokes case §2.2 app. A.3.3.

Gradient handling. The multipliers are not trained by gradients. In the loss they are held constant, so the back-propagation goes only into the network §2.2 allen-cahn_rba-sa.ipynb.

Training loop. Algorithm 1 is written for plain gradient descent. Each step updates the multipliers, forms the weighted loss, and takes a gradient step on the network weights alg. 1 eq. 11. The Allen-Cahn and Helmholtz experiments use Adam instead §3.1 §3.2.

# Residual-based attention: one training step  [alg. 1] [eq. 12]
# Input:  theta (network weights); x_r (N, 2) collocation points
# State:  lam (N,) RBA weights, zero at the start  [§2.2]
# Settings: eta (network rate), eta_star (RBA rate), gamma (decay)

def rba_step(theta, lam, x_r):
    # Residual of the PDE at each collocation point  [§2.1] [eq. 1]
    residual = pde_residual(theta, x_r)              # (N,)

    # Divide by the largest residual magnitude, so each term is in [0, 1]
    # [eq. 12]
    normalised = abs(residual) / max(abs(residual))  # (N,)

    # Decaying sum: lam is state, not a trained parameter  [eq. 12]
    lam = gamma * lam + eta_star * normalised        # (N,)

    # Weighted squared residual, lam held constant in the gradient  [alg. 1]
    # AIV only: use lam + offset here  [§2.2]
    loss = mean((lam * residual) ** 2)

    # Gradient step on the network  [eq. 11]; the experiments use Adam
    # [§3.1]
    theta = theta - eta * grad(loss, theta)
    return theta, lam

Key quantities. The bound in eq. 13 gives the largest multiplier the method can reach, so it is the number to check when choosing eta* and gamma. The paper reports that the mean multiplier settles at about 20% of this bound in the Allen-Cahn and Helmholtz runs app. A.1. In those two cases RBA acts on the PDE residual of Eq. 1 only §3.1 §3.2.

How RBA is applied in each case.

• 1D Allen-Cahn (dynamic, stiff): RBA on the PDE term only. The initial condition term has weight w_ic = 100 and the PDE term weight w_r = 1 §3.1.
• 2D Helmholtz (static): RBA on the PDE term only, with w_bc = 100 on the boundary §3.2. The second variant uses approximate distance functions and drops the boundary loss §3.2 eq. 32.
• 3D Navier-Stokes inverse problem (artificial intelligence velocimetry, AIV): RBA on the boundary, the Navier-Stokes residual and the data terms, with a global weight per term §3.3. The data terms use one multiplier per point. Boundary and Navier-Stokes points share a multiplier per training group, and training uses ordered mini-batches of one point per group §3.3 fig. 5 fig. 6.

Network and boundary tools used with RBA. The experiments add three existing tools. The notation below follows the paper.

• Modified MLP (mMLP): two encoders U and V, each a tanh layer on the input. Each hidden layer's output mixes them: (1 - alpha) * U + alpha * V §2.3.1 eq. 14 eq. 15 eq. 16 eq. 17.
• Weight normalisation: each weight vector is split into a length g and a direction v, theta = g v / ||v|| §2.3.2 eq. 18 eq. 19.
• Exact Dirichlet boundary conditions via an approximate distance function: u = g + phi * u_NN, where phi vanishes on the boundary §2.3.3 eq. 20 eq. 21.
• Exact periodic boundary conditions via Fourier features with period P, in one and two dimensions §2.3.3 eq. 22 eq. 23.

The Helmholtz experiments use a truncated two-dimensional Fourier embedding with m = n = 5 §3.2 eq. 31. The paper does not give the period P for this case. It also says the embedding is partly for comparison, because it contains the trigonometric structure of the solution §3.2.1.

Hyperparameters

Values are the paper's unless the Notes column says the code differs. Each case has its own block below. "not stated" means the paper and the released code give no value.

1D Allen-Cahn (PyTorch)

2D Helmholtz, Fourier features (JAX)

2D Helmholtz, approximate distance function

3D Navier-Stokes (AIV). No AIV or Navier-Stokes code is in the repository.

Results

Author-reported. All numbers are relative L2 errors unless stated. The baseline rows for Allen-Cahn and Helmholtz are values from the literature; only the SA rows were re-run by the authors, as five-seed means with their enhancements §3.1 tab. 1 tab. 3.

1D Allen-Cahn, comparison (PyTorch; the RBA row is the full model, RBA + mMLP + Fourier features; the best prior methods also use enhancements) §3.1 tab. 1

1D Allen-Cahn, ablation (single best seed of Table 1 is the full model) §3.1.1 tab. 2

2D Helmholtz, comparison (JAX; Fourier and ADF variants; the SA rows are the authors' own runs) §3.2 tab. 3

2D Helmholtz, ablation (Fourier and ADF; the RBA + mMLP + Fourier entry is the best seed, as the text states) §3.2.1 tab. 4

3D Navier-Stokes inverse problem (AIV), validation set. Error is on the 2D PTV velocity (u, v) only, for the validation split §3.3 tab. 5

AIV ablation §3.3.1 tab. 6

Convergence. On both benchmarks the paper reports that RBA makes the vanilla PINN converge about one order of magnitude faster, even without loss weights w. SA weights are about four times faster than vanilla in the same setup §3.4 fig. 8. For Allen-Cahn, the vanilla network converges only because the learning rate is scheduled §3.4.

Cost. Time per iteration, from the authors' table: Allen-Cahn 17 ms (vanilla) and 17 ms (RBA); Helmholtz 13.8 ms and 13.8 ms; AIV 325 ms and 329 ms. The authors note that the code was not optimised for speed app. A.2 app. A.3.3.

Paper vs code

Official repository: the paper links github.com/soanagno/rba-pinns §1. The repository is MIT-licensed, and its LICENSE names Sokratis Anagnostopoulos, the first author LICENSE. The commit is c8ca4cb5. It contains Allen-Cahn (PyTorch), Helmholtz (JAX) and a Helmholtz comparison notebook. It has no AIV or Navier-Stokes code and no AIV data.

• Allen-Cahn iterations. Paper: 3e5 §3.1. Enhanced notebook: 300000 allen-cahn_enhanced.ipynb. Agrees.
• Comparison notebook iterations. Paper: Section 3.4 uses the same settings as the ablations §3.4. The RBA-vs-SA notebook runs 30000 iterations, has no Fourier features and no initial-condition weight allen-cahn_rba-sa.ipynb allen-cahn_rba-sa.ipynb. Unexplained.
• Allen-Cahn RBA learning rate. Paper: eta* = 1e-2 app. A.2. Code: 1e-3, with the comment that 1e-2 would give an upper bound of 10, so the initial condition term is scaled by 10^2 instead allen-cahn_enhanced.ipynb. Because eta* scales every multiplier linearly, the paper's value makes the PDE term 100 times the code's at equal residuals, while both use an initial-condition weight of 100. The code's balance is not the paper's literal values.
• Allen-Cahn Fourier embedding. Paper gives the form eq. 22 but no frequency count. Code: M = 10 frequencies on period L = 2, and it also appends the time coordinate and a constant allen-cahn_enhanced.ipynb allen-cahn_enhanced.ipynb. Eq. 22 has no time input.
• Allen-Cahn boundary points. The code samples 200 boundary points at x = -1 and x = 1 allen-cahn_enhanced.ipynb. The loss uses them only when the symmetry flag is on allen-cahn_enhanced.ipynb, and it is off here. Periodicity therefore comes from the Fourier embedding alone, which the paper describes §2.3.3.
• Optimiser. Algorithm 1 and Eq. 11 are plain gradient descent. All runs use Adam with an exponential schedule allen-cahn_enhanced.ipynb allen-cahn_enhanced.ipynb. The text says Adam §3.1.
• mMLP gate order. Paper, Eq. 17: (1 - alpha) * U + alpha * V. PyTorch code matches allen-cahn_enhanced.ipynb. JAX code computes alpha * U + (1 - alpha) * V, which swaps the roles of U and V helmholtz_fourier.ipynb. Both train; the paper's equation is the literal reading.
• Helmholtz loss form. Paper, Eq. 10 and Algorithm 1: squared, weighted residual, (lambda * e)^2. PyTorch code: (lambda * e)^2 allen-cahn_enhanced.ipynb. JAX code: lambda * e^2, because its error function multiplies the squared residual by the weight helmholtz_fourier.ipynb helmholtz_fourier.ipynb. For multipliers in [0, 10] these differ, so the JAX loss is not the paper's loss.
• Helmholtz residual points. Paper: 1201 app. A.2. Code: 101 x 101 = 10201 uniform points helmholtz_fourier.ipynb. The ADF notebook uses the same grid helmholtz_adf.ipynb.
• Helmholtz L-BFGS. Paper: 2e3 iterations §3.2. Code: five runs of maxiter = int(1000 / 5) + 1 = 201, so at most 1005 iterations in total helmholtz_fourier.ipynb helmholtz_fourier.ipynb.
• Helmholtz Fourier features. Eq. 31 is written with paired indices. The code takes every pair of n and m in 1..5 for three product types, giving 75 inputs helmholtz_fourier.ipynb helmholtz_fourier.ipynb. The code hard-codes P = 2, so omega = pi; the paper gives no P.
• Helmholtz learning rates. Code matches the appendix for the ADF variant: eta* = 1e-2, gamma = 0.999 helmholtz_adf.ipynb helmholtz_adf.ipynb. Fourier variant: eta* = 1e-3, gamma = 0.9999 helmholtz_fourier.ipynb helmholtz_fourier.ipynb. Both match.
• Helmholtz BC weight in ADF mode. Code sets the boundary weight to 0 when the ADF mode is chosen helmholtz_fourier.ipynb, consistent with Eq. 32, which satisfies the boundary exactly.
• Navier-Stokes / AIV. Not in the repository; there is no AIV data either. Everything in §3.3 and the AIV appendix is paper-only.

Pitfalls

• Multiplier bounds. From Eq. 13, the upper bound is eta* / (1 - gamma) plus any offset. Derived values: Allen-Cahn 1e-2 / 1e-3 = 10 (paper's value); Helmholtz 1e-3 / 1e-4 = 10; AIV PTV and Navier-Stokes 0.2 / 1e-3 = 200 and boundary 0.02 / 1e-3 = 20, each plus the offset of 1.0 (paper gives no AIV bound; these are derived). Any value you pick should be checked against this bound.
• Normalisation by the maximum. The divisor is the largest residual in the current set. One large outlier shrinks every other new term. This is the intended effect in the paper, but it means the mask depends on the single worst point, and it is not stated how the paper handles outliers beyond the AIV observation §3.3 app. A.3.3.
• Do not differentiate the multipliers. They are state. The notebook detaches them allen-cahn_rba-sa.ipynb. If they are left in the graph, the training changes.
• Adam, not plain gradient descent. Algorithm 1 is written for gradient descent. Following it literally is not what the experiments did.
• Fixed Fourier periods. The embedding hard-codes the period to the domain length (L = 2 for Allen-Cahn, P = 2 for Helmholtz, both on [-1, 1]). A different domain needs a new period.
• Helmholtz features contain the answer's frequencies. With omega = pi and indices up to 5, the basis includes the frequencies of the exact solution (1 and 4). The paper says this is for comparison only §3.2.1. Derived from the code, not stated in the paper.
• Single-seed numbers in the ablations. The Allen-Cahn ablation uses the best seed of Table 1 §3.1.1. The Helmholtz ablation's best-seed number is 5.18e-6 §3.2.1. Use the five-seed means in Results for comparisons.
• Baselines are literature values. Apart from the SA rows, the baseline numbers in Tables 1 and 3 were not re-run §3.1 §3.2.
• Internal inconsistency, AIV split. The text says 60% of PTV data for training and 40% for validation §3.3. The appendix table lists 9.4e3 train and 1.4e4 validation, which is 40% and 60% app. A.3.3. It is not stated which is correct.
• Internal inconsistency, AIV iterations. The text says 2.5e5 iterations §3.3. The ablation figure caption says 2.5e4 fig. 7.
• Cost. RBA stores one multiplier per collocation point, which is O(N) memory. For AIV the authors share multipliers across training groups and use ordered mini-batches to reduce this §3.3 fig. 6.
• Vanilla convergence on Allen-Cahn. Vanilla PINNs converge there only because of the learning-rate schedule §3.4. Keep the schedule when you compare.
• Timing caveats. The authors say the code is not tuned for speed app. A.2. The Helmholtz and Allen-Cahn runs were on RTX 3090, and AIV on A100; the cross-case timings are not comparable.

Lineage

• Builds on: standard PINNs for PDEs (the "vanilla PINN" baseline) §1 §2.1; the exponentially weighted moving average of Adam and RMSprop, which the RBA update copies §2.2; weight normalisation §2.3.2; the modified MLP §2.3.1; approximate distance functions for exact Dirichlet conditions §2.3.3; Fourier features for periodic conditions §2.3.3.
• Closest method: self-adaptive weights, which train one multiplier per point by gradient ascent §1 §3.4. The paper compares against it throughout.
• Compared against: time marching; causal training (Wang et al.); DASA-PINNs (Zhang et al.); PECANN-AL (Basir et al.); AL-PINNs (Son et al.); learning-rate annealing (Wang et al.) §3.1 §3.2 §3.3 tab. 1 tab. 3 tab. 5.
• Superseded by: not stated in the paper or the repository.