Minimum-norm stochastic reconfiguration

Paper title: Empowering deep neural quantum states through efficient optimization

Paper authors: Ao Chen and Markus Heyl

arXiv:2302.01941 (2023)

Nat. Phys. 20, 1476-1481 (2024)

Related tutorials: Square J1-J2 model

Estimated cost:

This example shows how to solve the ground state of the 10x10 J1-J2 Heisenberg model. The implementation details are not identical to the original paper, as there are several improvements over recent years.

Due to the high cost of simulations, we recommend running the code on a GPU cluster. The script for different clusters may differ a lot. Here, we show an exemplary script to run the code on the Perlmutter supercomputer of NERSC.

#!/bin/bash -l
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=4
#SBATCH --gpu-bind=none
#SBATCH --mem=0
#SBATCH --output=gpu-out.%j
#SBATCH --error=gpu-err.%j
#SBATCH --time=12:30:00
#SBATCH --constraint=gpu
#SBATCH --qos=premium
#SBATCH -c 128
#SBATCH --account=<account_name>

source <path_to_quantax_env>/bin/activate

export XLA_FLAGS="--xla_gpu_shard_autotuning=false"

srun --cpu-bind=socket python -u <python_script>.py

There are several caveats before we start the simulation.

1. We start 1 process for each node (4 GPUs). jax.distributed.initialize doesn’t handle this case correctly by default. One needs some extra efforts to make it work.

2. We use reweight=1.0 in the Monte Carlo sampler, which improves the sampling distribution and the training accuracy in small systems.

import os
import numpy as np
import matplotlib.pyplot as plt
import jax

gpus_per_node = int(os.environ["SLURM_GPUS_ON_NODE"])

jax.distributed.initialize(
    num_processes=int(os.environ["SLURM_NTASKS"]),
    process_id=int(os.environ["SLURM_PROCID"]),
    local_device_ids=list(range(gpus_per_node)),
)

import jax.numpy as jnp
import jax.random as jr
import equinox as eqx
import quantax as qtx

lattice = qtx.sites.Square(10, Nparticles=(50, 50))
N = lattice.Nsites

H = qtx.operator.Heisenberg(J=[1.0, 0.5], n_neighbor=[1, 2], msr=True)

model = qtx.model.ResConv(
    nblocks=8,
    channels=32,
    kernel_size=3,
    final_activation=qtx.nn.sinhp1_by_scale,
)

state = qtx.state.Variational(model, max_parallel=16384)
sampler = qtx.sampler.SpinExchange(state, 10000, reweight=1.0, n_neighbor=[1, 2])
optimizer = qtx.optimizer.SR(state, H)

energy = qtx.utils.DataTracer()
VarE = qtx.utils.DataTracer()
step_norm = qtx.utils.DataTracer()

for n in range(10000):
    samples = sampler.sweep()
    step = optimizer.get_step(samples)
    norm = jnp.linalg.norm(step)
    state.update(step * 1e-3)

    energy.append(optimizer.energy)
    VarE.append(optimizer.VarE)
    step_norm.append(norm)

    if jax.process_index() == 0:
        print(n, optimizer.energy, optimizer.VarE, norm)

    if n % 10 == 0 and n:
        energy.plot(batch=10, start=-1000)
        plt.savefig("energy.pdf")
        plt.clf()
        VarE.plot(batch=10, start=-1000)
        plt.savefig("VarE.pdf")
        plt.clf()
        step_norm.plot(batch=1, start=-1000, logy=True)
        plt.savefig("step_norm.pdf")
        plt.clf()
    if n % 100 == 0:
        state.save(f"params_{n}.eqx")

state.save("params_J1J2_ResConv.eqx")

Then we can impose symmetries on the trained state. As explained in Square J1-J2 model, the symmetry group is C4v (Rotation and Mirror) x Z2 (Spin inverse). The translation symmetry is already encoded in the CNN architecture.