A Graph-Based Inspection and Intervention Tool for Assessing Mechanistic Learning in PINNs

AI summary

Overview

Research area: Mechanistic interpretability for physics-informed neural networks (PINNs), sitting at the intersection of Scientific Machine Learning, representation analysis, and out-of-distribution generalization.

Technical level: Intermediate. The paper's framing questions are accessible, but the method involves differential-equation benchmarks, sensitivity analysis, and layer-wise perturbation protocols.

Scope: The paper introduces GIIT (Graph-based Inspection and Intervention Tool), which maps physical variables onto internal network components and then tests whether those discovered correspondences survive temporal extrapolation beyond the training domain.

What This Paper Is About

PINNs are trained by embedding differential-equation residuals directly into the loss, so a trained model can achieve near-zero physics residual and low output error while its hidden layers encode a representation that bears little relation to the causal structure of the physical system. This paper asks a specific question: do physically meaningful correspondences discovered inside a trained model remain meaningful when the model is evaluated outside the conditions under which those correspondences were found? GIIT answers this by building an explicit graph of the governing physics, linking graph nodes to internal network layers, and then perturbing those layers to see whether the links hold up as the evaluation window moves from in-distribution interpolation into progressively more severe extrapolation.

Key Contributions

1. The GIIT pipeline. A tool that automatically links physical concepts (such as mass, damping, or viscosity) to specific internal layers of a neural network using sensitivity analysis, with a physical specification compiled into a computational physics dependency graph and mapped onto candidate network tensors.

2. Intervention testing. A procedure that applies controlled multiplicative perturbations to identified layers and tests whether altering those components changes the corresponding physical observables, producing what the authors call intervention-validated functional correspondences rather than merely correlational attributions.

3. Out-of-distribution analysis. A transfer protocol that tracks mapping discovery across four evaluation regimes, revealing a systematic shift toward earlier-layer winners as the evaluation window moves from full in-distribution through late-ID, boundary-crossing, and severe OOD regimes.

4. A model-agnostic framing. The framework starts from known physical relationships and tests whether a trained model is consistent with them, rather than discovering causal structure from data or representations as most prior work does.

Main Findings

• Low residuals do not guarantee mechanistic correctness. On the 1D viscous Burgers' equation, the PINN achieves a physics residual of 1.915 × 10⁻⁴ on full in-distribution (t ∈ [0, 1]) and 1.01 × 10⁻⁴ on an ID sub-window (t ∈ [0.5, 0.8]), yet the functional mapping changes even before leaving the training domain.

• Residual error grows with extrapolation severity. The residual rises to 2.81 × 10⁻² on the boundary-crossing window (t ∈ [0.5, 1.5]) and to 3.30 × 10⁻¹ on severe OOD (t ∈ [1, 2]).

• Predictive accuracy collapses late. L2 error is 3.855% on full ID, 5.413% on the ID sub-window, 5.883% on the boundary-crossing window, and 31.752% on severe OOD.

• Internal mappings shift leftward. The average winner layer index drops from 5.71 (full ID) to 3.00 (ID sub-window), 2.43 (boundary-crossing), and 2.29 (severe OOD) across the seven physical nodes, measured over 100 perturbation replicates at σ = 0.1.

• Mapping stability degrades before output metrics do. Only 2 of 7 physical nodes maintain stable layer assignments across the boundary-crossing window (28.6%), and only 1 of 7 under severe extrapolation (14.3%). The boundary-crossing regime shows substantial internal reorganization despite a relatively stable L2 of approximately 5.88%.

• Drift is already visible in-distribution. Curvature-related terms (u_xx, νu_xx) migrate to the earliest hidden layer (fch.0.0.weight) in the late-time ID sub-window, reducing stable mappings to 2/7. Under severe OOD, only velocity (u) retains its baseline modal assignment (h5w, 29/100 replicates).

• High sensitivity does not imply a physically valid mapping. Viscosity (ν) is consistently localized to fch.5.0.weight in-distribution (69/100 replicates), but its intervention response usually has the opposite of the expected physical direction.

• Residual-loss sensitivity rises under shift. ∂L_R/∂ν is −1.705 on full ID, −5.064 on the ID sub-window, −73.576 on the boundary-crossing window, and −154.115 on severe OOD, remaining negative throughout but not capturing the layer-wise correspondence shifts GIIT detects.

• The damped harmonic oscillator shows the same pattern. In-distribution L2 is 1.13% with a physics residual of 2.007 × 10²; under OOD extrapolation (t ∈ [1, 2]) L2 explodes to 2007.34% and the residual to 1.317 × 10⁵. The PINN stops oscillating near t ≈ 1.0 and freezes at an unphysical fixed point (x̂ ≈ −1.19, ẋ̂ ≈ 0), while in phase space the true trajectory spirals to the origin (0, 0) and the PINN terminates at (−1.19, −0.19).

• Residual-loss sensitivities flip sign under extrapolation for the oscillator. ∂L_R/∂m moves from +1377.28 (ID) to −1747.92 (OOD), and ∂L_R/∂c from −3.84 to +579.62, while ∂L_R/∂k moves from +111.77 to +458.12.

• Mass and damping collapse to a shared tensor out of distribution. Across 21 perturbation replicates, mass and damping both map to fch.1.0.weight at 10/21 (47.6%) in-distribution and both collapse to fch.0.0.weight at 16/21 (76.2%) under OOD, while the spring constant shifts from fch.0.0.weight (8/21, 38.1%) to fch.1.0.weight (12/21, 57.1%).

• Derivative localization is highly consistent in-distribution. Velocity (ẋ) and acceleration (ẍ) map to fch.0.0.weight in 81.0% and 76.2% of oscillator replicates respectively.

Methodology in Plain English

The researchers start from physics they already know rather than inferring structure from data. A domain expert (or source code, optionally with LLM assistance) supplies a directed graph whose nodes are physical variables — parameters like mass, damping, and viscosity, intermediate quantities like gradients, and outputs — and whose edges encode known directional dependencies, for example mass → acceleration. A small executable ruleset computes derived observables from model output using automatic differentiation, such as turning x into x_t and x_tt, or computing amplitude as A = max(x) − min(x).

GIIT then links graph nodes to internal network components in two ways. For quantities computable from model output through a known formula, it can compute gradient sensitivity, S_grad(W_ℓ, q) = mean |∂q/∂W_ℓ|, across candidate layers. For governing parameters like mass or damping, which are not direct network outputs, it uses trend-based perturbation mapping: it multiplies a layer's weights by (1 + ε) where ε = σ|N(0,1)| is strictly positive, runs a full forward pass, and measures the relative change in the target observable, S_pert = |q(W'_ℓ) − q(W_ℓ)| / (|q(W_ℓ)| + δ). Layers are ranked mainly by response magnitude, with directional sign checked where the physics has a known direction. The paper uses perturbation sensitivity rather than gradient sensitivity for mapping discovery because gradient sensitivity can also reflect bias parameters.

To test transfer, the study defines four regimes: full in-distribution (t ∈ [0, 1]), an ID sub-window (t ∈ [0.5, 0.8]) isolating late-time dynamics before the boundary, a boundary-crossing window (t ∈ [0.5, 1.5]), and severe extrapolation (t ∈ [1, 2]). The baseline mapping is discovered on full ID data, then independently re-evaluated in each other regime using the identical multi-replicate perturbation protocol. Drift is quantified by modal layer stability — whether the top-mapped component is preserved — plus systematic layer index shifts, alongside L2 error and mean squared physics residual.

Evaluation uses the 1D viscous Burgers' equation (u_t + uu_x = νu_xx) adapted from PINNsTorch, on x ∈ [−1, 1], and a damped harmonic oscillator appendix benchmark. The Burgers' PINN takes (x, t) through Linear(2→20) → tanh, seven sequential hidden blocks fch.0.0 through fch.6.0 each Linear(20→20) → tanh (8 layer stages total), and Linear(20→1) to a scalar velocity field, giving 3,021 trainable parameters with ν = 0.01/π ≈ 0.003183. L2 error and residuals are computed on a uniform 100 × 100 space-time grid against a high-resolution finite-difference reference, using 100 perturbation replicates at σ = 0.1 on one trained model. The oscillator uses 2,209 parameters, Linear(1→32) → tanh → Linear(32→32) → tanh → Linear(32→32) → tanh → Linear(32→1), with 500 collocation points in each of the ID and OOD regimes and 21 perturbation replicates.

Why This Matters

Impact on research. The paper argues that near-zero physics residual and low output error are a form of numerical constraint satisfaction, not a guarantee of mechanistic correctness. Existing XAI tools — saliency, SHAP-style attribution, counterfactual explanations, Causal Dependence Plots, DAG-based SHAP — operate on model inputs and stop at the input-output boundary. GIIT's contribution is testing correspondence at the level of internal neural representations, using known physical structure as the starting point rather than inferring structure from data. The finding that correspondence drift appears within the nominal training domain suggests that output-level accuracy and physics residuals alone may be insufficient for diagnosing changes in internal functional organization.

Real-world applications (note that the paper's own evidence is limited to two low-dimensional dynamical systems, so these are implications the framing opens up rather than validated deployments):

• Auditing surrogate models used in scientific simulation, where a model is trained on a limited time window and then queried beyond it.
• Validating physics-based digital twins, where the internal model must track the real system's governing mechanisms rather than merely fit observed outputs.
• Quality checks for engineering design tools that rely on low-data, high-simulation-cost PINN surrogates.
• Diagnostics for safety-critical extrapolation, where a model may look accurate on conventional metrics while its internal representation has silently reorganized.

Industry relevance. PINNs are attractive precisely where simulation is expensive and data is scarce, a regime common in scientific computing and engineering analysis. A tool that flags correspondence changes before output metrics degrade offers a potential early-warning signal for models headed toward out-of-distribution failure, and the framework is described as model-agnostic.

Future Directions

• Activation-level analysis. The observed leftward migration has at least one alternative explanation: deeper representations may simply become less sensitive during extrapolation, causing interventions to select earlier layers. The authors state that activation-level analysis is needed to distinguish this from a broader redistribution of sensitivity.

• Multiple independently trained models. The current evidence comes from a single trained model per benchmark; the authors call for testing whether these effects persist across independently trained models.

• Higher-dimensional and larger scientific systems. The evidence is limited to two low-dimensional dynamical systems, so generalization to larger or higher-dimensional SciML models remains untested.

• Sharpening the interpretive claim. The identified tensors should be read as intervention-validated functional correspondences, not unique causal encodings or isolated modular circuits. Perturbing a tensor demonstrates that altering an identified component changes the target observable, not that the tensor uniquely or exclusively encodes the physical concept. The viscosity result, where response magnitude was high but direction was wrong, shows why this distinction matters.

Target Audience

This paper benefits researchers working on mechanistic interpretability and explainability for scientific machine learning, PINN practitioners who need to trust models under extrapolation, and domain scientists applying neural surrogates in low-data regimes. It is also relevant to readers interested in causal abstraction and representation analysis, and to engineers evaluating whether a simulation surrogate's internal mechanisms match the physics it was trained on. Readers with a working grasp of neural networks and differential equations will get the most from the methods and results sections, while the framing and discussion are accessible to a broader technical audience. Code, dataset configurations, and model weights are available at the repository linked in the appendix.

Authors’ abstract

We ask whether physically meaningful correspondences discovered inside a trained scientific model remain meaningful outside the conditions under which they were discovered. We introduce GIIT (Graph-based Inspection and Intervention Tool), which represents governing physics as a computational physics dependency graph, maps graph nodes to internal network components via sensitivity- and trend-based discovery, and tests the resulting mapping under targeted intervention. Evaluating on temporal extrapolation out-of-distribution (OOD) regimes of increasing severity, we find that temporal extrapolation is associated with layer-wise correspondence drift and reduced functional correspondence. Specifically, while the model achieves low physics residual in-distribution (1.915 x 10-4 on full ID and 1.01 x 10-4 on an ID sub-window t in [0.5, 0.8]), the functional mapping changes even before leaving the training domain, and the shift continues as the evaluation window extends beyond the training domain: residual error increases from 2.81 x 10-2 on the boundary-crossing window (t in [0.5, 1.5]) to 3.30 x 10-1 on severe OOD (t in [1, 2]), accompanied by a systematic leftward shift of internal layer mappings, where the average winner layer index drops from 5.71 (ID) to 3.00 (ID sub-window) down to 2.29 (severe OOD). Only 2 of 7 physical nodes maintain stable layer assignments across the boundary-crossing window, and only 1 of 7 under severe extrapolation, revealing potential internal functional correspondence changes before severe degradation in conventional output metrics. Additional results for linear oscillatory systems are further detailed in the appendix.