Overview
PhysSAE is introduced as a mechanistic interpretability framework designed to analyze Physics-Informed Neural Networks (PINNs). The framework addresses the opacity of PINN internal representations, specifically the unknown nature of physical features encoded within hidden layers and their localized causal roles. PhysSAE employs overcomplete sparse autoencoders (SAEs) trained on PINN penultimate-layer activations. Evaluation of the discovered dictionary atoms is conducted through direct causal intervention, applied to the original frozen hidden state, bypassing the SAE decoder during this intervention.
Research Context
Physics-Informed Neural Networks (PINNs) incorporate partial differential equation (PDE) residuals into their neural network training process. Despite their utility, the internal mechanisms of PINNs remain largely uninterpretable. A key challenge lies in understanding what specific physical features are encoded by the hidden layers and whether these features exert localized causal influence within the network's operation.
Approach
The PhysSAE framework trains overcomplete sparse autoencoders (SAEs) on activations derived from the penultimate layers of PINNs. To evaluate the dictionary atoms identified by these SAEs, a direct causal intervention method is employed. This intervention operates on the original frozen hidden state, represented by the formula $h_{\mathrm{cf}} = h - \alpha z_k d_k$. This process explicitly avoids using the SAE decoder during the causal intervention step. The methodology was applied across six distinct PDE families. For each PDE family, the study utilized three different PINN seeds and three corresponding SAE seeds.
Findings
- Alignment with Physical Observables: Discovered SAE atoms demonstrated alignment with independently defined physical observables. The maximum Pearson correlation coefficient ($|r|$) observed was 0.951, consistently exceeding a permutation null.
- Causal Footprint Concentration: Ablation of top-aligned atoms resulted in a causal footprint that was 1.2 to 4.2 times more spatially concentrated than interventions performed using Principal Component Analysis (PCA) or Independent Component Analysis (ICA). This metric indicates a more precise localization of impact.
- Causal Localization Performance: Top-aligned atoms exhibited superior performance over matched random controls in causal localization for structured physical concepts. The observed ESF$_{80}$ advantage ranged from 0.04 to 0.44, indicating a better ability to isolate specific conceptual influences.
- Bilateral Representation Efficacy: The use of two-atom bilateral representations led to an improvement in concept regression R$^2$ values, with an increase of $\Delta R^2 = 0.05\text{-}0.15$ compared to single atoms. Conversely, random pairs of atoms caused a decrease in R$^2$ by up to 0.60, highlighting the specificity of effective atom combinations.
These findings collectively suggest that PINNs develop sparse, physically structured latent representations. Furthermore, these representations are identifiable and amenable to post-hoc causal interrogation.
Why This Matters
The ability to identify and causally interrogate physically structured latent representations within PINNs opens a path toward interpretability-aware scientific machine learning. This enhanced understanding of internal network mechanisms could facilitate further advancements in the application of neural networks to scientific domains.