Local updates

Updates in VMC are often local, for instance, flipping a few spins. If the variational wavefunction has specific internal structures, one doesn’t have to recompute the wavefunctions from scratch. Therefore, defining a different computing graph for local updates might greatly accelerate the VMC simulation.

The local update technique has been widely adopted in mean-field fermionic wavefunctions and tensor networks. In this tutorial, we will introduce how to define local updates in your wavefunction.

As a prerequisite, please read Build your network to understand how to build a network in Quantax.

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

qtx.set_default_dtype(jnp.float64)

L = 64
lattice = qtx.sites.Chain(L)

RefModel

Let’s consider a restricted Boltzmann machine (RBM) wavefunction,

\[\begin{split} \begin{aligned} h_i &= \sum_j W_{ij} s_j + b_i \\ \psi &= \prod_i \cosh h_i \end{aligned} \end{split}\]

In this example, we don’t consider possible overflow of \(\psi\) for simplicity.

class RBM(eqx.Module):
    linear: eqx.nn.Linear

    def __init__(self, M: int):
        key = qtx.get_subkeys()
        linear = eqx.nn.Linear(L, M, key=key)
        self.linear = qtx.nn.apply_lecun_normal(key, linear)

    def __call__(self, s: jax.Array) -> jax.Array:
        h = self.linear(s)
        return jnp.prod(jnp.cosh(h))

If \(s\) is flipped locally to generate \(s'\), then

\[\begin{split} \begin{aligned} h_i' &= \sum_j W_{ij} s_j' + b_i = h_i + \sum_{j \in \{ j|s_j' \neq s_j \}} W_{ij} (s_j' - s_j) \\ \psi' &= \prod_i \cosh h_i' \end{aligned} \end{split}\]

Assume the number of hidden units \(M\) is of order \(O(L)\). Then the local update reduces complexity from \(O(L^2)\) to \(O(L)\).

Quantax provides RefModel, a subclass of eqx.Module, for local updates. Here, we construct RBM_Ref as an RBM with local updates.

from typing import Union, Any
from jaxtyping import PyTree


class RBM_Ref(qtx.nn.RefModel):
    linear: eqx.nn.Linear

    def __init__(self, M: int):
        """The same as usual RBM"""
        key = qtx.get_subkeys()
        linear = eqx.nn.Linear(L, M, key=key)
        self.linear = qtx.nn.apply_lecun_normal(key, linear)

    def __call__(self, s: jax.Array) -> jax.Array:
        """The same as usual RBM"""
        h = self.linear(s)
        return jnp.prod(jnp.cosh(h))

    def init_internal(self, s: jax.Array) -> tuple[jax.Array, PyTree]:
        """Compute the initial hidden units for local updates in `ref_forward`."""
        h = self.linear(s)
        psi = jnp.prod(jnp.cosh(h))
        return psi, h
    
    @property
    def required_update_modes(self) -> tuple[str, ...]:
        """
        The required update modes for accelerated ref_forward pass.
        """
        return ("nflips",)

    def ref_forward(
        self,
        s: jax.Array,
        s_old: jax.Array,
        update_mode: dict[str, Any],
        internal: PyTree,
        return_update: bool,
    ) -> jax.Array | tuple[jax.Array, PyTree]:
        """
        Forward pass with reference to the old configuration and the number of flipped spins.
        This is the core function of local updates.
        """
        # A marker indicating that local updates are being compiled.
        print("Using local updates ...")

        nflips = update_mode["nflips"]
        diff = s - s_old
        idx_flip = jnp.flatnonzero(diff, size=nflips)
        h_diff = self.linear.weight[:, idx_flip] @ diff[idx_flip]
        h_new = internal + h_diff
        psi = jnp.prod(jnp.cosh(h_new))
        if return_update:
            return psi, h_new
        else:
            return psi

Be careful that init_internal and ref_forward will be automatically jitted in Quantax. In ref_forward, only update_mode and return_update will be treated as static arguments. Therefore, the shapes of arrays shouldn’t depend on other inputs. For instance, jnp.flatnonzero will trigger jit error if size=nflips is not specified.

The correctness of local updates is checked here.

model = RBM_Ref(4 * L)
s_old = qtx.utils.rand_states()
psi0, internal = model.init_internal(s_old)

update_mode = {"nflips": 1}
s_new = s_old.at[0].multiply(-1)

print("Testing direct forward ...")
psi_direct = model(s_new)
print("Direct forward psi: ", psi_direct)

print("Testing local updates ...")
psi_ref = model.ref_forward(s_new, s_old, update_mode, internal, return_update=False)
print("Local updates psi: ", psi_ref)
assert jnp.isclose(psi_direct, psi_ref)

Testing direct forward ...
Direct forward psi:  1.0715160422898089e+39
Testing local updates ...
Using local updates ...
Local updates psi:  1.0715160422898078e+39

VMC with local updates

To use RefModel in VMC, one needs to wrap it by Variational. Variational with RefModel provides batched and jitted init_internal(), ref_forward(), and ref_forward_with_updates() methods.

state = qtx.state.Variational(model)

s_old = qtx.utils.rand_states(16384)
psi0, internal = state.init_internal(s_old)

update_mode = {"nflips": 1}
s_new = s_old.at[:, 0].multiply(-1)

print("Testing direct forward ...")
psi_direct = state(s_new)

print("Testing local updates ...")
psi_ref = state.ref_forward(s_new, s_old, update_mode, internal)
assert jnp.allclose(psi_direct, psi_ref)
print("Local update test passed!")

Testing direct forward ...
Testing local updates ...
Using local updates ...
Local update test passed!

Then we can test that ref_forward is indeed faster than direct forward pass. The actual improvement depends on the utilized machine.

%timeit jax.block_until_ready(state(s_new))

%timeit jax.block_until_ready(state.ref_forward(s_new, s_old, update_mode, internal))

The samplers in Quantax automatically utilize local updates whenever possible. One should ensure that update_mode in the sampler matches required_update_modes in RefModel. Most built-in samplers include "nflips" in their update_mode. For instance, LocalFlip flips 1 spin at a time, so "nflips" is 1.

sampler = qtx.sampler.LocalFlip(state, nsamples=16384)
print(sampler.update_mode)
samples = sampler.sweep()

Using local updates ...
{'nflips': 1}

You might see memory overflow errors if internal stores too many values. To avoid it, please check the documentation of max_parallel in quantax.state.Variational.__init__().

The local updates are also utilized when computing local energies. To turn on local updates, please call apply_update_mode_filter() to reorder the internal data structure of Operator. Here is an example.

H = qtx.operator.Ising(h=1.0)
H.apply_update_mode_filter(qtx.operator.nflips_filter)
Eloc = H.Oloc(state, samples)

Using local updates ...

To disable local updates, one can set use_ref=False when defining Variational. Using this trick, we can check that local energies computed by local updates and direct forward passes are equivalent.

state_direct = qtx.state.Variational(state.model, use_ref=False)
Eloc_direct = H.Oloc(state_direct, samples)

assert jnp.allclose(Eloc, Eloc_direct)
print("Local energy test passed!")

Local energy test passed!