1 Introduction↩︎

Neural quantum states (NQS) have recently emerged as a powerful class of variational ansätze for quantum many-body systems [1]–[4]. By leveraging the expressivity of neural networks to parameterize many-body wavefunctions, NQS have achieved remarkable success across spin and fermionic systems, in both ground-state representation and real-time dynamics [4]–[21].

The expressivity of NQS, however, comes at the expense of transparency. Despite their empirical success, NQS remain largely “black boxes”. This raises a fundamental question: trained on variational objectives such as energy minimization alone, how do NQS accurately capture physical observables that have never been explicitly provided as optimization targets? Is such physical information somehow represented inside the network, and if so, where and in what form? Can such representations be extracted, interpreted, or causally manipulated? At present, such questions remain largely unanswered.

This opacity is not only conceptual, but also directly impacts practical aspects of NQS, especially regarding model trust and robustness. In many-body quantum systems, NQS are often used precisely in regimes where obtaining numerical benchmarks can be challenging. In such settings, one would also like to have independent, physically motivated probes on NQS models to assess whether the network has learned a physically sensible representation of the quantum state and whether its predictions can be validated and trusted.

Existing works have mostly approached the opacity of NQS externally, by imposing physically motivated architectural constraints [14], [19], [22]–[30], understanding their capacity to represent quantum states [31]–[36], and analyzing model parameters [37]–[42]. These efforts ask what physical structure should be built into, or can be represented by, neural quantum states, and are essential for designing more physically faithful NQS. However, a post hoc analysis of trained NQS models, which would reveal how networks internally organize physical information, remains missing. Probing the internal representations formed by NQS is important not only for understanding how neural networks represent quantum states, but also for building more informed and reliable variational tools for quantum many-body problems.

In this work, we present a systematic framework to bridge this gap using tools from mechanistic interpretability, a rapidly developing field that seeks to interpret neural networks by analyzing their activations and response to causal interventions [43]–[45]. In particular, we use a sparse autoencoder (SAE) to decompose NQS activations into a sparse set of interpretable features. Originally applied to large language models [46]–[50], SAEs have also successfully extracted various scientific concepts from protein language models and weather predictors [51]–[56].

Compared to these models, NQS offer a particularly ideal setting for interpretability analysis. NQS are not trained to learn patterns from an external dataset, but rather by minimizing physical objectives such as the variational energy. As a result, any structure found inside the network reflects how NQS internally organize information in order to solve quantum many-body problems such as ground state search, as opposed to being inherited from an external dataset. Furthermore, unlike feature labels in many language and scientific models, which are often semantic and difficult to validate uniquely, the relevant concepts in NQS are precisely defined physical observables. This allows feature interpretations to be tested quantitatively, through both observable correlations and causal interventions.

To benchmark our approach, we first analyze the internal representations of transformer-based NQS for both ground state of, and time evolution under, the 1D Transverse-Field Ising Model (TFIM). Applying SAEs to the final-layer residual stream, we extract sparse features that correlate nearly perfectly with physical observables, including the order parameter, absolute magnetization, and half-chain spin correlator. Remarkably, such feature learning is entirely autonomous: the SAE is not, a priori, informed of what features to look for; as a result, the discovery indicates that the model has formed its own internal representations of such physical observables during training. We further establish that the features play a causal role in predicting the physical observables through targeted intervention: rescaling the sparse feature smoothly steers the corresponding observable, while leaving the variational energy almost unchanged. We extend this analysis to the two-dimensional Heisenberg antiferromagnetic model, demonstrating that the approach generalizes beyond one-dimensional and exactly solvable systems.

Our results demonstrate that NQS do not merely approximate wavefunction amplitudes configuration by configuration, but rather organize physical quantities into their internal coordinates. This provides both a diagnostic tool for assessing whether trained NQS models have learned physically meaningful representations, as well as a route toward controlled, feature-level manipulation of NQS. More broadly, our approach opens the door to a two-way exchange between quantum many-body systems and mechanistic interpretability: NQS provide a setting in which interpretability claims and tools can be tested against ground truths; conversely, interpretability offers methods and probes to understand, diagnose, and ultimately improve the reliability of NQS models.

The rest of this work is organized as follows. In Sec. 2, we review the basics of neural quantum states (NQS) and sparse autoencoders (SAE), a tool from the interpretability community that extracts patterns from neural network activations. In Sec. 3, we demonstrate that SAE is able to extract sparse features in NQS residual streams that strongly correlate with, and causally control, physical observables, for both ground state and time-evolved states. We extend our analysis to the 2D Heisenberg Antiferromagnetic model in Sec. 4, demonstrating the method’s applicability beyond simple systems. We summarize our results and discuss directions for future work in Sec. 5.

2 Background ↩︎

2.1 Neural Quantum States↩︎

In this section, we briefly review neural quantum states (NQS), a neural-network-based variational ansätz for quantum many-body wavefunctions. For concreteness, we focus on systems with spin-\(1/2\) degrees of freedom. The computational basis consists of tensor products of Pauli-\(z\) eigenstates, \(| \boldsymbol{\sigma} \rangle = | \sigma_1,\sigma_2,\ldots,\sigma_L \rangle\), where \(L\) is the system size and \(\sigma_i\in\{-1,+1\}\) denotes the eigenvalue of the Pauli-\(z\) operator on site \(i\).

A quantum many-body wavefunction, therefore, can be viewed as a mapping from the configuration space (here bit-strings of length \(L\): \(\{-1, 1 \}^L\)) to the associated complex amplitude: \[\psi: \boldsymbol{\sigma} \rightarrow \langle \boldsymbol{\sigma} |\psi\rangle \in \mathbb{C}.\]

The central idea of NQS is to leverage the expressivity of neural networks to parametrize the function \(\psi\), while using only a polynomial number of parameters \(\boldsymbol{\theta}\). Among the neural network architectures proposed for this purpose, we focus on autoregressive quantum states [9]–[11], [17], [18], and in particular transformer-based representations. The transformer architecture used in this work is illustrated in Fig. 1 (a), with implementation details in Appendix 6.

In practice, the autoregressive ansätze represent the wavefunction amplitude \(\psi(\boldsymbol{\sigma})\) as a Born probability \(p(\boldsymbol{\sigma})\) and a phase \(\phi(\boldsymbol{\sigma})\), which are output separately: \[\psi_{\boldsymbol{\theta}}(\boldsymbol{\sigma}) = \sqrt{p_{\boldsymbol{\theta}}(\boldsymbol{\sigma})}\; e^{i\phi_{\boldsymbol{\theta}}(\boldsymbol{\sigma})}. \label{eq:mod95phase}\tag{1}\]

In the autoregressive transformer wavefunction ansätze, the Born probability distribution is factorized into a product of conditional probabilities, \[p_{\boldsymbol{\theta}}(\boldsymbol{\sigma}) = \prod_{i=1}^{L} p_{\boldsymbol{\theta}} (\sigma_i\,|\,\sigma_1,\ldots,\sigma_{i-1}), \label{eq:factorization}\tag{2}\] where \(p(\sigma_i\,|\,\sigma_1,\ldots,\sigma_{i-1})\) is the conditional probability of the \(i\)th spin, given the configuration of all preceding spins. The transformer therefore acts as a causal sequence model for spin configurations: at each site, the distribution of the next spin is predicted, conditioning on all previously generated spins. The phase is also decomposed analogously.

Since each conditional probability is normalized, the autoregressive factorization ensures that the full Born distribution is normalized by construction. This factorization also enables exact autoregressive sampling: each configuration is generated sequentially site-by-site, via sampling from the corresponding conditional probabilities. All configurations are independently drawn from the Born distribution \(|\psi_{\boldsymbol{\theta}}|^2\), avoiding Markov-chain inter-sample autocorrelation.

From such samples, expectation values of an observable \(O\) can be estimated by Monte Carlo, \[\langle O\rangle = \mathop{\mathbb{E}}_{\boldsymbol{\sigma}\sim |\psi_{\boldsymbol{\theta}}|^2} \bigl[O_{\mathrm{loc}}(\boldsymbol{\sigma})\bigr],\] where the local estimator \[O_{\mathrm{loc}}(\boldsymbol{\sigma}) = \sum_{\boldsymbol{\sigma}'} \langle \boldsymbol{\sigma} |O |\boldsymbol{\sigma'}\rangle \frac{\psi_{\boldsymbol{\theta}}(\boldsymbol{\sigma}')}{\psi_{\boldsymbol{\theta}}(\boldsymbol{\sigma})} \label{eq:local95estimator}\tag{3}\] can be efficiently evaluated when \(O\) is sparse in the computational basis.

In particular, to represent the ground state, the NQS parameters are optimized to minimize the variational energy \[E(\boldsymbol{\theta}) = \frac{\langle \psi_{\boldsymbol{\theta}}|H | {\psi_{\boldsymbol{\theta}}} \rangle }{\langle \psi_{\boldsymbol{\theta}} |\psi_{\boldsymbol{\theta}}\rangle}. \label{eq:var95energy}\tag{4}\] In practice, both the energy and its gradient are estimated stochastically during training.

NQS can also be used to simulate real-time dynamics through time-dependent parameters \(\boldsymbol{\theta}(t)\). The goal is to approximate exact Schrödinger evolution within the variational manifold. The network parameters are optimized so that the variational residual \[R(\dot{\boldsymbol{\theta}}) = \Bigl\|\, \sum_k \dot{\theta}_k\, \partial_{\theta_k}| \psi_{\boldsymbol{\theta}} \rangle + iH| \psi_{\boldsymbol{\theta}} \rangle \,\Bigr\| \label{eq:residual}\tag{5}\] is minimized at each time step. This yields the time-dependent variational principle (TDVP) equation of motion, \[\sum_{k'} S_{kk'}\,\dot{\theta}_{k'} = -i\,F_k, \label{eq:tdvp}\tag{6}\] where \(S_{kk'} = \langle O_k^{*}O_{k'}\rangle -\langle O_k^{*}\rangle\langle O_{k'}\rangle\) is the covariance matrix, and \(F_k = \langle O_k^{*}E_{\mathrm{loc}}\rangle -\langle O_k^{*}\rangle\langle E_{\mathrm{loc}}\rangle\) is the force vector. The quantity \(O_k = \partial_{\theta_k}\log\psi_{\boldsymbol{\theta}}\) is the logarithmic derivative of the wavefunction. Both \(S_{k k'}\) and \(F_k\) are Monte-Carlo estimates over samples from the instantaneous Born distribution \(|\psi_{\boldsymbol{\theta}(t)}|^2\). In what follows, we apply sparse autoencoders to analyze the internal representations formed by NQS in both settings.

2.2 Sparse Autoencoders↩︎

A sparse autoencoder (SAE) embodies a dictionary learning approach [50] that decomposes activation vectors of neural networks into sparse combinations of interpretable learned feature directions. This decomposition is motivated by the superposition hypothesis [57], which postulates that neural networks with residual-stream dimension \(d\) can encode far more than \(d\) distinct features by superimposing the features in activation space. The SAE approach is designed to disentangle this superposition; it has successfully identified many human-interpretable features in large language models [46]–[48] and more recently scientific concepts in the natural sciences. For example, SAEs have extracted biological motifs such as catalytic-site activation in protein language models [51]–[54], cyclone-related features in weather models [55], and enstrophy-related features in fluid-mechanical foundation models [56].

In this work, the inputs to the SAE are activation vectors within the residual stream of the transformer NQS, \(\mathbf{h}^{(m)} \in \mathbb{R}^{d_\text{model}}\). Here \(d_\text{model}\) is the dimension of the transformer residual-stream and \(m\) indexes the input. The SAE learns a dictionary \(\{ \mathbf{f}_k \}\), where \(k \in \{1, 2, \dots, d_{\text{SAE}} \}\). Each “feature" \(\mathbf{f}_k \in \mathbb{R}^{d_\text{model}}\) is unit-normalized, and the dictionary is overcomplete, in the sense that its number of elements exceeds the residual stream dimension, \(d_\text{SAE} > d_\text{model}\). The SAE is trained to find both the feature directions \(\mathbf{f}_k\) and sparse, input-dependent coefficients \(z_k^{(m)}\), such that each activation vector \(\mathbf{h}^{(m)}\) can be approximated as \[\mathbf{h}^{(m)} \simeq \sum_k z_k^{(m)} \, \mathbf{f}_k,\] with only a small number of non-zero coefficients \(z_k^{(m)}\). Thus, the SAE represents each activation vector as a sparse linear combination of feature directions.

Architecturally, the SAE is an autoencoder with a single hidden layer of dimension \(d_\text{SAE}\). In contrast to conventional autoencoders [58], which generally compress inputs down to lower dimensions, the hidden layer in SAEs is deliberately higher dimensional than its input. The SAE first lifts an activation vector \(\mathbf{h^{(m)}}\in\mathbb{R}^{{d_\text{model}}}\) to \(\mathbf{z}^{(m)}\in\mathbb{R}^{d_{\mathrm{SAE}}}\) via an encoder: \[\mathbf{z}^{(m)} = \sigma\!\left(W_{\mathrm{enc}} \, \mathbf{h}^{(m)}+\mathbf{b}_{\mathrm{enc}}\right), \label{eq:SAE95encoding}\tag{7}\] where \(W_{\mathrm{enc}}\) is the encoder weight matrix, \(\mathbf{b}_{\mathrm{enc}}\) the encoder bias, and \(\sigma\) a pointwise non-linear activation function. The SAE then projects \(\mathbf{z}^{(m)}\) back to a reconstructed activation \({\mathbf{h'}^{(m)}}\in\mathbb{R}^{d_\text{model}}\) via a decoder: \[{\mathbf{h'}}^{(m)} = W_{\mathrm{dec}}\,\mathbf{z}^{(m)} + \mathbf{b}_{\mathrm{dec}} = \sum_{k=1}^{d_{\mathrm{SAE}}} z_k^{(m)}\,\mathbf{f}_k + \mathbf{b}_{\mathrm{dec}}, \label{eq:SAE95decoding}\tag{8}\] where the columns of \(W_{\mathrm{dec}}\) are normalized feature directions \(\mathbf{f}_k\). The encoder weights \(W_{\mathrm{enc}}\), both biases \(\mathbf{b}_{\mathrm{enc}},\mathbf{b}_{\mathrm{dec}}\), and features (or equivalently \(W_{\mathrm{dec}}\)) are all learned during training. The concept and architecture of SAEs are illustrated in Fig. 1 (b).

The SAE is trained by minimizing the following loss function summed over all residual-stream activations: \[\mathcal{L}_\text{SAE} = \sum_m \left( \|\mathbf{h}^{(m)} - \mathbf{h}'^{(m)}\|_2^2 + \lambda\|\mathbf{z}^{(m)}\|_1 \right),\] where \(\| \cdot \|_2\) is the Euclidean (\(L_2\)) norm, and \(\|\cdot \|_1\) is the \(L_1\) norm. The first term enforces faithful reconstruction of the activation vectors, while the second term imposes a sparsity penalty that forces all but a few entries of \(\mathbf{z}^{(m)}\) to vanish. The competition between the two terms in the loss function ensures that the network finds the right balance between reconstruction fidelity and sparsity. The hyperparameter \(\lambda\) controls the relative scale between the two terms. Details about the SAE hyperparameters are in Appendix 7.

Importantly, the SAE is trained in an entirely unsupervised manner on the activation vectors, without any access to physical labels. The feature directions must be discovered solely from the transformer’s internal representations. As we demonstrate in the next two sections, despite this seemingly blind decomposition, the SAE recovers features that correlate strongly with physical observables.

3 Results: Transverse Field Ising Model↩︎

We start with the one-dimensional Transverse-Field Ising Model (TFIM) [59], with the Hamiltonian: \[H = -J\sum_{i=1}^{L} \sigma^z_i \sigma^z_{i+1} - h\sum_{i=1}^{L} \sigma^x_i,\] where \(\sigma^x_i\) and \(\sigma^z_i\) are the Pauli \(x\)- (\(z\)-) operators on site \(i\), \(J\) is the nearest-neighbor Ising coupling, and \(h\) the transverse field strength. We impose periodic boundary conditions for a chain of \(L=40\) sites. The TFIM has a critical point at \(h/J = 1\): for \(h < J\), the system is ferromagnetically ordered; with \(h > J\), the system is paramagnetic. We set \(J \equiv 1\) throughout the rest of this work.

The TFIM provides an ideal starting point for mechanistically interpreting NQS. The model is exactly solvable and well-understood, and therefore provides a controlled setting in which the ground truth is known. For example, the ground state energy density can be exactly solved given \(J\) and \(h\), which allows us to assess convergence of NQS training. Furthermore, because the relevant physical observables are known, we can directly test whether the features autonomously extracted by SAE genuinely correspond to physically meaningful quantities.

In this section, we study transformer-NQS representations of both the TFIM ground state and real-time dynamics. In the next section, we turn to a two-dimensional, non-integrable model, the 2D Heisenberg Antiferromagnet, to demonstrate that our approach is not restricted to exactly solvable one-dimensional systems.

3.1 Feature Identification for Ground State↩︎

We begin by applying the SAE analysis to transformer-based representations of the TFIM ground state. As introduced in Sec. 2.1, the transformer is trained to model the ground state by stochastically minimizing the variational energy. Details about the transformer architecture and model parameters are in Appendix 6 and Appendix 7, respectively.

After training, we autoregressively sample \(M\) spin configurations \(\boldsymbol{\sigma}^{(m)}\) (with \(m\in\{1,2,\dots,M\}\)) from \(|\psi(\boldsymbol{\sigma})|^2\). Because the wavefunction is represented autoregressively, these spin samples are drawn from the Born distribution independently, without Monte-Carlo autocorrelation. For each configuration, we store the final-layer residual stream activations of the transformer as a matrix \(\mathbf{h}^{(m)}\in\mathbb{R}^{L\times d_\text{model}}\), where \(L\) is the system size and \(d_\text{model}\) is the residual-stream dimension. The row vector \(\mathbf{h}_i^{(m)}\) is the residual stream activation on site \(i\) for the \(m\)th configuration. We focus on the final-layer residual stream, because these activations enter the unembedding layer (Fig. 1 (a)) and directly determine the model’s output.

The SAE is trained on the collection of site-level residual stream activations, using the architecture and loss function as described in Sec. 2.2. After training, the SAE decomposes each activation into a sparse linear combination of feature directions \[\mathbf{h}_i^{(m)} \approx \sum_k z_{ik}^{(m)} \, \mathbf{f}_k,\] where the coefficient \(z_{ik}^{(m)}\) denotes the strength of the \(k\)th feature, on the \(i\)th site, for the \(m\)th configuration.

To associate each SAE feature with a single number per spin configuration, we “mean-pool” over the spin chain and average the activation over all sites, \[Z_k^{(m)} = \frac{1}{L}\sum_{i=1}^L z_{ik}^{(m)}.\] Thus, \(Z_k^{(m)} \in \mathbb{R}\) is a scalar feature strength. Collecting \(Z_k^{(m)}\) across all configurations yields a vector \(\mathbf{Z}_k \in \mathbb{R}^M\), which can be directly correlated against physical observables evaluated on the same set of sampled spin configurations.

For the TFIM, we consider three physically motivated observables. Naturally, we consider the order parameter of the model, namely the magnetization \[\begin{align} M_z(\boldsymbol{\sigma}) &= \frac{1}{L} \sum_{i=1}^L \sigma_i^z, \end{align}\] together with the absolute magnetization \(|M_z|\).

We also probe long range order using the translationally averaged half-chain correlator: \[C_{zz}(L/2) = \frac{1}{L} \sum_i \sigma_i^z \sigma_{i+L/2}^z,\] where site indices are modulo \(L\) due to periodic boundary conditions. \(C_{zz}\) probes spin-spin order and saturates in the ferromagnetically ordered phase.

For each observable \(O \in \{M_z,\,|M_z|,\, C_{zz}(L/2)\}\), we evaluate \(O^{(m)} \equiv O(\boldsymbol{\sigma}^{(m)})\) for all sampled spin configurations. We collect the results, forming vectors \(\mathbf{O} \in \mathbb{R}^M\). We then compute the correlation between each mean-pooled SAE feature \(\mathbf{Z}_k\) and physical observables \(\mathbf{O}\): \[r_k{(O)} = \mathrm{corr}(\mathbf{Z}_k,\, \mathbf{O}).\]

We use the Pearson correlation coefficient, which is defined as: \[r_k(O) := \frac{\sum_m \bigl(Z_k^{(m)}-\overline{Z_k}\bigr)\bigl(O^{(m)}-\overline{O}\bigr)}{\sqrt{\sum_m \bigl(Z_k^{(m)}-\overline{Z_k}\bigr)^2}\,\sqrt{\sum_m \bigl(O^{(m)}-\overline{O}\bigr)^2}},\] where the overline denotes an average over the \(M\) samples. A value \(|r_k(O)| = 1\) indicates perfect linear correlation between the \(k\)th feature and the observable \(\mathbf{O}\), while \(r_k(O)=0\) indicates no linear correlation. For each observable \(O\), we compute \(r_k(O)\) for all \(d_\text{SAE}\) features and rank the features by \(|r_k(O)|\). We denote the maximally correlated feature by \(k^*\). We emphasize that we do not claim a one-to-one correspondence between SAE features and physical observables. In this work, we focus on the maximally correlated feature direction \(k^*\) because our primary goal is to establish that the residual stream contains at least one interpretable direction corresponding to physical observables.

As shown in Fig. 2, the SAE autonomously discovers sparse features that correlate strongly with physical observables across the TFIM phase diagram. In the paramagnetic phase (\(h=1.5)\), feature 87 correlates with the magnetization \(M_z\) at \(r=0.99\). In the ferromagnetic phase (\(h=0.5\)), the autonomously discovered feature 106 correlates with the half-chain correlator \(C_{zz}\) at \(r=-0.97\), despite \(C_{zz}\) being a non-linear observable in the spin configuration. The negative sign of the correlation indicates that the learned feature direction is aligned with \(-C_{zz}\) instead of \(C_{zz}\). Remarkably, even at the critical point \(h=1.0\), the SAE method successfully discovers a feature (\(f193\)) that strongly correlates with \(|M_z|\).

These remarkably high correlations suggest that the SAE features act as nearly linear internal coordinates for the corresponding physical observables. Importantly, neither the NQS nor the SAE is trained using these observables as labels: the NQS is trained only through variational energy minimization, while the SAE is trained only to reconstruct residual-stream activations in a sparse way. To confirm that the observed correlations are not artifacts of the transformer or the SAE architectures, we repeat the same analysis for a randomly initialized NQS in Appendix 9. In that case, no feature shows any meaningful correlation with physical observables. Thus, the feature-observable alignment emerges from the trained NQS representation rather than from architectural bias.

3.2 Observable Steering through Sparse Features↩︎

In the previous section, we demonstrated that SAEs identify sparse features that correlate strongly with physical observables across different phases of the TFIM. Although these correlations are highly suggestive, they do not by themselves establish that the features actually control the observables predicted by the NQS. To probe causation directly, we perform targeted intervention [60] on the NQS: we rescale the top-correlating (\(k^*\)th) feature and measure how the corresponding observable responds. Related interventions have been successful in natural language processing, steering the behaviors of large language models such as GPT models and Claude 3 Sonnet [49], [61]. Similar ideas have also been applied in scientific domains to amplify interpretable concepts in protein and weather models [51], [55].

In our setting, we carry out intervention through activation steering [62], [63]. For each configuration \(\boldsymbol{\sigma}^{(m)}\), we modify the corresponding final-layer residual stream activations \(\mathbf{h}^{(m)}\) in three steps. First, the activation \(\mathbf{h}^{(m)}\) is encoded to the sparse activation strengths \(\mathbf{z}^{(m)}\) by the SAE encoder (Eq. 7 ). Next, the activation strength of the top-correlated feature is rescaled, \(z_{k^*} \rightarrow \widetilde{z_{k^*}} \equiv \alpha z_{k^*}\), while all other activations \(k \neq k^*\) remain unchanged. Finally, the modified sparse activation vector \(\widetilde{\mathbf{z}}\) is passed through the SAE decoder (Eq. 8 ), to reconstruct a modified residual stream \(\widetilde{\mathbf{h}}\), which is then passed through the unembedding layer to define a modified NQS \(\widetilde{{\psi}}\). The scale factor \(\alpha\) therefore controls the strength of this intervention: \(\alpha = 1\) leaves all features unchanged; while \(\alpha < 1\) and \(\alpha > 1\) suppress and amplify the selected feature \(k^*\), respectively. The concept of feature-steering procedure is illustrated in Fig. 1 (c).

Fig. 3 shows how the magnetization \(M_z\) and the half-chain correlator \(C_{zz}(L/2)\) respond to the feature scale \(\alpha\) in the paramagnetic \((h=1.5)\) and ferromagnetic \((h=0.5)\) phases, respectively. In both cases, rescaling the top-correlated feature produces a smooth and monotonic shift in the corresponding physical observables. At the same time, the variational energy remains almost unchanged, with a relative energy deviation \(|\Delta E/E|\) remaining below \(0.02 \%\) in both cases. This indicates that the intervention is localized in representation space: modifying a single sparse feature can steer the target observable alone. The identified features \(k^*\) therefore do not merely correlate with physical observables, but play a causal role in how these observables are represented by the NQS.

There is a crucial distinction between previous intervention studies and our case. For large language models and protein language models, evaluating the effect of steering requires choosing a behavioral metric [43]–[45], which can by itself be subjective. In contrast, interventions on sparse features in NQS can be objectively measured by expectation values of unambiguously defined physical observables, with no interpretive ambiguity.

We note that there is an important subtlety about sampling. To isolate the effect of feature steering itself, we evaluate the steering experiment using the same set of sampled spin configurations \(\{ \boldsymbol{\sigma}^{(m)}\}\) for all values of \(\alpha\). All configurations are thus drawn from the original Born distribution \(|\psi|^2\). As a result, \(\alpha\) cannot be taken arbitrarily large or small; otherwise the modified Born distribution \(|\widetilde{\psi}|^2\) would differ significantly from \(|\psi|^2\), and the fixed sample set will no longer provide reliable estimates of expectation values. In Appendix 8, we quantify the range of admissible \(\alpha\) and show that the values of \(\alpha\) considered in Fig. 3 are well within the regime where importance-sampling estimates are reliable.

Such feature-based steering also has potential practical applications. For example, in situations where the exact ground state is unknown, but certain observables are available experimentally, the sparse features identified by SAE provide an interpretable handle for adjusting the NQS toward desired physical behavior, without changing diagnostics such as the variational energy. Importantly, this control acts directly on the residual stream: once a sparse feature strongly associated with a target observable has been identified, the observable can be tuned by simply adjusting a single scalar coefficient immediately before the unembedding layer. The post-training, low-dimensional intervention differs from other approaches such as fine-tuning [11], [38], [64]–[66], which would require additional optimization and backward passes and can be substantially more expensive for larger models.

3.3 Real-Time Dynamics↩︎

In this section, we extend our analysis to transformer NQS that capture real-time dynamics. Compared to the ground state analysis, this setting is more challenging, as both the quantum state and associated Born distribution evolve continuously in time. It is not obvious whether a fixed set of sparse features can track physical observables throughout the entire trajectory.

To test this, we study real-time dynamics under the TFIM. The system is initialized in the ferromagnetically aligned state \[|\psi_0\rangle = |\!\downarrow\downarrow\cdots\downarrow\rangle,\] and evolved under the TFIM Hamiltonian in the paramagnetic regime with \(h = 1.5\). As discussed in Sec. 2.1, the time-dependent state \(| \psi(t) \rangle\) is represented by a transformer NQS with time-dependent parameters \(\theta(t)\), optimized using the TDVP over the interval \(t \in [0, T]\) with \(TJ = 0.5\).

The SAE analysis proceeds analogously to the ground state case. At each time step \(t\), we draw \(M\) spin configurations \(\boldsymbol{\sigma}^{(m)}(t)\) from the instantaneous Born distribution \(|\psi(\boldsymbol{\sigma}; t)|^2\) and collect the final-layer residual stream activations \(h^{(m)}(t)\). We then pool together the sampled spin configurations and activation vectors across all times, and train a single SAE on the full dynamical dataset. Crucially, the SAE dictionary is not trained separately at each time step. Instead, the learned feature directions \(\mathbf{f}_k\) remain fixed, while the activation strengths \(\mathbf{z}_k(t)\) vary over time.

For each time step, we then analyze the correlation between physical observable \(\mathbf{O}(t)\) and mean-pooled activations \(\mathbf{Z}_k(t)\) for all features, just as in the ground state setting (Sec. 3.1). This construction asks whether a sparse feature continues to align with the same observable during the trajectory. Features that maintain a high correlation \(|r(t)|\) over time can therefore be interpreted as the internal coordinates that the NQS uses to represent the dynamics of that given observable.

As shown in Fig. 4 (a), a single sparse feature (\(f24\)) maintains a strong correlation with the instantaneous magnetization \(M_z(t)\) throughout the entire evolution. Specifically, the Pearson correlation \(|r(\tilde{z}_{k^*}, M_z)|\) remains above 0.92 for all \(tJ \in [0, 0.5]\), with an average correlation value close to 0.95. Thus, this feature not only captures the magnetization near the initial polarized state, but also continues to track the magnetization faithfully, even as the magnetization evolves significantly away from its initial value. Remarkably, time and magnetization are not explicitly input into the SAE, which must infer from the residual-stream activation vectors alone which direction is associated with the magnetization dynamics.

We next test whether feature-steering is applicable to the dynamical state. We intervene on the identified SAE feature by scaling its activation strength across all times, \(\mathbf{z}_{k^*}(t) \rightarrow \alpha \,\mathbf{z}_{k^*}(t)\), while leaving all remaining feature activations unchanged. Fig. 4 (b) shows the resulting time-dependent magnetization \(\langle M_z(t) \rangle\) for three representative values of the scale factor: \(\alpha = 1\) (original state), \(\alpha = -1\) (feature suppressed), and \(\alpha = 3\) (feature amplified). Rescaling the single feature produces a clear and systematic effect on the magnetization: amplifying \(f24\) accelerates the departure from the initial magnetization, while suppressing \(f24\) keeps the state more polarized. The separation between the resulting trajectories is persistent throughout the full time interval, indicating that the feature captures magnetization dynamics over the entire evolution instead of at isolated times.

These results demonstrate that the SAE analysis is not limited to static wavefunction representations. Even when the NQS parameters evolve in time, the SAE is able to learn a static basis, whose time-dependent activation strengths both correlate with and causally steer physical observables over the time interval. This suggests that the transformer NQS encodes dynamical physical information in terms of a static internal coordinate system. More generally, the constant scaling considered here is the simplest intervention protocol. One could, for example, introduce a time-dependent control function \(\alpha(t)\), and implement a more refined steering of the relevant direction, \(\mathbf{z}_{k^*}(t) \rightarrow \alpha(t) \mathbf{z}_{k^*}(t)\). In such cases, the sparse feature would serve as a dynamical control knob for targeted manipulation of time-dependent neural quantum states.

4 2D Model: Heisenberg Antiferromagnet↩︎

While the TFIM provides a controlled setting for benchmarking our approach, it is also important to test whether the SAE analysis extends beyond one-dimensional, integrable models. In this section, we turn to a more challenging quantum many-body system, the spin-\(1/2\) Heisenberg Antiferromagnet (AFM) model in two dimensions [67], [68]. This model is a paradigmatic model for quantum magnetism. Its ground state exhibits antiferromagnetic long-range correlations, as well as strong quantum fluctuations, making it a significantly more challenging benchmark for testing whether SAE features can identify physically meaningful observables.

Concretely, we consider the Heisenberg AFM on a square lattice of size \(L_x \times L_y\) with periodic boundary conditions. The Hamiltonian is: \[H = J \sum_{\langle ij \rangle} \left( \sigma_i^x \sigma_j^x + \sigma_i^y \sigma_j^y + \sigma_i^z \sigma_j^z \right),\] where \(\langle ij \rangle\) denotes nearest-neighbor bonds on the square lattice, \(\sigma_i^x, \sigma_i^y, \sigma_i^z\) are Pauli operators on site \(i\). We set \(J > 0\), corresponding to antiferromagnetic exchange.

The order parameter is the staggered magnetization, which is defined as: \[M_\mathrm{stag} = \frac{1}{L} \sum_i (-1)^{x_i + y_i} \sigma_i^z,\] where \(i = (x_i, y_i)\) labels the site coordinate on the square lattice, \(L = L_x L_y\) is the total number of sites, and the factor \((-1)^{x_i+y_i}\) assigns opposite signs to the two sublattices.

As in the TFIM case, we train a transformer NQS to approximate the ground state, through minimizing the variational energy, for a lattice of size \(6 \times 6\). We determine training convergence by benchmarking the variational energy against the ground state energy density determined from Quantum Monte Carlo (QMC) [69]. The SAE analysis proceeds in the same way as Sec. 3.1: after training, we sample \(M\) spin configurations and collect the corresponding final-layer residual-stream activation vectors; we then train an SAE to extract sparse features from the activations; and finally compute the Pearson correlation between \(\mathbf{Z}_k\) and \(\mathbf{M}_\text{stag}\), ranking the features by the correlation strength.

Fig. 5 (a) shows that the SAE identifies a sparse feature (\(f132\)) that is strongly correlated with the staggered magnetization, with a correlation coefficient \(|r|=0.97\) (the negative slope means that the feature is aligned with \(-M_\text{stag}\) instead of \(+M_\text{stag}\)). To probe whether the activation strengths causally control \(M_\text{stag}\), we follow the same activation-steering procedure used in Sec. 3.2: rescaling the activation strength of \(f132\) by \(\alpha\), while leaving all other feature strengths unchanged, and modifying the residual stream of transformer NQS. As shown in Fig. 5 (b), increasing and decreasing the strength of the top-correlating feature monotonically changes the predicted staggered magnetization, establishing that the SAE again successfully identifies a sparse feature that tracks the staggered order parameter.

Our results establish that the interpretability approach holds beyond simple, one-dimensional systems, and also beyond simple quantities in terms of spin configurations. Indeed, compared to magnetization in the TFIM, the staggered magnetization is not a uniform sum over all sites, but a spatially structured quantity that distinguishes the two sublattices. The fact that a single SAE feature tracks this quantity indicates that the transformer NQS can internally form representations for quantities beyond global spin information. Even in a two-dimensional interacting quantum magnet, the trained NQS model organizes physically relevant information in sparse residual stream directions, by encoding the sublattice pattern associated with the antiferromagnetic long-range order into its activation patterns.

5 Discussion↩︎

Do neural networks dream of Schrödinger’s Cat? In this work, we show that for neural quantum states, the answer is affirmative. Using sparse autoencoders to analyze the residual-stream of transformer NQS, we have successfully identified physically interpretable features from the NQS’s internal representation. The autonomously identified features not only strongly correlate with physical observables, but in fact causally control them. By adjusting the strengths of the top-correlating feature activations, we can tune the corresponding observable in a controlled way. Importantly, our analysis holds for both real-time dynamics, where a static feature tracks time-dependent magnetization across the entire time window, and for a two-dimensional, non-integrable quantum antiferromagnet.

Our results provide evidence that NQS models, trained solely on variational objectives, form internal representations that encode relevant physical quantities of the system. Furthermore, our approach provides a diagnostic tool for NQS. If order parameters and relevant observables are represented by internal coordinates, this provides evidence that the NQS could have learned a physically sensible representation; conversely, if such correlations are absent, it could help diagnose failure modes of the model and potentially point to better architectures or physical constraints.

Our findings open up a frontier of future directions. In this work, we have focused on the transformer architecture exclusively. Whether other neural architectures, such as Restricted Boltzmann Machines (RBM) [4] or recurrent neural networks [18], organize physical information in a similar way remains open. Such comparisons could help identify architectural choices that make NQS more interpretable or physically faithful.

The SAE model in this work also only decomposes activation vectors from the final-layer residual stream. This is a natural starting point, as these activations feed directly into the unembedding layer; however, a systematic, layer-resolved analysis of the NQS and especially the self-attention mechanism could also carry physical information. For instance, do attention weights distinguish short-range from long-range order across phase transitions? Probes of this kind, using interpretability tools, could potentially provide novel diagnostics of phase boundaries.

As SAE was originally designed for foundation models, it would be interesting to see if our analysis could provide further insight on foundational NQS models, which are trained across families of Hamiltonians or driving protocols [65], [66], [70], [71]. For example, certain sparse features could represent order parameters across phase transitions. If such features exist, they could provide insights into how foundation NQS models organize information across families of Hamiltonians rather than individual wavefunctions.

More broadly, our work suggests a two-way exchange between the communities of neural quantum states and mechanistic interpretability. NQS provide a setting where interpretability claims can be tested against clearly defined physical observables and exact results; conversely, mechanistic interpretability provides tools for opening the black box of variational neural wavefunctions. As NQS take on a larger role in quantum simulations, understanding their internal representations will become essential for diagnosing failures, building trust, and potentially even extracting novel physical insights from trained models.

Acknowledgments↩︎

We thank Junkai Dong, Toni Liu, Yang Peng, and Miles Stoudenmire for helpful discussions, and especially Yang Peng and Junkai Dong for reading through the manuscript. The authors acknowledge the use of ChatGPT 5.5 and Sonnet 4.6 for code writing and manuscript polishing. ZQ gratefully acknowledges support from the Simons Center for Geometry and Physics, Stony Brook University at which some of the research for this work was performed during the program Complexity, Information, and Tractable Simulations of Quantum Many-body Dynamics.

6 Details of Transformer Architecture↩︎

The transformer architecture [72] has recently become a powerful ansätz for quantum wavefunctions [10], [11], [70], [71]. Its self-attention mechanism allows the wavefunction representation to capture long-range correlations in many quantum states. In this work, we use a decoder-only autoregressive transformer as the NQS architecture.

As illustrated in Fig. 1 (a), inputs to the transformer are spin configurations \(\boldsymbol{\sigma} = (\sigma_1, \sigma_2, \dots, \sigma_L)\); these are the transformer’s tokens. The spin tokens are first lifted to a latent space with dimension \(d_\text{model}\), via an embedding operation: \[\mathbf{e}_i = W_E \, \sigma_i,\] where \(W_E\) is an embedding weight matrix shared across all sites. Each site is also augmented with a learnable positional encoding \(\mathbf{p}_i \in \mathbb{R}^{d_\text{model}}\). The input to the transformer blocks is therefore: \[\mathbf{x}_i = \mathbf{e}_i +\mathbf{p}_i, \quad i\in \{1, 2, \dots, L\}.\]

The sequence \(\{\mathbf{x}_i\}_{i=1}^L\) is then passed through \(L_T\) transformer layers. Each layer consists of two parts: a multi-head self-attention layer followed by a feedforward multi-layer perceptron, with residual connections around both. For notational simplicity, we suppress the transformer-block index in the following.

Let \(X \in \mathbb{R}^{L \times d_\text{model}}\) denote the residual stream entering a transformer block. The first sublayer is masked multi-head self-attention. For each attention head \(a=1,\ldots,n_h\) (where \(n_h\) is the number of heads), the normalized residual stream is projected into queries, keys, and values: \[Q^{(a)} = X W_Q^{(a)}, \quad K^{(a)} = X W_K^{(a)}, \quad V^{(a)} = X W_V^{(a)} , \label{eq:self95attn95proj}\tag{9}\] where \[W_Q^{(a)}, W_K^{(a)}, W_V^{(a)} \in \mathbb{R}^{d_\text{model} \times d_h}\] are the projection weight matrices, and \(d_h = d_\text{model}/n_h\) is the dimension per attention head. At each position, the query determines what information that position is seeking; the key determines whether this position is relevant to queries from all other positions; and the value contains the information that can be passed to other positions.

For each head, we compute the attention score \[S^{(a)}= \frac{ Q^{(a)} {K^{(a)}}^T}{ \sqrt{d_h}}+ \mathcal{M}, \label{eq:attention95scores}\tag{10}\] where \(\mathcal{M}\in\mathbb{R}^{L\times L}\) is the mask: \[\mathcal{M}_{ij} = \begin{cases} 0, & j \leq i, \\ -\infty, & j > i. \end{cases} \label{eq:causal95mask}\tag{11}\] The mask is causal in the sense that it sets the weight of every future position to zero, so the representation at position \(i\) can only use information from previous spin positions \(j\leq i\).

The attention matrix is obtained by applying a row-wise softmax, \[A^{(a)}_{ij} = \frac{ \exp \left(S^{(a)}_{ij}\right)}{ \sum_{m=1}^{L} \exp \left(S^{(a)}_{im}\right)}. \label{eq:attention95weights}\tag{12}\]

The output of head \(a\) is a weighted sum of value vectors, \[O^{(a)} = A^{(a)} V^{(a)} \in \mathbb{R}^{L \times d_h}. \label{eq:attention95output95single95head}\tag{13}\] The outputs of all heads are concatenated and projected back to the embedding dimension: \[\mathrm{Attn}(X) = \mathrm{Concat} \left( O^{(1)},\ldots,O^{(n_h)} \right) W_O + \mathbf{b}_O, \label{eq:self95attn95concat}\tag{14}\] where \(W_O\in\mathbb{R}^{d_\text{model}\times d_\text{model}}\) and \(\mathbf{b}_O\in\mathbb{R}^{d_\text{model}}\). The self-attention result is added back to the input, \(X' = X + \text{Attn}(X)\).

The second sublayer is a position-wise feedforward network that mixes the internal feature channels. For the input matrix \(X'\in\mathbb{R}^{L\times d_\text{model}}\), the feedforward network is \[\mathrm{FFN}(X') = \sigma \left( X' W_1 + \mathbf{b}_1 \right) W_2 + \mathbf{b}_2, \label{eq:ffn}\tag{15}\] where \(W_1\in\mathbb{R}^{d_\text{model}\times d_f}, \,W_2\in\mathbb{R}^{d_f\times d_\text{model}}\) are the weight matrices of each layer. Here \(d_f\) is the feedforward hidden dimension. The feedforward network first expands the channel dimension from \(d_\text{model}\) to \(d_f\), applies the nonlinear function \(\sigma\), and then projects back to \(d_\text{model}\).

The output of the transformer layer is therefore \[X_{\mathrm{out}} = X' + \mathrm{FFN} \left( X' \right). \label{eq:ffn95residual95update}\tag{16}\] The same structure is repeated across all \(L_T\) transformer layers. After the final layer, we obtain the final-layer residual stream \[H \equiv X^{(L_T)}_\text{out} \in \mathbb{R}^{L \times d_\text{model}}.\]

This is the representation analyzed by the SAE in this work, right before the unembedding layer. The activation vector at each position, \(\mathbf{h}_i \in \mathbb{R}^{d_\text{model}}\), is linearly mapped to logits over the spin dictionary: \[\mathbf{l}_i = \mathbf{h}_i W_U + \mathbf{b}_U, \qquad \mathbf{l_i} \in \mathbb{R}^2\] where \(\mathbf{l}_i\) is the logit on site \(i\), \(W_U \in \mathbb{R}^{d_\text{model}\times 2}\) and \(\mathbf{b}_U \in \mathbb{R}^2\) are unembedding weight and bias.

The conditional probability for the physical spin \(\sigma_i\) is obtained by applying a softmax to these logits \[p_\theta(\sigma_i=s\mid \sigma_{<i}) = \frac{ \exp(l_{i,s}) }{ \sum_{s'\in \{\pm 1\}}\exp(l_{i,s'}) }, \quad s\in \{-1, +1\}. \label{eq:softmax95conditional}\tag{17}\]

A second linear head maps the same residual stream to the conditional phase. The wavefunction amplitude is assembled from both the amplitude and phase.

Thus, the final residual-stream activation vectors are not themselves wavefunction amplitudes. They are first unembedded into logits, and the logits are then normalized by a softmax to produce the autoregressive conditional distributions. The SAE, therefore, does not directly analyze the probability distribution outputted by the transformer NQS, but an internal representation immediately upstream of the unembedding layer.

7 Hyperparameters of NQS and SAE↩︎

To ensure reproducibility, we list the hyperparameters used in the NQS and SAE in Table 1 and Table 2, respectively. The same choice of parameters is used across the TFIM ground state, TFIM real-time evolution, and Heisenberg AFM.

In addition to the \(L_1\) loss in the main text, one may also train a top-K SAE [73], in which sparsity is imposed directly by retaining only the \(K\) largest feature activations and setting the remaining coefficients to zero. We tested both loss functions for training the SAE and found no significant differences in the feature-observable correlations. The results in this work use the \(L_1\) penalized loss function.

8 Admissible Range of Steering Parameter \(\alpha\)↩︎

In Sec. 3, we demonstrate the ability to control observables by rescaling the top-correlated features by \(\alpha\): \(\mathbf{z}_{k^*} \rightarrow \alpha \mathbf{z}_{k^*}\). To isolate the effect of this intervention from any resampling artifact, we used the same set of spin configurations sampled from the original (unmodified) Born distribution \(|\psi|^2\). As a result, the modified distribution \(|\widetilde{\psi}|^2\) must remain sufficiently close to the original distribution, for the fixed sample set to provide reliable estimates.

Observables under the modified wavefunction \(\widetilde{\psi}(\alpha)\) are estimated by importance sampling. For each spin configuration, we define the weight \[w(\boldsymbol{\sigma}) = \frac{|\widetilde{\psi}(\boldsymbol{\sigma})|^2}{|\psi(\boldsymbol{\sigma})|^2}.\] The expectation value of an observable \(O\) under the modified wavefunction is estimated as: \[\langle O\rangle_{\widetilde{\psi}} \approx \frac{\sum_\sigma w(\sigma)O(\sigma)}{\sum_\sigma w(\sigma)}.\]

The quality of this estimate depends on whether the weights are well-behaved: if a small number of samples dominate the average, then the estimator has a high variance, and the estimated observables are no longer reliable. We monitor this effect using the effective sample size (ESS), \[\mathrm{ESS} = \frac{\bigl(\sum_{\boldsymbol{\sigma}} w(\boldsymbol{\sigma})\bigr)^2}{\sum_{\boldsymbol{\sigma}} w(\boldsymbol{\sigma})^2},\] expressed as a fraction of total sample size \(M\). Intuitively, the ESS estimates the effective number of useful samples from the original distribution that would contribute correctly to estimating expectation values under the modified distribution \(|\widetilde{\psi}|^2\). Therefore, when \(\mathrm{ESS}/M\) is close to 1, the modified distribution is close to the original one, and nearly all \(M\) samples contribute effectively, yielding reliable estimates. As \(\alpha\) moves away from 1, the weights become increasingly uneven, and the variance of the estimator increases, leading to incorrect expectation values.

Fig. 6 shows the ratio \(\mathrm{ESS}/M\) for both the paramagnetic case (\(h=1.5)\) and ferromagnetic case (\(h=0.5\)) for the range of \(\alpha\) considered in Fig. 3. For the values of \(\alpha\) considered, the ratio ESS\(/M\) remains above 0.90, which confirms that the importance-sampling estimates are reliable; therefore, the observed changes in observables are not artifacts of weight degeneracy.

9 Randomly Initialized NQS↩︎

In Fig. 2, we have demonstrated strong correlation between certain sparse features with physical observables such as magnetization. To verify that this correlation is not an artifact of the transformer architecture, the sampling procedure, or the SAE itself, we repeat the analysis using a randomly initialized NQS of the same architecture and hyperparameters as the trained model. The SAE is trained on the final-layer residual-stream activations of this random NQS. We draw the same number of configurations and repeat the same SAE analysis as in the main text.

Fig. 7 shows the scatter plot between the top-correlated feature (\(f79\)) and the magnetization \(M_z\), in the paramagnetic regime \(h=1.5\). The maximum correlation is only \(|r|=0.16\), which is in sharp contrast to the trained model, where the top feature tracks magnetization with \(|r|=0.99\) (Fig. 2 (a)). This confirms that the high feature-observable correlations observed in the trained models emerge during NQS learning. The NQS must learn an accurate representation of the quantum state for the SAE to recover physically meaningful features.

References↩︎