Hands on implementation of neural-network quantum states

Stochastic reconfiguration:

Theory of neural-network quantum states (NQS)
State representations
Implementation details
Open tasks

Theory of neural-network quantum states (NQS)

Variational Quantum Monte Carlo (VQMC)

Given a quantum system with Hamiltonian $H$ , VQMC aims to solve for the ground state of the system by minimizing the energy of the the system $\mathcal E$ which is always greater then the ground state energy $\mathcal{E}_0$ .

To do this, one first propose a trial wave function $\ket{\psi}$ with parameters $\theta$ , the energy is thus given by

\begin{aligned} \mathcal{E}[\psi] &=\frac{\langle\psi|\mathbf{H}| \psi\rangle}{\langle\psi | \psi\rangle}=\frac{\sum_{\mathbf{x}, \mathbf{x}^{\prime}}\langle\psi | \mathbf{x}\rangle\left(\mathbf{x}|\mathbf{H}| \mathbf{x}^{\prime}\right\rangle\left\langle\mathbf{x}^{\prime} | \psi\right\rangle}{\langle\psi | \psi\rangle}=\frac{\sum_{\mathbf{x}}|\psi(\mathbf{x})|^{2}\left(\sum_{\mathbf{x}^{\prime}} H_{\mathbf{xx}^{\prime}} \frac{\psi\left(\mathbf{x}^{\prime}\right)}{\psi(\mathbf{x})}\right)}{\sum_{\mathbf{x}}|\psi(\mathbf{x})|^{2}} \\ &=E\left[\sum_{\mathbf{x}^{\prime}} H_{\mathbf{x} \mathbf{x}^{\prime}} \frac{\psi\left(\mathbf{x}^{\prime}\right)}{\psi(\mathbf{x})}\right] = E\left[\mathcal{E}_{\rm loc}\right] \end{aligned}

where we have defined the local energy $\mathcal{E}_{\rm loc}$ as

\mathcal{E}_{\rm loc} = \sum_{\mathbf{x}^{\prime}} H_{\mathbf{x} \mathbf{x}^{\prime}} \frac{\psi\left(\mathbf{x}^{\prime}\right)}{\psi(\mathbf{x})}

To minimize the energy, we find the gradient of $\mathcal E$ with respect to the parameters $\theta$ in the trial wave function

\dfrac{\partial \mathcal{E}[\psi]}{\partial \theta_{j}}

and we perform gradient descent update

\theta_{i} \leftarrow \theta_{i}-\alpha \frac{\partial \mathcal{E}[\psi]}{\partial \theta_{i}}

Sampling

One of the key procedure is to evaluate the energy of the system. Since a direct evaluation is often impossible, one can evaluate the energy by Monte Carlo method, i.e., generating a probability distribution according to $|\psi(\mathbf {x})|^2$ . Here, for simplicity, we consider Metropolis-Hastings algorithm where the states are updated according to

P\left(\mathbf{x}^{(k)} \rightarrow \mathbf{x}^{(k+1)}\right)=\min \left(1,\left|\frac{\psi\left(\mathbf{x}^{(k+1)} ; \theta\right)}{\psi\left(\mathbf{x}^{(k)} ; \theta\right)}\right|^{2}\right)

Other choices of the sampler include

NQS as the trial wave function

We consider a restricted Boltzmann machine a our neural network given by

\psi(\mathbf{x}, \mathbf{h} ; \theta)=\exp \left(\mathbf{a}^{\mathrm{T}} \mathbf{x}+\mathbf{b}^{\mathrm{T}} \mathbf{h}+\mathbf{h}^{\mathrm{T}} \mathbf{W} \mathbf{x}\right)

After summing up the hidden units, the marginal wave function with respect to the visible units $\mathbf{x}$ (physical spin) is given by

\psi(\mathbf{x} ; \theta)=\sum_{\mathbf{h}} \psi(\mathbf{x}, \mathbf{h} ; \theta)=\exp \left(\mathbf{a}^{\mathrm{T}} \mathbf{x}\right) \prod_{i=1}^{M} 2 \cosh \left(b_{i}+\sum_{j=1}^{N} W_{i j} x_{j}\right)

The above wave function would be our trial wave function in VQMC.

State representations

In order to derive the equation, the spin states are represented as 1 and -1 (for cosh). For actual implementation, the state can be represented explicitly as a vector of 1 or -1. Alternatively, one could also represent the state as an integer and operate on it via bit operations. The bit operations could be highly efficient if one make use of the extension sets such of BMI, ABM, AVX512.

A nice tutorial has been made by Phillip Weinberg and Marin Bukov (Authors of QuSpin pacakge) to explain the bit basis representations: A tutorial on QuSpin's user_basis

Implementation details

Constructing the RBM (NQS)

mutable struct RBM{T <: Number}
  _nv::Int
  _nh::Int
  a::Vector{T}
  b::Vector{T}
  W::Matrix{T}
  theta::Vector{T}
end

n_visible(rbm::RBM) = rbm._nv
n_hidden(rbm::RBM) = rbm._nh
v_nodes(rbm::RBM) = rbm.a
h_nodes(rbm::RBM) = rbm.b
weight(rbm::RBM) = rbm.W
thetas(rbm::RBM, x::Vector{Int})  = rbm.b + rbm.W * x

# Constructor for RBM
function RBM(nv, nh; sigma=0.01) 
  d = Normal(0, sigma)
  a = init_weight * (rand(d, nv)  + im * rand(d, nv)) 
  b = init_weight * (rand(d, nh)  + im * rand(d, nh))
  W = init_weight * (rand(d, nh, nv)  + im * rand(d, nh, nv))
  return RBM(a, b, W, nv, nh)
end

\psi(\mathbf{x} ; \theta)=\sum_{\mathbf{h}} \psi(\mathbf{x}, \mathbf{h} ; \theta)=\exp \left(\mathbf{a}^{\mathrm{T}} \mathbf{x}\right) \prod_{i=1}^{M} 2 \cosh \left(b_{i}+\sum_{j=1}^{N} W_{i j} x_{j}\right)

# Get the probability
function ψ(rbm::RBM, x::Vector{Int})
  a = v_nodes(rbm)
  b = h_nodes(rbm)
  W = weight(rbm)
  reduce(*, [exp(sum(a .* x)); 2 * cosh.(b + W * x)])
end

# Derivatives on RBM
function derivatives(rbm::RBM, state::Vector{Int})
  ∂a = state 
  ∂b = tanh.(rbm.b + rbm.W * state)
  ∂W = transpose(∂a * transpose(∂b))
  [∂a; ∂b; ∂W[:]]
end

Implementing the sampler

Sampler struct and constructor

mutable struct Sampler
  rbm::RBM 
  n_thermal::Int
  n_sweeps::Int
  state::Vector{Int}  
end

function Sampler(rbm::RBM, n_thermal, n_sweeps)
  state = rand(-1:2:1,rbm._nv)
  return Sampler(rbm::RBM, n_thermal, n_sweeps, state)
end

n_thermal(s::Sampler) = s.n_thermal
n_sweeps(s::Sampler) = s.n_sweeps
n_bins(s::Sampler) = s.bins
get_state(s::Sampler) = s.state

Thermalization and sweep

The samples have to be taken after "system" thermalizes, i.e., there is no correlation with the given random states (check this). This check be checked by measure the auto-correlation.

function thermalize!(s::Sampler)
  L = n_visible(s.rbm)
  s.state = rand(-1:2:1, L)
  for i = 1:s.n_thermal
    single_sweep!(s)
  end
end

function single_sweep!(s::Sampler)
  L = n_visible(s.rbm)
  x = s.state
  for i = 1:L
    x_trial = copy(x)
    site = rand(1:L);
    x_trial[site] = -x[site];
    p = abs2(ψ(s.rbm, x_trial) / ψ(s.rbm,x))
    if rand() < min(1,p)
      s.state = x_trial
    else 
      s.state = x;
    end
  end
    s.state 
end

Implementing the VQMC

Evaluation of local energy

\mathcal{E}_{\rm loc} = \sum_{\mathbf{x}^{\prime}} H_{\mathbf{x} \mathbf{x}^{\prime}} \frac{\psi\left(\mathbf{x}^{\prime}\right)}{\psi(\mathbf{x})}

Given an input state $\mathbf{x}$ we want to know the connected elements $H_{\mathbf{xx'}}$ and the corresponding amplitude ratio.

The implementation below is only applicable for transverse Ising field model. The coupled_states function gives all the connected states $\mathbf{x}'$

function E_local(rbm::RBM, initial_state) :: ComplexF64
  n = length(initial_state)
  X = coupled_states(initial_state) 
  E_loc = 0
  for i = 1:n-1 
    E_loc +=  J * initial_state[i] *  initial_state[i+1];
  end
  for k = 2:n+1
    new_state = X[k,:]
    E_loc += ψ(rbm, new_state)/ψ(rbm, initial_state) * h
  end
  return E_loc
end

# Given a state x, output the coupled states x'
function coupled_states(x)
  n = length(x);
  X = zeros(Int,n+1,n);
  X[1,:] = x;
  for i = 1:n
    X[i+1,:] = x;
    X[i+1,i] = -x[i];
  end
  return X;
end

Implementing parameter updates and optimization

run function computes the averaged gradient

\braket{\dfrac{\partial \mathcal{E}[\psi]}{\partial \theta_{j}}}

and is then passed to update! to update the network parameters.

function Base.run(s::Sampler) 
  thermalize!(s)
  E_loc, ∂θ, E_loc∂θ  = measure(s.rbm, s.state)
  E_loc_avg = zero(E_loc) 
  ∂θ_avg = zero(∂θ)
  E_loc∂θ_avg = zero(E_loc∂θ)
  gNQS = 0
  for j = 1:s.n_sweeps
    single_sweep!(s)
    E_loc, ∂θ, E_loc∂θ  = measure(s.rbm, s.state)
    E_loc_avg += E_loc
    ∂θ_avg += ∂θ
    E_loc∂θ_avg += E_loc∂θ
  end
  E_loc_avg /= s.n_sweeps
  ∂θ_avg /= s.n_sweeps
  E_loc∂θ_avg /=s.n_sweeps
  gNQS =  E_loc∂θ_avg - E_loc_avg * conj(∂θ_avg)
  E_loc_avg, gNQS
end

function update!(rbm::RBM, α, ∂θ )
  rbm.a -= α * ∂θ[1:n_visible(rbm)]
  rbm.b -= α * ∂θ[n_visible(rbm)+1:n_visible(rbm) + n_hidden(rbm)] 
  rbm.W -= α * reshape(∂θ[n_hidden(rbm)+n_visible(rbm) + 1:end], (n_hidden(rbm), n_visible(rbm)))
  return 
end

Open tasks

Hamiltonian interface
Other Machine/NQS such as FNN, MPS
Other samplers such as Swedsen-Wang
Boson bit

^[1]G. Carleo and M. Troyer, Science 355, 602 (2017)
^[2]??, ??{http://netket.org}
^[3]A Neural Network Approach to the Quantum Many-Body Problem ^[4]??{http://weinbe58.github.io/QuSpin/user_basis.html}