NetKet and JAX: How to Build a Transformer Model for Frustrated Spin Systems

Transformer architectures are beginning to find their place in computational quantum physics: a new practical guide demonstrates how to assemble a complete Neural Quantum States pipeline using NetKet and JAX for a complex problem — the frustrated J1-J2 Heisenberg spin chain. This is not abstract theory, but a reproducible framework where the model, sampler, optimization, and accuracy verification are brought together in a single research loop. The core idea of the guide is that the Transformer architecture is well-suited for describing many-body quantum states, where long-range correlations between particles are important.

Conventional numerical methods quickly run into the curse of dimensionality of the state space, especially if the system is frustrated, meaning competing interactions prevent it from reaching a simple ordered energy minimum. Under such conditions, Neural Quantum States allow the wave function to be represented as a parametrized neural network, which is then optimized through Variational Monte Carlo. NetKet serves as a ready-made environment for quantum computations, while JAX acts as the engine for high-precision accelerated optimization.

The guide first sets up the basic physical part of the problem. The author defines a one-dimensional chain of length L with periodic boundary conditions, where nearest neighbors interact with coefficient J1, and their next-nearest neighbors with J2. This combination is what creates the frustration that makes the problem interesting and non-trivial.

To describe the system, a graph in NetKet is used, a Hilbert space of spin-1/2 particles with fixed total projection, and a Hamiltonian operator assembled through GraphOperator. In parallel, 64-bit precision is enabled in JAX, which is essential for stable calculations in this class of problems. Then machine learning in its pure form begins.

The wave function is defined by a custom TransformerLogPsi model on Flax: spin configurations are encoded as tokens, receive embeddings and positional representations, then pass through several blocks of self-attention and feed-forward layers. The example uses a hidden dimension of 96, four attention heads, and six Transformer layers. The model returns the complex logarithm of the wave function amplitude — this is critical because a quantum state cannot be adequately described by a real scalar alone.

After aggregating information across the entire chain through averaging, the network obtains a global representation of the configuration and can express more complex correlations than local ansatze. A particular value of the guide is that it does not stop at model definition. For training, the author assembles a complete VMC loop: a MetropolisExchange sampler, variational state MCState, Adam optimizer, and Stochastic Reconfiguration as an analog of natural gradient descent for quantum states.

The example configuration uses 4096 samples, rejection of initial states in chains, and roughly 250 optimization iterations. Such a stack is necessary not only to achieve low energy, but also to control convergence. The code saves trajectories of mean energy and variance, so one can see how stably the model moves toward a good solution.

After training, the pipeline is used as a research tool. The author runs calculations for several values of J2 in the range from 0 to 0.7 for a chain of length 24 nodes, records the final energies, and estimates the peak of the structure factor.

This allows not just to tune neural network parameters, but to observe how the physical behavior of the system changes as frustration increases and where transitions between different phases of magnetic order might appear. For additional quality verification, the model is compared against exact diagonalization on a smaller system of size 14 nodes using the Lanczos method. Energy comparison provides a clear numerical benchmark: how close the variational Transformer is to the exact solution where exact calculation is still feasible.

The practical significance of the material is that it bridges the gap between two worlds — modern deep learning architectures and real computational physics problems. For ML engineers, it is a good example of how a Transformer can be used outside of text, images, and standard tabular data. For physicists, it is an understandable template for transitioning from the abstract idea of Neural Quantum States to a reproducible experiment with concrete metrics, benchmarks, and observable quantities.

And for those working at the intersection of these fields, the guide provides a foundation that can be extended further: moving to larger lattices, adding symmetries, studying entanglement, or building more complex time simulations. What this means: the Transformer approach is gradually becoming a working tool not only for classical AI tasks, but also for modeling quantum systems, where the cost of error is high and exact methods quickly run out. If NetKet and JAX are already in your working stack, this material provides a practically ready starting point for research-level experiments.