Design principles of deep translationally-symmetric neural quantum states for frustrated magnets

TL;DR

This work introduces a ConvNext-based neural quantum state (NQS) optimized for translational symmetry to study frustrated magnets. By embedding sublattice priors and employing a depthwise-convolution encoder with a compact readout, the authors demonstrate that ConvNext closely aligns with the factored vision transformer (fViT) and can achieve state-of-the-art-like variational energies for the Shastry-Sutherland and J$_1$-J$_2$ models. Through systematic hyperparameter studies—patching, kernel size, depth, expansion ratio, and readout size—the study provides a concrete blueprint for designing translationally-symmetric NQS capable of tackling challenging ground-state problems in frustrated magnetism. The results highlight the practical viability of translationally-symmetric architectures and their potential extension to larger systems, higher dimensions, and spectroscopy, with implications for identifying exotic quantum spin liquids and related phases.

Abstract

Deep neural network quantum states have emerged as a leading method for studying the ground states of quantum magnets. Successful architectures exploit translational symmetry, but the root of their effectiveness and differences between architectures remain unclear. Here, we apply the ConvNext architecture, designed to incorporate elements of transformers into convolutional networks, to quantum many-body ground states. We find that it is remarkably similar to the factored vision transformer, which has been employed successfully for several frustrated spin systems, allowing us to relate this architecture to more conventional convolutional networks. Through a series of numerical experiments we design the ConvNext to achieve greatest performance at lowest computational cost, then apply this network to the Shastry-Sutherland and J1-J2 models, obtaining variational energies comparable to the state of the art, providing a blueprint for network design choices of translationally-symmetric architectures to tackle challenging ground-state problems in frustrated magnetism.

Figures (8)

• Figure 1: The $J_1$ (black) and $J_2$ couplings of the (red) Shastry-Sutherland and (grey and red) $J_1-J_2$ models. The sites in the unit cell are highlighted for the Shastry-Sutherland (red) and $J_1-J_2$ (blue) models.
• Figure 2: Illustration of a patched translationally-symmetric NQS. The patching and embedding of the input means the remainder of the network operates on sublattice-dependent quantities, a useful bias for learning the sign structure of ground states. Convolution-like encoding blocks then construct a hidden representation out of (high-order) correlators of these quantities, before projection to a momentum, $\mathbf{k}$, and readout of a complex amplitude.
• Figure 3: The ConvNext architecture adapted for use as an NQS. The right-hand side displays the shape and type of the data as it passes through the network. The number of features, $N_f$, expansion ratio, $\alpha$, number of encoder blocks, $n_b$, as well as convolutional kernel width, $k$ and size ($n_o, \alpha_o$) of the feedforward neural network (FFNN) in the readout block are all hyperparameters. The number of patches along the Cartesian axes are denoted $P_x, P_y$. The embedding and encoder have real-valued parameters, with a complex-valued FFNN in the readout head resulting in a complex output.
• Figure 4: Two different ways of patching and embedding the input. The numbers label the patches, which are made up of the variables in red boxes. (Left) The data is first reshaped (patched) to $\mathbb{Z}_2^{P \times N_x}$, where $N_x$ is the number of sublattices created by the patching, before embedding to $\mathbb{R}^{P \times N_f}$. (Right) The data is reshaped to its lattice positions, before a strided convolution simultaneously performs the patching and embedding operations to $\mathbb{R}^{P_x \times P_y \times N_f}$. The two are mathematically equivalent.
• Figure 5: Optimization of the $(4,3, 2)[1,1]$ ConvNext with and without patching for $L = 6$. The respective final energies ($E/N = -0.451717(6)$ and $-0.451633(6)$) are shown with dotted lines and the ground state energy $E_{\mathrm{GS}}/N = -0.4517531$viteritti2024 dashed. Relative errors along with measures of the computational cost are shown in fig. \ref{['fig:summary']}. The annotations refer to the symmetrization stages of the optimization (see Appendix \ref{['app:optimization']}).