---
title: Transformer Wave Functions in Quantum Systems
url: https://www.emergentmind.com/topics/transformer-wave-functions
type: topic
---

# Transformer Wave Functions in Quantum Systems

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)\) [2211.05504][2405.18874]. 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 [2211.05504]. 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 [2405.18874], 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 [2503.10462].

## 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 [2405.18874][2211.05504].

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 [2405.18874][2503.10462].

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\)-\(J_2\) Heisenberg chain, the two-dimensional \(J_1\)-\(J_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 [2211.05504][2405.18874][2503.10462].

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 [2211.05504]. In the one-dimensional \(J_1\)-\(J_2\) chain, a spin string is split into non-overlapping patches of size \(b\), typically \(b=4\), producing a sequence of patch vectors \((\boldsymbol{x}_1,\dots,\boldsymbol{x}_N)\). Each Transformer block applies multi-head attention, a linear projection, and an elementwise \(\log[\cosh(\cdot)]\) nonlinearity, and the final outputs are summed to produce \(\log \Psi_{\rm ViT}(\sigma)\) [2211.05504]. 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 \(6\times 6\) \(J_1\)-\(J_2\) model uses a similar pipeline but makes the encoder more explicitly modern-ViT-like [2405.18874]. The lattice is partitioned into \(b\times b\) patches, typically \(b=2\), each patch is linearly embedded into a \(d\)-dimensional real vector, and the encoder consists of \(n_l\) identical blocks with multi-head self-attention, a two-layer MLP with GeLU activation, Pre-LayerNorm, and residual connections. The outputs \((\mathbf y_1,\dots,\mathbf y_n)\) are summed into a global representation
\[
\mathbf z=\sum_{i=1}^n \mathbf y_i,
\]
after which a shallow complex network returns
\[
\log\Psi_\theta(\sigma)=\sum_{\alpha=1}^d g\big(b_\alpha+\mathbf w_\alpha\cdot \mathbf z\big),
\qquad g(x)=\log\cosh(x),
\]
with complex parameters \(b_\alpha,\mathbf w_\alpha\) [2405.18874].

The convolutional transformer wave function (CTWF) introduces a hybrid design intended to combine local inductive bias with genuine content-dependent attention [2503.10462]. 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 [2503.10462].

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 [2211.05504]. In the 2024 study, the same ViT scaffold is used to isolate the contribution of different attention mechanisms [2405.18874]. In the 2025 CTWF, convolutional blocks explicitly handle local structure while MHSA is reserved for long-range content-dependent interactions [2503.10462].

## 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 [2405.18874]. 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 \(p_{i-j}\) | input-dependent, normalized |
| Decoupled Attention | softmax dot-product term plus additive \(p_{i-j}\) | input-dependent + additive positional term |
| Factored Attention | \(\alpha_{ij}^{\mathrm F}=p_{i-j}\) | 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 [2405.18874].

The 2024 analysis assigns both physical and computational meaning to Factored attention. Physically, it encodes the strength with which patch \(j\) contributes to the representation of patch \(i\) 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 \(Q\mathbf x_i\), \(K\mathbf x_j\), and their pairwise dot products; the per-attention-weight cost becomes \(O(1)\), whereas T5 and Decoupled attention require \(O(n d^2)\) for projections plus \(O(n^2 d)\) for the dot-product matrix per head and layer [2405.18874].

The 2022 ViT wave function had already implemented a closely related position-only scheme in one dimension. There, the attention vectors are written as
\[
\boldsymbol{A}^{\mu}_i=\sum_{j=1}^N \alpha^\mu_{ij}\,V^\mu \boldsymbol{x}_j,
\]
with learned positional weights \(\alpha^\mu_{ij}\) depending on positions \(i,j\) and head index \(\mu\), but not on the spin content of the patches [2211.05504]. 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 [2211.05504].

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 [2503.10462]. In the CTWF, one-head attention is defined as
\[
A=\mathrm{softmax}\left(\frac{QK^T+P}{\sqrt d}\right)V,
\]
where \(P_{ij}=p_{x_i-x_j,\;y_i-y_j}\) is a trainable translation-invariant relative positional encoding on the periodic lattice [2503.10462]. The explicit addition of \(QK^T\) 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
\[
E[\Psi_\theta]=E_\theta=
\frac{\langle\Psi_\theta|\hat H|\Psi_\theta\rangle}
{\langle\Psi_\theta|\Psi_\theta\rangle}\ge E_0,
\]
and the corresponding local energy is
\[
E_{\rm loc}(\sigma)=
\frac{\langle \sigma|\hat H|\Psi_\theta\rangle}{\Psi_\theta(\sigma)}
=
\sum_{\sigma'} H_{\sigma,\sigma'}\,
\frac{\Psi_\theta(\sigma')}{\Psi_\theta(\sigma)}.
\]
Expectation values are estimated by Monte Carlo sampling from \(|\Psi_\theta(\sigma)|^2\) [2405.18874].

Optimization is typically carried out with stochastic reconfiguration (SR), equivalent to natural gradient descent in parameter space [2405.18874]. Defining logarithmic derivatives
\[
O_\alpha(\sigma)=\frac{\partial}{\partial\theta_\alpha}\log\Psi_\theta(\sigma),
\]
the energy gradient is expressed as a covariance and the parameter update is
\[
\delta\theta_\alpha=\tau\sum_\beta (S^{-1})_{\alpha\beta}F_\beta,
\]
where the quantum geometric tensor or Fubini–Study metric is
\[
S_{\alpha\beta}=
\Re\left[
\left\langle
(\hat O_\alpha-\langle\hat O_\alpha\rangle)^\dagger
(\hat O_\beta-\langle\hat O_\beta\rangle)
\right\rangle
\right].
\]
Both \(F\) and \(S\) are estimated with MCMC, and log-derivatives are computed by automatic differentiation in JAX [2405.18874].

The 2022 one-dimensional ViT wave function uses the same SR principle, with a regularized covariance matrix and Monte Carlo sampling restricted to the \(S^z=0\) sector using local spin-exchange updates [2211.05504]. Typical sample sizes per optimization step are \(O(10^3\!-\!10^4)\), and the matrix inversion is stabilized by adding a small diagonal regularization \(\varepsilon\sim 10^{-3}\!-\!10^{-4}\) [2211.05504].

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) [2503.10462]. For the \(6\times 6\) benchmark, networks are trained with SR using \(10^4\) Monte Carlo samples for \(10^4\) steps; for the larger \(10\times 10\) problem, the work uses MinSR, described there as a large-scale stabilized SR variant [2503.10462]. 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 \(J\tau=10^{-3}\) [2503.10462].

Symmetry handling is a major architectural constraint. The 2022 and 2024 ViT constructions emphasize translation-friendly relative positional structure [2211.05504][2405.18874]. The CTWF makes this explicit by using periodic-lattice relative positional encoding \(P_{ij}=p_{x_i-x_j,\;y_i-y_j}\), together with convolutional layers with periodic boundaries, so that translation symmetry is built into the architecture [2503.10462]. It then adds symmetry projections for spatial reflections, rotations, and spin-flip, amounting to 16 symmetry group elements in total [2503.10462].

## 5. Benchmark systems and empirical performance

The first systematic benchmark for a ViT wave function was the one-dimensional antiferromagnetic \(S=\tfrac12\) \(J_1\)-\(J_2\) Heisenberg chain with periodic boundary conditions,
\[
\hat H=
J_1\sum_R \hat{\boldsymbol S}_R\cdot \hat{\boldsymbol S}_{R+1}
+
J_2\sum_R \hat{\boldsymbol S}_R\cdot \hat{\boldsymbol S}_{R+2},
\]
which contains a gapless phase at small \(J_2/J_1\), a gapped dimerized phase at large \(J_2/J_1\), a critical point at \((J_2/J_1)_c=0.241167\pm 0.000005\), and incommensurate spin correlations for \(J_2/J_1>0.5\) [2211.05504]. On \(L=100\), benchmark DMRG energies reported for periodic boundary conditions are \(E/J_1=-0.4432295\) at \(J_2/J_1=0\), \(E/J_1=-0.3803882\) at \(J_2/J_1=0.4\), and \(E/J_1=-0.4216664\) at \(J_2/J_1=0.8\) [2211.05504]. With a single attention layer and fewer than \(10^3\) parameters, the ViT wave function achieves \(\Delta\varepsilon\lesssim 0.1\%\) even at \(J_2/J_1=0.8\), and \(\Delta\varepsilon\approx 0.01\%\) at \(J_2/J_1=0.4\) [2211.05504]. It also reproduces spin–spin correlations, the structure factor \(S(k)\), and the saturation of connected dimer–dimer correlations in the dimerized regime [2211.05504].

The 2024 case study moves to the two-dimensional spin-\(\tfrac12\) \(J_1\)-\(J_2\) Heisenberg antiferromagnet on a \(6\times 6\) square lattice with periodic boundary conditions, in both the unfrustrated regime \(J_2/J_1=0\) and the frustrated regime \(J_2/J_1=0.5\), where exact-diagonalization benchmarks are available [2405.18874]. 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 [2405.18874]. A representative comparison at \(J_2/J_1=0.5\) with a four-layer ViT gives

| Attention | Energy \(E/J_1\) | Parameters |
|---|---:|---:|
| T5 | \(-0.50318(1)\) | 184260 |
| Decoupled | \(-0.50324(1)\) | 184260 |
| Factored | \(-0.50321(1)\) | 154980 |

with exact ground-state energy \(E_0/J_1=-0.503810\) [2405.18874]. 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 [2405.18874].

The CTWF study uses the maximally frustrated point \(J_2/J_1=0.5\) of the square-lattice Heisenberg model as a demanding benchmark in both \(6\times 6\) and \(10\times 10\) geometries [2503.10462]. On \(6\times 6\), a complex translationally symmetric RBM gives \(\epsilon_{\mathrm{rel}}\approx 1.5\times 10^{-2}\), a CNN with GELU gives \(\epsilon_{\mathrm{rel}}\approx 2.4\times 10^{-3}\), \(\sigma^2/N\approx 4.7\times 10^{-3}\), \(I\approx 2.3\times 10^{-2}\), and a Transformer with factored attention gives \(\epsilon_{\mathrm{rel}}\approx 4.0\times 10^{-3}\) [2503.10462]. The best CTWF design—linear \(Q/K/V\), convolutional IRFFN, and relative positional encoding—reaches \(\epsilon_{\mathrm{rel}}\approx 2.3\times 10^{-3}\), \(\sigma^2/N\approx 5.3\times 10^{-3}\), and \(I\approx 1.4\times 10^{-2}\) [2503.10462]. On \(10\times 10\), 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 [2503.10462].

The same CTWF architecture is also tested dynamically on the two-dimensional transverse-field Ising model at the critical point \(h/J=3.04438(2)\), starting from the fully polarized paramagnetic state \(|\psi_0\rangle=|\rightarrow\rangle^{\otimes N}\) [2503.10462]. On \(6\times 6\), CTWF and CNN (GELU) track the exact evolution of \(\langle \sigma_i^x(t)\rangle\) very closely and remain stable and accurate for significantly longer times than the earlier CNN-based TDVP baseline; on \(8\times 8\), the two newer architectures remain mutually consistent beyond the time where the older baseline has already drifted [2503.10462].

## 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 [2405.18874]. In a one-layer ViT with embedding dimension \(d=60\), \(h=10\) heads, and the \(6\times 6\) frustrated \(J_1\)-\(J_2\) model at \(J_2/J_1=0.5\), T5 attention maps differ across inputs at initialization but are essentially identical across different input configurations after optimization [2405.18874]. For Decoupled attention, the input-dependent component averaged over 6000 samples is close to an identity matrix plus a constant offset, \(\alpha_{ij}^{\rm input-dep}\approx c+\delta_{ij}\), while the meaningful off-diagonal structure resides entirely in the positional term \(p_{i-j}\) [2405.18874]. The paper therefore concludes that both T5 and Decoupled attention spontaneously reduce to the position-only Factored form under VMC optimization [2405.18874].

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 [2405.18874]. In that exact construction, one head uses \(\alpha^{(1)}_{ij}=\delta_{ij}\), the other uses \(\alpha^{(2)}_{ij}=\delta_{j,S(i)}\), and the attention maps encoding which patches are entangled are completely positional; spin dependence enters only through values and nonlinearities [2405.18874]. Second, invoking the cluster property
\[
\lim_{|i-j|\to\infty}\langle \hat A_i\hat A_j\rangle
=
\langle \hat A_i\rangle\langle \hat A_j\rangle
\]
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 \(Q\) and \(K\) should become ineffective in the large-distance limit [2405.18874].

The CTWF work presents a different conclusion, but on different architectural grounds [2503.10462]. 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 [2503.10462].

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 \(6\times 6\) lattices [2405.18874]. 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 [2503.10462]. 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 [2211.05504][2503.10462]. 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 [2405.18874]. 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 [2405.18874][2503.10462].

Source: https://www.emergentmind.com/topics/transformer-wave-functions