Physics-informed Transformers for Electronic Quantum States

TL;DR

This work tackles the basis-dependence and interpretability challenges of neural quantum states by introducing a physics-informed variational Monte Carlo approach that centers a reference state (RS) from Hartree-Fock or a strong-coupling limit and uses a Transformer to learn corrections to the RS. The method yields accurate ground-state energies while sampling a small subset of the Hilbert space and reveals an energetic and latent-ordering structure among basis states, with the RS weight α indicating proximity to the RS across phases. The Transformer's hidden representations encode excitations on the RS, providing physically meaningful interpretability and a hierarchical organization of corrections that enhances efficiency, especially away from criticality. Together, these results open a path toward more efficient, interpretable neural quantum-state representations and motivate extensions to other Hamiltonians and learning architectures.

Abstract

Neural-network-based variational quantum states in general, and more recently autoregressive models in particular, have proven to be powerful tools to describe complex many-body wave functions. However, their performance crucially depends on the computational basis chosen and they often lack physical interpretability. To mitigate these issues, we here propose a modified variational Monte-Carlo framework which leverages prior physical information to construct a computational second-quantized basis containing a reference state that serves as a rough approximation to the true ground state. In this basis, a Transformer is used to parametrize and autoregressively sample the corrections to the reference state, giving rise to a more interpretable and computationally efficient representation of the ground state. We demonstrate this approach using a non-sparse fermionic model featuring a metal-insulator transition and employing Hartree-Fock and a strong-coupling limit to define physics-informed bases. We also show that the Transformer's hidden representation captures the natural energetic order of the different basis states. This work paves the way for more efficient and interpretable neural quantum-state representations.

Figures (9)

• Figure 1: General methodology. (a) First, we choose a $\hat{H}_{0}$ approximating the target Hamiltonian $\hat{H}$, e.g., via a mean-field approximation or by taking the strong-coupling limit. We use the groundstate $\ket{\text{RS}}$ and excited states $\ket{\boldsymbol{s}}$ of $\hat{H}_{0}$ to define a physics-informed, interpretable basis for the transformer (b) in Eq. (\ref{['eq:tqsmod']}); as long as the dominant weight of the ground state of $\hat{H}$ is in the low-energy part of the spectrum $E_0(\boldsymbol{s})$ of $\hat{H}_{0}$, this further improves sampling efficiency and the expressivity of the ansatz. (c) We sample the states $\boldsymbol{s}$ using the batch-autoregressive sampler barrettAutoregressiveNeuralnetworkWavefunctions2022bmalyshev2023autoregressiveMalyshev2024Aug. It is controlled by the batch size $N_{s}$ and the number of partial unique strings $n_{U}$, and directly produces the relative frequencies $r(\boldsymbol{s})$ associated with each state in a tree structure format. Back to (b), the states $\boldsymbol{s}$ are then mapped to a high-dimensional representation of size $d_{ \text{emb}}$ and passed through $N_{\text{dec}}$ decoder-layers zhang2023tqs, containing $N_{h}$ attention heads, which produce correspondent representations $\boldsymbol{h}\left(\boldsymbol{s}\right)\in \mathbb{R}^{d_{\text{emb}}}$ in latent space. As discussed in the main text, the wavefunctions $\psi_{\boldsymbol{\theta}}(\boldsymbol{s})=\sqrt{q_{\boldsymbol{\theta}}(\boldsymbol{s})}e^{i\phi_{\boldsymbol{\theta}}(\boldsymbol{s})}$ can be directly obtained from these vectors. A new set of states $\mathcal{C}$ is then obtained, according to the updated $q_{\boldsymbol{\theta}}(\boldsymbol{s})$, and the process is repeated until the convergence of $\{\boldsymbol{\theta},\alpha\}$ according to Eq. (\ref{['eq:loss']}).
• Figure 2: Performance of HF-TQS for different system sizes and $t/U$. (a) Difference between the ground state energy per site and HF for $x$ as ED (solid lines) and TQS (markers) in $\delta E_{\text{HF}}=\left|E_{x}-E_{\text{HF}}\right|$ at various system sizes $N_e$. The inset shows the absolute value of the relative error $\delta E_{\text{ED}}=\left|E_{\text{TQS}}-E_{\text{ED}}\right|$. The corresponding converged $\alpha$ weights [according to Eq. (\ref{['eq:expenergy']})] are shown in (b). The gray regions indicate the vicinity of the metal-insulator transition. Convergence of the ground state energy per site (c) and of $\alpha$ (d) as a function of epochs for $t/U=0.16$ and $N_{e}=30$. $n_{U}^{f}$ is the total number of unique states kept by the Transformer from the value $n_{U}$ set initially in Fig. \ref{['fig:generalmethodology']}(c). Training time refers to one NVIDIA H100 GPU with networks parameters defined in Fig. \ref{['fig:generalmethodology']}b (see Appendix \ref{['sec:vmc']}). The total number of network parameters $\boldsymbol{\theta}$ used by the TQS is indicated by $\# \boldsymbol{\theta}$.
• Figure 3: Performance comparison of the TQS in different bases. (a) Difference between the ground state energy per site and HF for $x$ as ED (solid lines) and TQS (markers) in $\delta E_{\text{HF}}=\left|E_{x}-E_{\text{HF}}\right|$ for $N_{e}=10$. The inset shows the absolute value of the relative error $\delta E_{\text{ED}}=\left|E_{\text{TQS}}-E_{\text{ED}}\right|$. (b) Converged $\alpha$ for the TQS (markers) in (a) as a function of $t/U$, with dashed lines as a guide to the eye. (c) Histograms indicate the total relative frequencies, according to Eq. (\ref{['eq:totalweight']}), for the excitation classes $\mathcal{E}(\boldsymbol{s})$ from Eq. (\ref{['eq:labeldof']}). From top to bottom, these refer to the insulating, critical and metallic regions indicated by the respective gray boxes in (a). (d) HF-TQS results (markers) for the momentum-resolved fermionic bilinears $\mathcal{N}^j_{\boldsymbol{k}}$, defined in Eq. (\ref{['eq:observablen']}), in comparison to those obtained from HF (dashed lines) for $N_{e}=12$.
• Figure 4: Visualization of the Transformer's latent space. Results are shown for $N_{e}=10$ electrons at different values of $t/U$ for the band, chiral and HF bases. Each point represents a basis state $\boldsymbol{s}$, which is colored according to the class label $\mathcal{E}(\boldsymbol{s})$ [cf. Eq. (\ref{['eq:labeldof']})], and has been obtained by projecting the respective latent space features $\boldsymbol{H}(\boldsymbol{s})$ onto two dimensions using PCA. All simulations use embedding dimension $d_{\text{emb}}=300$ with single attention head and decoder layer $N_{\text{h}}=N_{\text{dec}}=1$.
• Figure 5: Observables for the fermionic model from the perspective of Hartree-Fock for $N_{e}=100$ electrons. Band structure for the fermionic model at $t/U=0$ (a) and $t/U=10$ (b) obtained according to the expression \ref{['eq:hfunit']}. (c) Order parameter $\xi$ as a function of $t/U$. (d) Momentum-resolved fermionic bilinears (according to Eq. (\ref{['eq:observablen']})) as a function of $\boldsymbol{k}$ for different values of $t/U$.