Mechanistic Interpretability and Causal Feature Steering of Neural Quantum States via Sparse Autoencoders

Technical Summary: Mechanistic Interpretability and Causal Feature Steering of Neural Quantum States via Sparse Autoencoders

Problem Statement
Neural Quantum States (NQS) have emerged as powerful variational ansätze for representing quantum many-body wavefunctions, achieving success in both ground-state representation and real-time dynamics across spin and fermionic systems. However, despite their empirical efficacy, NQS remain largely "black boxes." A fundamental gap exists in understanding their internal mechanisms: trained solely on variational objectives (such as energy minimization) without explicit optimization for specific physical observables, it is unclear how NQS accurately capture quantities like order parameters or correlation functions. This opacity hinders model trust and robustness, particularly in regimes where numerical benchmarks are difficult to obtain. Existing approaches have largely focused on external architectural constraints or parameter analysis, leaving a lack of post hoc analysis regarding how networks internally organize physical information.

Methodology
The authors propose a systematic framework using Sparse Autoencoders (SAEs), a tool from mechanistic interpretability, to decompose the internal activations of NQS into sparse, interpretable features.

1. NQS Architecture: The study utilizes autoregressive transformer-based NQS. The wavefunction amplitude $\psi(\sigma)$ is factorized into a Born probability $p(\sigma)$ and a phase $\phi(\sigma)$. The transformer acts as a causal sequence model, predicting spin configurations site-by-site.
2. SAE Application: The SAE is applied to the residual stream of the final transformer layer. It learns an overcomplete dictionary of feature directions $\{f_k\}$ and sparse coefficients $z_k$ such that activation vectors $h$ are approximated as a sparse linear combination: $h \approx \sum z_k f_k$.
3. Unsupervised Training: Crucially, the SAE is trained entirely unsupervised on the residual stream activations. It receives no physical labels (e.g., magnetization values) and is optimized solely to reconstruct activations with a sparsity penalty ($L_1$ norm).
4. Feature-Observable Correlation: After training, the authors "mean-pool" the sparse coefficients over the spin chain to obtain scalar feature strengths. These are correlated with physical observables (e.g., magnetization, staggered magnetization, half-chain correlators) computed from the same sampled configurations.
5. Causal Intervention (Steering): To establish causality, the authors perform targeted interventions. They rescale the activation strength of a specific feature ($z_{k^*} \to \alpha z_{k^*}$) while keeping others fixed, reconstruct the modified residual stream, and observe the resulting change in the physical observable. This is done while monitoring the variational energy to ensure the intervention remains localized.

Key Contributions and Results

• Discovery of Interpretable Features: In the 1D Transverse-Field Ising Model (TFIM), the SAE autonomously discovered sparse features that correlate nearly perfectly ($|r| > 0.97$) with physical observables, including magnetization ($M_z$), absolute magnetization ($|M_z|$), and half-chain correlators ($C_{zz}$). This discovery occurred across the paramagnetic phase, ferromagnetic phase, and the critical point, despite the SAE having no prior knowledge of these quantities.
• Causal Control: The authors demonstrated that these features play a causal role in the NQS predictions. Rescaling a single feature strength ($\alpha$) smoothly and monotonically steered the corresponding observable (e.g., magnetization) while leaving the variational energy nearly unchanged (relative deviation $< 0.02\%$). This confirms that the features are not merely correlated artifacts but are internal coordinates that control physical predictions.
• Real-Time Dynamics: The framework was extended to real-time dynamics under the TFIM. A single static feature direction was found to track the time-dependent magnetization $M_z(t)$ throughout the entire evolution trajectory. Furthermore, scaling this feature's activation strength across time systematically altered the magnetization dynamics, proving that NQS encode dynamical information in a static internal coordinate system.
• Generalization to 2D Systems: The approach was validated on the 2D Heisenberg Antiferromagnetic model (a non-integrable, 2D system). The SAE identified a feature strongly correlated with the staggered magnetization (a spatially structured order parameter distinguishing sublattices). Steering this feature successfully controlled the staggered magnetization, demonstrating the method's applicability beyond 1D and exactly solvable systems.

Significance and Claims
The paper claims that NQS are not merely functional approximators of wavefunction amplitudes but encode rich, interpretable internal representations of physical information.

• Diagnostic Tool: The approach provides a method to assess whether a trained NQS has learned a physically sensible representation. The presence of features that causally control order parameters serves as evidence of successful learning, while their absence could diagnose model failure.
• Intervention Tool: The work establishes a route for controlled, feature-level manipulation of NQS. By adjusting a single scalar coefficient in the residual stream, one can tune specific physical observables without retraining or fine-tuning the entire network.
• Two-Way Exchange: The authors posit a synergistic relationship between quantum many-body physics and mechanistic interpretability. NQS provide a setting with ground-truth physical observables to test interpretability claims, while interpretability tools offer methods to open the "black box" of variational wavefunctions, potentially leading to more reliable and transparent quantum simulations.

The authors remain modest regarding the scope, noting that this work focuses exclusively on transformer architectures and final-layer residual streams, leaving open questions regarding other architectures (e.g., RBMs, RNNs) and layer-resolved self-attention mechanisms.