Learning in PINNs: Phase transition, diffusion equilibrium, and generalization

In plain words

Training a physics-informed neural network means teaching a model to satisfy a physics equation at many points at once. Each point is like a student that has to be marked correct. Training does not go smoothly: it moves through rough stages, and some points can be left behind while the average looks fine.

The authors track how well the gradient pushes from different points agree with each other. Agreement is measured as a signal-to-noise ratio. They see the early stages familiar from information bottleneck theory, plus a third stage they call diffusion equilibrium. There the pushes line up, the ratio settles, and the error drops fastest. Convergence is best when the remaining error is spread evenly across the points instead of piling up in a few.

To encourage that spread, they propose residual-based attention. Points whose errors stay large get larger weights, so they cannot stall. On the Allen-Cahn test problem the reweighted model reaches a 10% error about ten times sooner, and it avoids a stuck solution on the cavity-flow problem. The idea is not tied to physics-informed networks, but the evidence comes from a small set of physics test problems.

TL;DR

Full-batch Adam on PINNs goes through fitting and diffusion phases, then a third phase, "diffusion equilibrium" (DE), where per-sample gradients align and the test error falls fastest §1.4 §4. Reaching DE also needs the residuals to be homogeneous across collocation points §2.4. Residual-based attention (RBA) reweights points by their recent normalised residuals and reaches a given error sooner on Allen-Cahn fig. 4. The results are author-reported on five test problems; the repository holds code for three of them.

Problem

PINNs approximate a PDE solution by minimising residuals at collocation points §1.2. Non-convex training is often dominated by discrete phases §1.1. The authors ask why training can settle at a solution where some samples are still unresolved. With a quadratic loss, the batch-averaged gradient can point in a consistent direction while individual sample gradients cancel out, and the relative L2 test error is very sensitive to outliers, so every sample must be resolved §2.1 §2.3. The analysis assumes full-batch Adam and fully-connected networks §2.1 §3.1.

Method

The paper measures training dynamics, names the phases, links them to residual homogeneity, and proposes a reweighting scheme. Its information-compression analysis is a separate, descriptive part.

1. Loss and noise model. The mean squared residual is minimised with mini-batch updates (1)–(3). Each batch gradient is the true gradient plus zero-mean noise (4), so each step splits into a deterministic and a stochastic part (5) §2.1.

   \min_\theta \mathcal{L}(\theta) = \frac{1}{n} \sum_{i} \mathcal{R}(f(x_i;\theta), y_i)^2 eq. 1

2. Batch-wise SNR. The SNR is the norm of the mean batch gradient divided by the norm of its batch standard deviation, over the weights (7). The paper computes batch gradients at each iteration without updating, so it can see sample-level disagreement under full-batch training §2.1.

   \mathrm{SNR} = \frac{\lVert \mathbb{E}_{\mathcal{B}}[\nabla_\theta \mathcal{L}_{\mathcal{B}}] \rVert_2}{\lVert \mathrm{std}_{\mathcal{B}}[\nabla_\theta \mathcal{L}_{\mathcal{B}}] \rVert_2} eq. 7

3. Phases. Fitting has high SNR and a fast decline. Diffusion has a falling SNR while the error barely moves. DE is a sudden rise to a plateau where the SNR is flat (∂SNR/∂t ≈ 0), and the error then drops steeply. The paper reports these phases for PINNs trained with Adam §2.1 fig. 2 §3.1.

4. Adam's view. Adam's per-parameter step scale is driven by a similar ratio. The paper defines SNR_Adam with the first and second moments (13) and finds it tracks the batch SNR and reaches equilibrium at DE §2.2 §3.2 fig. 7.

5. Order parameter. With φ = ‖mean‖ / ‖rms‖ of batch gradients (16), the paper shows SNR = φ / √(1 − φ²) (17) §2.1 app. B. SNR > 1 holds for φ ∈ (√2/2, 1) (18), and SNR ≈ φ for small φ (19).

6. Residual homogeneity. Zero gradient (23) and Σ_i R_i ∂R_i/∂θ = 0 (24) do not force the residual magnitudes to be equal. A sample can keep a large residual while its gradient is near zero §2.4 §2.5.

7. RBA weights. The residual term is weighted per point, λ_i² R_i² (25). Each λ_i is a moving average of the largest-normalised residual (26), with γ = 1 − η*, so λ_i stays in [0, 1] §2.5. The order parameter for residuals is φ_RBA (28) §3.2.

   \lambda_i^{t+1} = \gamma\lambda_i^{t} + \eta^{*} \frac{\lvert \mathcal{R}_i \rvert}{\lVert \mathcal{R} \rVert_\infty} eq. 26

8. Compression. The paper computes mutual information as the entropy of the discretised layer activations (29). Each neuron uses 30 bins over the layer's mean activation range. Saturation is the share of activations at the bounds of tanh. DE coincides with saturation, deeper layers become close to binary, and the middle layers saturate most. The paper argues geometric compression is not necessary for good generalisation §3.3 fig. 8 app. C.

The paper's loop, as code, is ordinary full-batch Adam with the RBA weights updated once per step before the loss is formed. Pseudocode below follows the repository (see Paper vs code).

# Residual-based attention: weights and weighted loss  [§2.5] [eq. 25-26]
# Input:  residuals (N,) at all collocation points; rsum (N,) from last step
# Output: updated rsum, weighted residual loss

def rba_update(residuals, rsum, eta_star=0.001, gamma=0.999):
    # Scale |R| by its largest value over all points  [eq. 26]
    relative = abs(residuals) / max(abs(residuals))
    # Moving average; no gradient flows through the weights  [eq. 26]
    rsum = gamma * rsum + eta_star * relative
    # Residual term weighted per point; IC/data term keeps weight 1  [eq. 25]
    loss_residual = mean((rsum * residuals) ** 2)
    return rsum, loss_residual

# Batch-wise gradient SNR over weight matrices  [§2.1] [eq. 7]
# Input:  weights (weight matrices only), batches (index sets)
# Output: scalar SNR

def batch_snr(weights, batches):
    # Gradient of each batch at the same parameters; no optimizer step  [§2.1]
    grads = [batch_gradient(weights, batch) for batch in batches]
    # Per-element mean and population std over batches  [§2.1] [app. B]
    mu = mean(grads, axis=batches)
    sigma = std(grads, axis=batches)
    # Global norms over all weight elements  [eq. 7]
    return norm(mu) / norm(sigma)

Hyperparameters

Results

Author-reported. Setup: PINNs trained with full-batch Adam on Allen-Cahn, Helmholtz, lid-driven cavity (Re = 1000), a 3D cerebral CSF flow, and Burgers; vanilla versus RBA; relative L2 test error §3.1 app. C app. D. The paper has no tables. The numbers below come from figure captions and text.

Only the Allen-Cahn and cavity numbers are stated in the text or captions. The paper reports no seed spread, confidence intervals or compute cost.

Paper vs code

Checked against the repository at code_commit. Code is MIT-licensed, by the first author LICENSE. The repository holds the Allen-Cahn, Helmholtz and cavity scripts and the SNR and mutual-information utilities, but not the Burgers or cerebral-flow scripts.

• Cavity collocation count. Paper: 25,600 points app. D. Code: N_r = 10000 CV.py.
• RBA normalisation. Paper divides by ‖R_i‖_∞ and bounds λ in [0, 1] §2.5 eq. 26. Allen-Cahn and cavity divide by the largest |R| over all points AC.py CV.py. Helmholtz first subtracts the minimum and divides by the range HM.py. The appendix says this shift was only "for better visualization" and "did not affect performance" app. D. In the code the shifted weights enter the training loss HM.py.
• First RBA step. Helmholtz sets init = 2, so the first update uses η = 1 HM.py HM.py. Allen-Cahn and cavity use η = 0.001 from the first step AC.py CV.py. The paper describes no special first step.
• RBA constants. Paper: γ = 1 − η*, no numbers §2.5. Code: η* = 0.001, γ = 0.999 AC.py.
• Cavity velocity convention. Paper: u = ∂ψ/∂y, v = −∂ψ/∂x app. D. Code: u = −∂ψ/∂t, v = ∂ψ/∂x, where t is the second coordinate CV.py. The signs differ. Flipping both u and v changes the viscous term, so follow the code's convention and do not mix the two.
• Cavity boundary loss. Paper: sums of squared boundary residuals app. D. Code: squared sum of absolute velocity errors on the prescribed points CV.py.
• Cavity lid. Paper: lid velocity (1, 0) on the top edge, corners smoothed app. D. Code: reference-solution values on the first grid row and zeros on the other three walls CV.py. The paper's lid profile is not used, and the card did not verify which physical wall that row is.
• Cavity error metric. Code's reference is the velocity magnitude √(u² + v²) CV.py. The paper does not say which field its cavity relative L2 uses §2.1.
• SNR parameters. Paper: "all network parameters" §2.1. Code: weight matrices only, biases excluded snr.py func.py.
• SNR under RBA. Code's batch gradients use the RBA-weighted loss AC.py. The paper does not say whether SNR uses weighted or unweighted residuals.
• Smoothing window. Paper: 100 past iterations §3.1. Notebook: a window of 60 stored rows Paper_plots.ipynb. The code stores a row every 100 iterations, and every 10 before iteration 300 [code: utils/snr.py#L21]. So the window covers thousands of iterations, not 100.
• Adam SNR. Paper: simplified update without bias correction §2.2. Code: bias-corrected moments and includes biases snr.py. The diagnostic is off by default AC.py.
• Mutual-information bins. Paper: 30 equidistant bins per neuron §3.3. Code: bin width is the mean range divided by 30, so each neuron gets range / width + 1 bins, about 30 on average mi.py. The diagnostic is off by default AC.py.
• Helmholtz forcing. Paper's forcing has "+k sin..." while the PDE has k² app. D. Code uses k² HM.py. With k = 1 they agree.
• Allen-Cahn initial data. Paper: u(0, x) = x² cos(πx) app. D. Code: reference values at t = 0 from AC.mat, 512 points AC.py. The code has no analytic formula.
• Scheduler. Paper: "exponential learning rate scheduler" with no constants app. D. Code: ×0.9 every 5000 iterations AC.py.
• Residual-order figure. The notebook reads residual files that the code writes only in commented lines AC.py AC.py. The notebook also shuffles residuals into groups of 100 points [code: snr_plots/Paper_plots.ipynb#L400-L411]. This figure cannot be regenerated from the current code.
• Phase boundaries. DE onsets are typed into the notebook, for example 91,700 and 11,100 for Allen-Cahn Paper_plots.ipynb Paper_plots.ipynb. No code detects DE.
• Missing code. The Burgers and cerebral-flow scripts and the loss-landscape visualisation fig. 3 are not in the repository. The repository has precomputed Burgers and cerebral results BN.

Pitfalls

• RBA warm-up. With η* = 0.001, the weights start near η* and approach their steady value with time constant 1/(1 − γ) = 1000 iterations. Derived from the update: λ_t ≈ r̂ (1 − 0.999^t), where r̂ is the normalised residual [code: AC/AC.py#L282-L287]. For the first few hundred iterations the residual term is almost switched off, so the early dynamics are close to unweighted training of the data term alone.
• RBA only down-weights the residual term. Since λ ≤ 1, the residual term is never stronger than vanilla. The data term keeps weight 1 [code: AC/AC.py#L287-L293]. The balance between data and residual therefore shifts over training. This is inferred from the code, not stated in the paper.
• Global normalisation. Dividing by the largest residual over all points lets one outlier set the scale for every weight. The Helmholtz min-max variant changes this again HM.py. Inferred from code.
• Population std. The identity SNR = φ / √(1 − φ²) app. B needs the population standard deviation (ddof = 0). np.std uses ddof = 0 by default snr.py, so keep that. A sample std breaks the identity.
• Diagnostic cost. Each SNR logging step runs one backward pass per batch, 100 in total, every 100 iterations and every 10 before iteration 300 [code: utils/snr.py#L21-L33]. After iteration 300 this adds about one full-batch backward pass per 100 iterations. Before iteration 300 the rate is ten times higher. Derived from the code.
• Batch construction. The SNR depends on batch size and layout. The paper gives a batch-independent SNR_S = SNR / √m app. A and effective batch sizes m/2 for Allen-Cahn and m/16 for Helmholtz to account for symmetry app. A. The code does not implement these.
• Late-training step sizes. With the code's schedule, the learning rate after 3×10^5 iterations is 1e-3 × 0.9^60 ≈ 1.8e-6 (derived). Late dynamics are driven by very small steps, which matters when reading the late SNR curves.
• DE is found by eye. The paper does not give a detector for the DE onset, and it states that the first-order nature of the onset is unexplained §4. Plots use hand-set vertical lines Paper_plots.ipynb.
• One architecture. All cases use 6 × 128 tanh networks and no sweep was run §3.1. The DE onset and its bounds may depend on architecture and problem stiffness §3.1.
• Evaluation grids differ. The relative L2 error is computed on different grids: 512 × 201 for Allen-Cahn AC.py, 1001 × 1001 for Helmholtz HM.py, 256 × 256 for the cavity [code: CV/CV.py#L161]. Compare errors only on the same grid.
• Log labels. The L_inf value printed in the training log uses the same L2 formula AC.py. Do not read it as the max-norm error.
• Single seed. The code fixes one seed AC.py and the paper reports no spread. Treat the 10× speed-up as a single-run result.
• Version. The source is arXiv v2, which the paper's note says matches the journal version. Figure numbers here were counted from the LaTeX source, in order. Check them against the PDF before citing.

Lineage

• Builds on: the Adam optimiser (arXiv:1412.6980) §2.2; the edge-of-stability self-stabilisation analysis (arXiv:2209.15594) §1.1; information bottleneck theory §1.3, which the paper's bibliography cites without arXiv ids.
• The RBA scheme is attributed to the authors' earlier paper (anagnostopoulos2024residual in the bibliography), which has no arXiv id in this source §2.5.
• Compared against: vanilla PINN training with uniform residual weights §3.1.
• Superseded by: not known from this source.