Accurate Self-Attention Wavefunctions at Large Scale

Technical Summary: Accurate Self-Attention Wavefunctions at Large Scale

Problem Statement
Quantum many-body simulations face a fundamental bottleneck: the exponential growth of the Hilbert space limits the scalability of classical numerical methods. While traditional Variational Monte Carlo (VMC) and Diffusion Monte Carlo (DMC) techniques have been successful, their accuracy is often constrained by the quality of the trial wavefunctions, which typically rely on fixed analytic forms (e.g., Slater determinants augmented by Jastrow and backflow correlations). Machine learning-based Neural Quantum States (NQS) offer a route to more expressive wavefunctions by learning correlations directly from energy minimization. However, self-attention architectures, which provide high expressivity by modeling non-local, particle-dependent correlations without hand-crafted orbital ansätze, suffer from high computational complexity. The central question addressed is whether self-attention wavefunctions can retain their accuracy and superior expressivity at large system sizes ($N_e \approx 100$), or if the computational cost renders them impractical for the thermodynamic limit.

Methodology
The authors apply a self-attention neural network ansatz to the two-dimensional homogeneous electron gas (HEG), a benchmark system for strongly correlated physics.

• Architecture: The wavefunction consists of a stack of $L$ self-attention and perceptron layers. Unlike message-passing networks that apply backflow transformations to fixed plane-wave orbitals, this architecture generates generalized single-particle orbitals $\phi^{(k)}_i(r_j, \{r_{\neq j}\})$ directly from the many-body electron configuration. These orbitals are combined into $N_{det}$ Slater determinants to form the antisymmetric fermionic wavefunction.
• Scalability Mechanism: To address the $O(N_e^2)$ scaling of the all-to-all attention mechanism, the authors introduce a "spatially weighted particle attention." An exponential damping factor $e^{-r_{ij}/\lambda}$ is applied to the attention scores, where $r_{ij}$ is the interparticle distance and $\lambda$ is a trainable length scale. This biases the network to focus on local correlations while allowing long-range information to propagate through the depth of the network layers.
• Simulation Setup: The system is solved using Variational Monte Carlo (VMC) with periodic boundary conditions on a triangular lattice. The authors simulate spin-polarized electrons with $N_e = 91$ and $N_e = 169$ particles at interaction strengths $r_s = 30$ and $r_s = 40$, regimes near the liquid-crystal transition where correlation effects are strongest.
• Benchmarking: Results are compared against state-of-the-art fixed-node Diffusion Monte Carlo (DMC) using optimized Slater-Jastrow-backflow (SJB) wavefunctions. The DMC results serve as the reference for the nodal structure, while the neural network (NN-VMC) is trained to minimize the variational energy.

Key Contributions and Results

1. Superior Variational Energies: The self-attention wavefunctions achieve variational energies systematically lower than state-of-the-art fixed-node DMC results across all tested system sizes and interaction strengths. For $N_e = 169$ at $r_s = 30$, the NN-VMC energy is lower than the SJB-DMC benchmark, demonstrating that the neural ansatz captures correlation energy more effectively than traditional methods. The improvement over traditional SJB-VMC is quantified as approximately 1.4% in correlation energy.
2. Scalability and Convergence: Crucially, the accuracy of the self-attention wavefunction does not deteriorate as the system size increases from $N_e = 91$ to $N_e = 169$. This indicates that the method successfully scales to sizes approaching the thermodynamic limit without loss of fidelity.
3. Access to Collective Modes: Because the neural wavefunction provides direct access to the ground state, the authors compute observables beyond energy, including the momentum distribution $n(k)$ and the static structure factor $S(q)$.
   • Phase Characterization: The method correctly distinguishes between the metallic liquid phase ($r_s=30$, showing a sharp Fermi surface discontinuity) and the crystalline Wigner crystal phase ($r_s=40$, showing Bragg peaks in $S(q)$ and triangular density modulation).
   • Dispersion Relations: Using the Feynman relation $\epsilon(q) \sim q^2/S(q)$, the authors extract the collective mode dispersion. They recover the full excitation branch of the liquid phase, including the expected plasmon dispersion ($\epsilon \propto \sqrt{q}$) at small momenta and a roton-like minimum near $q \approx 2k_F$.
4. Thermodynamic Limit Agreement: The extracted dispersions for $N_e = 91$ and $N_e = 169$ are in near-perfect agreement, suggesting that the self-attention wavefunction has already converged to the thermodynamic limit behavior for these system sizes.

Significance
The paper claims that self-attention architectures can simultaneously meet the dual requirements of high numerical accuracy and scalability for large system sizes in continuum fermionic systems. By successfully solving the homogeneous electron gas—a cornerstone of many-body physics—at $N_e = 169$ with energies surpassing fixed-node DMC, the work demonstrates that neural quantum states are not merely theoretical curiosities but practical tools for quantitative physics. The ability to recover accurate collective excitations and phase transitions without hand-crafted ansätze suggests that this approach can be extended to explore new physics in more complex electronic systems, such as semiconductors, and potentially aid in materials discovery and quantum technology applications.