Restricted Boltzmann machine

Paper title: Solving the quantum many-body problem with artificial neural networks

Paper authors: Giuseppe Carleo and Matthias Troyer

arXiv:1606.02318 (2016)

Science 355, 602 (2017)

In this example, we solve the ground state of 10x10 Heisenberg model by utilizing a restricted Boltzmann machine with channel number \(\alpha=16\).

Related tutorials: Quick start

Estimated cost: 1 RTX4090 x 10 min

The restricted Boltzmann machine (RBM) represents the wavefunction with a single hidden layer. For a spin configuration \(s = (s_1, \dots, s_N)\) with \(s_j = \pm 1\), one introduces \(M = \alpha N\) binary hidden units coupled to the visible spins. Tracing out the hidden units analytically gives the amplitude

\[ \psi(s) = e^{\sum_j a_j s_j} \prod_{i=1}^{M} 2\cosh\!\left(b_i + \sum_j W_{ij}\, s_j\right), \]

where \(a_j\) are the visible biases, \(b_i\) the hidden biases, and \(W_{ij}\) the weights connecting the two layers. The hidden-unit density \(\alpha = M / N\) controls the expressiveness of the ansatz.

In Quantax, RBM_Conv() realizes this as a single convolutional layer with a \(\cosh\) activation, \(\psi(s) = \prod \cosh(\mathrm{Conv}(s))\). The convolution shares weights across lattice translations, so the ansatz is translationally invariant and the channel number plays the role of the hidden-unit density \(\alpha\).

RBM is unstable under TF32 precision, so we turn off TF32 in this example.

import jax
jax.config.update("jax_default_matmul_precision", "float32")

import quantax as qtx
import matplotlib.pyplot as plt
from IPython.display import clear_output
%config InlineBackend.figure_format = 'svg'

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

H = qtx.operator.Heisenberg(msr=True)
model = qtx.model.RBM_Conv(channels=16)
state = qtx.state.Variational(model, max_parallel=10000*32)
sampler = qtx.sampler.SpinExchange(state, nsamples=10000)
solver = qtx.optimizer.auto_shift_eig(rshift=1e-6, ashift=1e-8)
optimizer = qtx.optimizer.SR(state, H, solver=solver)

E_QMC = -268.62107

energy = qtx.utils.DataTracer()

for i in range(1000):
    samples = sampler.sweep()
    step = optimizer.get_step(samples)
    state.update(step * 2e-3)
    energy.append(optimizer.energy)

    if i % 10 == 0:
        clear_output(wait=True)
        energy.plot(start=-200, batch=10, baseline=E_QMC)
        plt.show()

The relative error of variational accuracy, given by

\[ \epsilon_\mathrm{rel} = (E_\mathrm{NQS} - E_\mathrm{QMC}) / |E_\mathrm{QMC}|, \]

is similar to the \(10^{-3}\) result presented in Fig. 3(C) of the original paper

E = energy[-50:].mean()
rel_err = (E - E_QMC) / abs(E_QMC)
print(rel_err)

0.0011072264

Here we reproduce Fig. 2 of the original paper, which shows the weights in RBM. The scale looks different due to training details, but the patterns are similar.

import jax.numpy as jnp
from matplotlib.colors import TwoSlopeNorm

W = state.model.layers[1].weight

# single symmetric color scale (centered at 0)
v = jnp.max(jnp.abs(W))
norm = TwoSlopeNorm(vmin=-v, vcenter=0.0, vmax=v)

fig, axes = plt.subplots(4, 4, figsize=(8, 8), constrained_layout=True)

for i, ax in enumerate(axes.flat):
    im = ax.imshow(W[i, 0], cmap='RdYlBu_r', norm=norm)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title(rf'$W^{({i+1})}$', fontsize=10, fontstyle='italic')

# one horizontal colorbar under all panels
cbar = fig.colorbar(im, ax=axes, orientation='horizontal', pad=0.08, shrink=0.9)

plt.show()