Building Transformer-Based NQS for Frustrated Spin Systems with NetKet

Building Transformer-Based NQS for Frustrated Spin Systems with NetKet

This framework implements a research-grade Variational Monte Carlo pipeline using NetKet and JAX to solve frustrated J1-J2 Heisenberg spin chains. It utilizes a Transformer-based architecture to capture complex quantum correlations in high-dimensional systems.

Why This Matters

Many-body physics often hits a computational wall when dealing with frustrated systems because traditional methods cannot efficiently handle the exponential complexity of quantum states. By using a Transformer architecture as a variational ansatz, engineers can leverage self-attention to model long-range correlations that are typically invisible to local-update algorithms.

Key Insights

• Stochastic Reconfiguration (SR) provides natural gradient descent for optimizing quantum wavefunctions in NetKet.
• Transformers utilize multi-layer self-attention to model complex correlations in J1-J2 Heisenberg spin chains.
• The GraphOperator in NetKet efficiently represents interacting spin systems using custom colored graph representations.
• Benchmarking against exact Lanczos diagonalization allows for calculating the absolute energy gap and verifying VMC accuracy.
• Structure factor analysis using FFT on spin-spin correlations helps detect phase transitions in quantum systems.

Working Examples

Hamiltonian construction for the frustrated J1-J2 spin chain.

def make_j1j2_chain(L, J2, total_sz=0.0): J1 = 1.0; edges = []; for i in range(L): edges.append([i, (i+1)%L, 1]); edges.append([i, (i+2)%L, 2]); g = nk.graph.Graph(edges=edges); hi = nk.hilbert.Spin(s=0.5, N=L, total_sz=total_sz); bond_ops = [(J1*mszsz).tolist(), (J2*mszsz).tolist(), (-J1*exchange).tolist(), (J2*exchange).tolist()]; bond_colors = [1,2,1,2]; H = nk.operator.GraphOperator(hi, g, bond_ops=bond_ops, bond_ops_colors=bond_colors); return g, hi, H

Transformer-based Neural Quantum State architecture implemented in Flax.

class TransformerLogPsi(nn.Module): L: int; d_model: int = 96; @nn.compact; def __call__(self, sigma): x = (sigma > 0).astype(jnp.int32); tok = nn.Embed(num_embeddings=2, features=self.d_model)(x); pos = self.param("pos_embedding", nn.initializers.normal(0.02), (1, self.L, self.d_model)); h = tok + pos; for _ in range(self.n_layers): h_norm = nn.LayerNorm()(h); attn = nn.SelfAttention(num_heads=self.n_heads, qkv_features=self.d_model, out_features=self.d_model)(h_norm); h = h + attn; h = nn.LayerNorm()(h); pooled = jnp.mean(h, axis=1); out = nn.Dense(2)(pooled); return out[...,0] + 1j*out[...,1]

VMC training routine using Stochastic Reconfiguration and Metropolis sampling.

def run_vmc(L, J2, n_iter=250): g, hi, H = make_j1j2_chain(L, J2, total_sz=0.0); model = TransformerLogPsi(L=L); sampler = nk.sampler.MetropolisExchange(hilbert=hi, graph=g, n_chains_per_rank=64); vs = nk.vqs.MCState(sampler, model, n_samples=4096, n_discard_per_chain=128); opt = nk.optimizer.Adam(learning_rate=2e-3); sr = nk.optimizer.SR(diag_shift=1e-2); vmc = nk.driver.VMC(H, opt, variational_state=vs, preconditioner=sr); log = vmc.run(n_iter=n_iter, out=None); return vs, np.array(log["Energy"]["Mean"]), np.array(log["Energy"]["Variance"])

Practical Applications

• Quantum Magnetism Simulation: Using NetKet to explore phases beyond the reach of classical exact methods. Pitfall: Insufficient sampling in MetropolisExchange leading to poor convergence.
• Ordering Transition Analysis: Computing structure factor peaks to detect quantum phase changes. Pitfall: Failing to update JAX to enable float64 can lead to numerical instability.

References:

• https://www.marktechpost.com/2026/04/16/transformer-nqs-netket-j1j2-guide/