Sparse probes and murky physics: a case study of interpretability challenges in a foundation model for continuum dynamics

Technical Summary: Sparse Probes and Murky Physics

Problem Statement
Generative AI emulators are increasingly deployed in scientific domains where strong physical theories and benchmarks already exist. This creates a critical interpretability challenge: when a foundation model successfully reproduces known continuum dynamics, it remains unclear whether the model has learned the underlying physical mechanisms or merely memorized statistical correlations. Specifically, researchers need to determine if the model's internal representations correspond to known physical principles, how these representations relate to the model's successes or failures, and whether "different" internal structures are genuinely informative or merely effective shortcuts. The paper addresses these questions using the Walrus foundation model, a 1.3B-parameter transformer trained on diverse physical scenarios, focusing on the gap between learned internal representations and physical intuition.

Methodology
The authors employ a mechanistic interpretability approach guided by physical principles, specifically utilizing a Sparse Autoencoder (SAE) to probe the internal activations of the Walrus model.

• Model and Data: The study focuses on Walrus, trained on nineteen physical scenarios including classical fluids. The analysis is restricted to a 2D shear-flow simulation suite governed by the incompressible Navier-Stokes equations. The dataset varies by Reynolds number ($Re$) and Schmidt number ($Sc$), with specific attention given to three simulations: Sim3 ($Re=10^4, Sc=0.1$), Sim50 ($Re=10^4, Sc=0.2$), and Sim56 ($Re=5\times10^4, Sc=0.1$).
• Physical Grounding: To triage the large feature set generated by the SAE, the authors use enstrophy ($E$), defined as the integral of the squared vorticity, as a physically grounded metric. They also analyze the time derivative of enstrophy ($\dot{E}$) to assess the balance of injection and dissipation.
• SAE Implementation: A Top-K Sparse Autoencoder was trained on activations from the middle block ($l=20$) of the Walrus transformer, which contains 2,816 activations. The SAE uses an expansion factor of 8, resulting in 22,528 features. The training utilized a two-part loss function (Mean Squared Error and an auxiliary loss for dead neurons) to stabilize training and prevent feature collapse.
• Analysis Strategy: The authors rank SAE features based on their Spearman rank correlation with global enstrophy metrics. They then visualize feature activation heatmaps against tracer fields and enstrophy maps to assess spatial coherence and consistency across different simulation setups and timesteps.

Key Results
The investigation yields several findings regarding the consistency and interpretability of the model's internal representations:

1. Feature Identification and Correlation: While a top-ranked feature showed a high correlation ($\rho = 0.85$) with enstrophy in the reference simulation (Sim50), approximately 10% of all 22,528 features exceeded the significance threshold. This suggests that correlation with a single physical metric like enstrophy is insufficient to uniquely isolate mechanistically meaningful features.
2. Intermittent Consistency: The study finds evidence of "piecewise consistency." Subsets of features recur in similar roles across different shear-flow setups (e.g., specific features activating in channel interfaces or vortex regions). However, this structure is highly intermittent; it does not map cleanly onto standard physical decompositions and often fails to persist across timesteps.
3. Diffusivity and Discrepancies: Direct comparisons between the numerical simulation and the Walrus emulator reveal systematic output-level discrepancies. Walrus tends to underpredict energy at mid-wavenumbers and overpredict at large wavenumbers, indicating the production of excessive small-scale structures or "too much" small-scale diffusion.
4. Representational vs. Physical Diffusivity: The loss of feature coherence as simulations evolve may stem from two sources:
   • Physical: Walrus may over-smooth flow fields, causing energy spectra to dissipate structure at scales where the reference simulation retains it.
   • Representational: Even when physical structure remains, the model may encode it across many weakly active features rather than concentrating it in a few interpretable ones.
5. Lack of Stable Mapping: Restricting analysis to top-ranked features within a single simulation does not recover a stable mapping onto familiar physical quantities. Patterns that appear meaningful at early timesteps often dissolve into noise or shift unpredictably as the flow evolves.

Significance and Claims
The paper does not claim to have solved the interpretability of foundation models or to have proven that Walrus has learned the governing equations of fluid dynamics. Instead, it serves as a case study highlighting open questions for the scientific foundation model community:

• Prioritization: It is difficult to robustly prioritize mechanistically meaningful features when a large fraction of features correlates with physical metrics.
• Artifact vs. Structure: There is a significant challenge in separating stable internal structure from analysis artifacts, including limitations inherent to single-layer probing and SAE methodology.
• Benchmarking: Established benchmarks (like output-level accuracy) are insufficient to determine if internal representations are genuinely informative. A model can reproduce dynamics effectively while utilizing internal representations that do not align with human physical intuition.

The authors conclude that claims about what such models have "learned" should be treated cautiously. They suggest that building trust in scientific foundation models may require starting from simpler, more interpretable architectures and scaling up systematically, rather than attempting to reverse-engineer billion-parameter transformers post-hoc. The work underscores the necessity of distinguishing between models that have captured governing relationships and those that have merely memorized statistical regularities.