Transformer Wave Functions in Quantum Systems

• Transformer wave functions are neural-network quantum states that parameterize many-body wave functions using Transformer architectures to capture both local and global correlations.
• They employ varied attention mechanisms—including T5, decoupled, and factored forms—to efficiently encode spin configurations while integrating lattice symmetry and translation invariance.
• Hybrid architectures combining convolutional and genuine self-attention, alongside variational optimization via stochastic reconfiguration, enable precise ground-state approximations and dynamic simulations.

Searching arXiv for recent and foundational papers on Transformer wave functions to ground the article in published work. Transformer wave functions are neural-network quantum states in which the many-body wave-function amplitudes are parameterized by a Transformer or Vision Transformer (ViT)-style architecture. In the standard formulation, one writes

$$|\Psi_\theta\rangle=\sum_{\sigma}\Psi_\theta(\sigma)\,|\sigma\rangle,$$

with $\Psi_\theta(\sigma)$ a complex-valued function of a basis configuration $\sigma$, and the Transformer is used to represent either $\Psi_\theta(\sigma)$ directly or $\log \Psi_\theta(\sigma)$ (Viteritti et al., 2022, Rende et al., 2024). Within the broader neural-network quantum state framework, Transformer wave functions were introduced as variational ansätze for frustrated quantum many-body systems, especially lattice spin models, because self-attention provides a direct mechanism for mixing local and global information and for encoding long-range correlations (Viteritti et al., 2022). Subsequent work has refined this picture in two directions: one line showed that position-only attention can match standard dot-product attention for benchmark ground-state problems (Rende et al., 2024), while another argued that earlier factored-attention constructions are mathematically equivalent to convolutional layers and proposed hybrid convolutional-transformer architectures that restore genuine content-dependent multi-head self-attention while preserving lattice symmetries (Chen et al., 13 Mar 2025).

1. Conceptual definition and scope

A Transformer wave function belongs to the class of neural-network quantum states (NQS), where a variational ansatz is defined by a neural network over basis configurations. For spin-$\tfrac12$ systems, the computational basis is typically written as

$$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$

or, in one dimension,

$$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$

and the variational state is expanded as a sum over configurations with complex amplitudes supplied by the network (Rende et al., 2024, Viteritti et al., 2022).

The defining architectural move is to treat a spin configuration as a sequence of tokens or as an “image” partitioned into patches. In the ViT formulation, the lattice is split into non-overlapping patches, each patch is flattened and embedded into a vector, and the resulting token sequence is processed by Transformer blocks. The network finally returns a complex scalar representing $\log \Psi_\theta(\sigma)$, or an equivalent complex amplitude after a final complex activation (Rende et al., 2024, Chen et al., 13 Mar 2025).

The physical target has primarily been the approximation of ground states of interacting Hamiltonians and, more recently, non-equilibrium dynamics. The benchmark models emphasized in the literature include the one-dimensional $J_1$-$\Psi_\theta(\sigma)$0 Heisenberg chain, the two-dimensional $\Psi_\theta(\sigma)$1-$\Psi_\theta(\sigma)$2 Heisenberg antiferromagnet on the square lattice, the Shastry–Sutherland model in an exactly solvable regime, and the two-dimensional transverse-field Ising model for dynamical simulations (Viteritti et al., 2022, Rende et al., 2024, Chen et al., 13 Mar 2025).

This suggests that “Transformer wave function” is best understood not as a single architecture but as a family of variational ansätze sharing a Transformer backbone while differing in their treatment of tokenization, attention, symmetry, and complex output parameterization.

2. Architectural realizations

The first explicit ViT-based construction for frustrated spin systems adapts the Vision Transformer to a complex-valued variational ansatz (Viteritti et al., 2022). In the one-dimensional $\Psi_\theta(\sigma)$3-$\Psi_\theta(\sigma)$4 chain, a spin string is split into non-overlapping patches of size $\Psi_\theta(\sigma)$5, typically $\Psi_\theta(\sigma)$6, producing a sequence of patch vectors $\Psi_\theta(\sigma)$7. Each Transformer block applies multi-head attention, a linear projection, and an elementwise $\Psi_\theta(\sigma)$8 nonlinearity, and the final outputs are summed to produce $\Psi_\theta(\sigma)$9 (Viteritti et al., 2022). All parameters are complex, which allows the ansatz to encode both amplitude and phase without separating them into distinct subnetworks.

A later two-dimensional ViT formulation for the $\sigma$0 $\sigma$1-$\sigma$2 model uses a similar pipeline but makes the encoder more explicitly modern-ViT-like (Rende et al., 2024). The lattice is partitioned into $\sigma$3 patches, typically $\sigma$4, each patch is linearly embedded into a $\sigma$5-dimensional real vector, and the encoder consists of $\sigma$6 identical blocks with multi-head self-attention, a two-layer MLP with GeLU activation, Pre-LayerNorm, and residual connections. The outputs $\sigma$7 are summed into a global representation

$\sigma$8

after which a shallow complex network returns

$\sigma$9

with complex parameters $\Psi_\theta(\sigma)$0 (Rende et al., 2024).

The convolutional transformer wave function (CTWF) introduces a hybrid design intended to combine local inductive bias with genuine content-dependent attention (Chen et al., 13 Mar 2025). Each main block contains a convolutional unit, a multi-head self-attention layer, and an inverted residual feed-forward network (IRFFN). The MHSA uses translation-symmetric relative positional encoding, while the convolutional units maintain periodic-lattice inductive bias. Before each sub-block, the architecture applies a normalization step by dividing inputs by their expected initial standard deviation rather than using LayerNorm, a choice reported to work better in that setting (Chen et al., 13 Mar 2025).

Across these formulations, a common architectural principle is the alternation of local processing and global mixing. In the 2022 ViT wave function, the local element is the patch embedding and the global element is the all-to-all positional attention (Viteritti et al., 2022). In the 2024 study, the same ViT scaffold is used to isolate the contribution of different attention mechanisms (Rende et al., 2024). In the 2025 CTWF, convolutional blocks explicitly handle local structure while MHSA is reserved for long-range content-dependent interactions (Chen et al., 13 Mar 2025).

3. Attention mechanisms and their interpretations

Attention is the central technical point on which the literature diverges most strongly. Standard self-attention forms queries, keys, and values from token embeddings and computes normalized overlaps to produce content-dependent mixing (Rende et al., 2024). In the two-dimensional ViT case study, three relative-position attention schemes were compared systematically:

Mechanism	Weight definition	Characterization
T5 Attention	dot-product softmax with relative position bias $\Psi_\theta(\sigma)$1	input-dependent, normalized
Decoupled Attention	softmax dot-product term plus additive $\Psi_\theta(\sigma)$2	input-dependent + additive positional term
Factored Attention	$\Psi_\theta(\sigma)$3	position-only, input-independent

In this setting, T5 attention is a standard dot-product attention with relative position bias, Decoupled attention separates the input-dependent softmax term from the positional term, and Factored attention removes queries and keys entirely so that attention depends only on relative lattice position (Rende et al., 2024).

The 2024 analysis assigns both physical and computational meaning to Factored attention. Physically, it encodes the strength with which patch $\Psi_\theta(\sigma)$4 contributes to the representation of patch $\Psi_\theta(\sigma)$5 as a fixed function of relative lattice position; all configuration dependence then flows through the value projections and subsequent nonlinear layers. Computationally, Factored attention removes the need to compute $\Psi_\theta(\sigma)$6, $\Psi_\theta(\sigma)$7, and their pairwise dot products; the per-attention-weight cost becomes $\Psi_\theta(\sigma)$8, whereas T5 and Decoupled attention require $\Psi_\theta(\sigma)$9 for projections plus $\log \Psi_\theta(\sigma)$0 for the dot-product matrix per head and layer (Rende et al., 2024).

The 2022 ViT wave function had already implemented a closely related position-only scheme in one dimension. There, the attention vectors are written as

$\log \Psi_\theta(\sigma)$1

with learned positional weights $\log \Psi_\theta(\sigma)$2 depending on positions $\log \Psi_\theta(\sigma)$3 and head index $\log \Psi_\theta(\sigma)$4, but not on the spin content of the patches (Viteritti et al., 2022). The motivation given is that which patches interact strongly depends mainly on their relative distance, not on their specific spin content, especially for distant patches (Viteritti et al., 2022).

The CTWF work reinterprets this earlier factored-attention line more critically. It shows that factored attention used in previous ViT-like NQSs for lattices is mathematically equivalent to a pair of convolutional layers rather than genuine attention (Chen et al., 13 Mar 2025). In the CTWF, one-head attention is defined as

$\log \Psi_\theta(\sigma)$5

where $\log \Psi_\theta(\sigma)$6 is a trainable translation-invariant relative positional encoding on the periodic lattice (Chen et al., 13 Mar 2025). The explicit addition of $\log \Psi_\theta(\sigma)$7 restores content dependence, which the paper identifies as the genuinely transformer-like component.

A plausible implication is that the literature distinguishes two different meanings of “Transformer” in this context: one tied to positional all-to-all mixing that may collapse to convolutional structure, and another tied to true content-dependent multi-head self-attention.

4. Variational optimization and symmetry structure

Transformer wave functions are trained variationally by minimizing the expectation value of the Hamiltonian. In the two-dimensional ViT formulation, the cost function is

$\log \Psi_\theta(\sigma)$8

and the corresponding local energy is

$\log \Psi_\theta(\sigma)$9

Expectation values are estimated by Monte Carlo sampling from $\tfrac12$0 (Rende et al., 2024).

Optimization is typically carried out with stochastic reconfiguration (SR), equivalent to natural gradient descent in parameter space (Rende et al., 2024). Defining logarithmic derivatives

$\tfrac12$1

the energy gradient is expressed as a covariance and the parameter update is

$\tfrac12$2

where the quantum geometric tensor or Fubini–Study metric is

$\tfrac12$3

Both $\tfrac12$4 and $\tfrac12$5 are estimated with MCMC, and log-derivatives are computed by automatic differentiation in JAX (Rende et al., 2024).

The 2022 one-dimensional ViT wave function uses the same SR principle, with a regularized covariance matrix and Monte Carlo sampling restricted to the $\tfrac12$6 sector using local spin-exchange updates (Viteritti et al., 2022). Typical sample sizes per optimization step are $\tfrac12$7, and the matrix inversion is stabilized by adding a small diagonal regularization $\tfrac12$8 (Viteritti et al., 2022).

The CTWF retains SR for ground-state optimization and extends the same variational geometry to real-time dynamics through the time-dependent variational principle (TDVP) (Chen et al., 13 Mar 2025). For the $\tfrac12$9 benchmark, networks are trained with SR using $$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$0 Monte Carlo samples for $$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$1 steps; for the larger $$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$2 problem, the work uses MinSR, described there as a large-scale stabilized SR variant (Chen et al., 13 Mar 2025). For non-equilibrium dynamics in the transverse-field Ising model, the variational parameters are evolved with TDVP using a second-order Heun method with time step $$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$3 (Chen et al., 13 Mar 2025).

Symmetry handling is a major architectural constraint. The 2022 and 2024 ViT constructions emphasize translation-friendly relative positional structure (Viteritti et al., 2022, Rende et al., 2024). The CTWF makes this explicit by using periodic-lattice relative positional encoding $$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$4, together with convolutional layers with periodic boundaries, so that translation symmetry is built into the architecture (Chen et al., 13 Mar 2025). It then adds symmetry projections for spatial reflections, rotations, and spin-flip, amounting to 16 symmetry group elements in total (Chen et al., 13 Mar 2025).

5. Benchmark systems and empirical performance

The first systematic benchmark for a ViT wave function was the one-dimensional antiferromagnetic $$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$5 $$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$6-$$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$7 Heisenberg chain with periodic boundary conditions,

$$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$8

which contains a gapless phase at small $$|\sigma\rangle = |\sigma_1^z,\dots,\sigma_N^z\rangle,\qquad \sigma_i^z=\pm 1,$$9, a gapped dimerized phase at large $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$0, a critical point at $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$1, and incommensurate spin correlations for $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$2 (Viteritti et al., 2022). On $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$3, benchmark DMRG energies reported for periodic boundary conditions are $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$4 at $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$5, $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$6 at $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$7, and $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$8 at $$\sigma=(\sigma_1,\dots,\sigma_L),\qquad \sigma_R=2S_R^z=\pm 1,$$9 (Viteritti et al., 2022). With a single attention layer and fewer than $\log \Psi_\theta(\sigma)$0 parameters, the ViT wave function achieves $\log \Psi_\theta(\sigma)$1 even at $\log \Psi_\theta(\sigma)$2, and $\log \Psi_\theta(\sigma)$3 at $\log \Psi_\theta(\sigma)$4 (Viteritti et al., 2022). It also reproduces spin–spin correlations, the structure factor $\log \Psi_\theta(\sigma)$5, and the saturation of connected dimer–dimer correlations in the dimerized regime (Viteritti et al., 2022).

The 2024 case study moves to the two-dimensional spin-$\log \Psi_\theta(\sigma)$6 $\log \Psi_\theta(\sigma)$7-$\log \Psi_\theta(\sigma)$8 Heisenberg antiferromagnet on a $\log \Psi_\theta(\sigma)$9 square lattice with periodic boundary conditions, in both the unfrustrated regime $J_1$0 and the frustrated regime $J_1$1, where exact-diagonalization benchmarks are available (Rende et al., 2024). For fixed architecture size, all three attention mechanisms—T5, Decoupled, and Factored—reach essentially the same energy in both regimes, and the convergence dynamics are nearly identical (Rende et al., 2024). A representative comparison at $J_1$2 with a four-layer ViT gives

Attention	Energy $J_1$3	Parameters
T5	$J_1$4	184260
Decoupled	$J_1$5	184260
Factored	$J_1$6	154980

with exact ground-state energy $J_1$7 (Rende et al., 2024). The same comparison reports 10 h on an A100 GPU for T5 and Decoupled versus 6 h for Factored attention, corresponding to about 40% speed-up without degrading energy accuracy (Rende et al., 2024).

The CTWF study uses the maximally frustrated point $J_1$8 of the square-lattice Heisenberg model as a demanding benchmark in both $J_1$9 and $\Psi_\theta(\sigma)$00 geometries (Chen et al., 13 Mar 2025). On $\Psi_\theta(\sigma)$01, a complex translationally symmetric RBM gives $\Psi_\theta(\sigma)$02, a CNN with GELU gives $\Psi_\theta(\sigma)$03, $\Psi_\theta(\sigma)$04, $\Psi_\theta(\sigma)$05, and a Transformer with factored attention gives $\Psi_\theta(\sigma)$06 (Chen et al., 13 Mar 2025). The best CTWF design—linear $\Psi_\theta(\sigma)$07, convolutional IRFFN, and relative positional encoding—reaches $\Psi_\theta(\sigma)$08, $\Psi_\theta(\sigma)$09, and $\Psi_\theta(\sigma)$10 (Chen et al., 13 Mar 2025). On $\Psi_\theta(\sigma)$11, CTWF and CNN (GELU) both improve significantly over factored attention in relative energy error and variance, and the paper states that CTWF and CNN (GELU) achieve energies extremely close to the best known variational estimates obtained via zero-variance extrapolation (Chen et al., 13 Mar 2025).

The same CTWF architecture is also tested dynamically on the two-dimensional transverse-field Ising model at the critical point $\Psi_\theta(\sigma)$12, starting from the fully polarized paramagnetic state $\Psi_\theta(\sigma)$13 (Chen et al., 13 Mar 2025). On $\Psi_\theta(\sigma)$14, CTWF and CNN (GELU) track the exact evolution of $\Psi_\theta(\sigma)$15 very closely and remain stable and accurate for significantly longer times than the earlier CNN-based TDVP baseline; on $\Psi_\theta(\sigma)$16, the two newer architectures remain mutually consistent beyond the time where the older baseline has already drifted (Chen et al., 13 Mar 2025).

6. Analytical interpretation, controversies, and broader significance

A central analytical result of the 2024 study is that standard input-dependent attention becomes effectively input-independent during optimization for the tested ground-state problems (Rende et al., 2024). In a one-layer ViT with embedding dimension $\Psi_\theta(\sigma)$17, $\Psi_\theta(\sigma)$18 heads, and the $\Psi_\theta(\sigma)$19 frustrated $\Psi_\theta(\sigma)$20-$\Psi_\theta(\sigma)$21 model at $\Psi_\theta(\sigma)$22, T5 attention maps differ across inputs at initialization but are essentially identical across different input configurations after optimization (Rende et al., 2024). For Decoupled attention, the input-dependent component averaged over 6000 samples is close to an identity matrix plus a constant offset, $\Psi_\theta(\sigma)$23, while the meaningful off-diagonal structure resides entirely in the positional term $\Psi_\theta(\sigma)$24 (Rende et al., 2024). The paper therefore concludes that both T5 and Decoupled attention spontaneously reduce to the position-only Factored form under VMC optimization (Rende et al., 2024).

That numerical observation is supported there by two analytical arguments. First, a single-layer ViT with Factored attention and two heads can exactly represent a Shastry–Sutherland ground state that is a product of singlets on a pattern of dimers (Rende et al., 2024). In that exact construction, one head uses $\Psi_\theta(\sigma)$25, the other uses $\Psi_\theta(\sigma)$26, and the attention maps encoding which patches are entangled are completely positional; spin dependence enters only through values and nonlinearities (Rende et al., 2024). Second, invoking the cluster property

$\Psi_\theta(\sigma)$27

for local operators, the paper argues that in large systems attention between distant patches should reduce to a function only of distance rather than of the local token values, implying that $\Psi_\theta(\sigma)$28 and $\Psi_\theta(\sigma)$29 should become ineffective in the large-distance limit (Rende et al., 2024).

The CTWF work presents a different conclusion, but on different architectural grounds (Chen et al., 13 Mar 2025). It argues that factored attention in earlier lattice ViTs degenerates to convolutional structure, and that this means such constructions have not used the full power of transformers. The proposed remedy is not to abandon attention, but to combine translation-symmetric relative positional encoding with genuine content-dependent MHSA and convolutional units. On its benchmarks, this design exhibits superior performance in ground-state search and non-equilibrium dynamics compared to previous reported transformer results (Chen et al., 13 Mar 2025).

These two lines are not identical in scope. One studies whether queries and keys matter inside a specific ViT wave-function family for local Hamiltonian ground states on $\Psi_\theta(\sigma)$30 lattices (Rende et al., 2024). The other studies whether a translation-symmetric hybrid architecture with genuine MHSA can outperform earlier transformer NQS designs on frustrated two-dimensional benchmarks and TFIM dynamics (Chen et al., 13 Mar 2025). This suggests that the relevance of content-dependent attention is contingent on the ansatz class, the symmetry constraints, and whether the task is ground-state approximation or real-time evolution.

Within the broader NQS landscape, Transformer wave functions sit alongside RBMs, CNNs, RNNs, autoregressive models, and tensor-network-inspired constructions (Viteritti et al., 2022, Chen et al., 13 Mar 2025). Prior works by Luo et al., Pfau et al., von Glehn et al., and Viteritti et al. are noted in the literature as examples where transformer-based NQS have reached state-of-the-art energies in one-dimensional and two-dimensional spin systems, bosonic systems, and ab initio electrons (Rende et al., 2024). The main methodological lesson emerging from the present literature is narrower: locality, translation symmetry, and relative positional structure are indispensable inductive biases; whether full query-key semantics is necessary depends on how much of the physical correlation structure is already fixed by geometry and on whether one seeks efficient ground-state representation or broader expressive capacity (Rende et al., 2024, Chen et al., 13 Mar 2025).