Neural backflow states for the Bose-Hubbard ground state

[chroli1]

Hi,

I have been playing with the code accompanying this paper for the ground state of the 2D square-lattice Bose-Hubbard model with fixed particle number. I have been struggling to get the parameters right for large systems (say, an 8x8 lattice).

In particular, I find that the energy and gradients for my initial wavefunction blow up, and the optimiser fails. The issue seems to be with the Monte Carlo sampler: if the Markov chains are poorly thermalised, states with vanishing amplitude are sampled, and the 1/\psi factor used when calculating expectation values and gradients diverges.

Is there any general advice/references on setting the MCMC parameters correctly? I think the main ones are:

• The total number of samples
• The number of parallel chains
• The sweep size (i.e. how often I record a sample as I move along the chain)
• The burn-in time (i.e. how many samples to discard at the start of the chain

[PhilipVinc]

cc @Z-Denis (who is possibly on leave AFAIK)

[Z-Denis]

Hi @chroli1 ! I will need a bit more information to be able to pinpoint the origin of the issue. For what nonlinearity is this happening? The stability of the approach will crucially depend on the regime. In the superfluid regime (U/t << 18), you should not encounter any issue with Monte Carlo and your chains should thermalise pretty fast with very low tau corr. Close to the phase transition, the situation is much trickier since the state is roughly a Kronecker delta around the Mott state. This means that, unless you do some fancy importance sampling, your acceptance will dramatically drop and, depending on the number of samples, you might even encounter an ill-defined QGT, causing your preconditioned forces to blow up. Finally, in the Mott regime, standard VMC is expected to fail.

I can list a few of the tricks I used in the paper:

• I first train a bare 2-body Jastrow. Then I create the Jastrow-backflow state, initialise the Jastrow weights to the learnt ones and initialise the weights of the last linear projection (a_\mu in the paper IIRC) to something really perturbative (say a_\mu ~N(\mu=0, \sigma = 1e-4)). This reduces the initial variance quite a lot when starting to train the backflow.
• I run the VMC on a GPU with 1024 chains. Having as many chains as possible is crucial here, especially close to the critical point, to mitigate intra-chain correlation times. Since I have only a few samples per chain, I probably bumped a bit the number of discarded samples per VMC step.
• I always keep an eye on the tau corr (make sure you set: netket.config.netket_experimental_fft_autocorrelation = True at the top of your script). You should increase the sweep size as you become less and less superfluid.
• Make sure you use the right Metropolis rule. If you work at fixed number of particles, you should use the one I provide in the accompanying repo. If you code it yourself, be really really careful with the log corrections of the rule.
• Don't forget to thermalise your state, also after adding the backflow; if you initialise wrongly your weights this thermalisation might take ages.

This is all I can think of, but I am happy to help you if you provide more information. I can send you a job file, if you want.

[Z-Denis]

Forgot to mention it, but the optimisation curves in these Bose-Hubbard experiments (with residual connections) are probably the cleanest I have ever obtained with VMC so the slightest instability should be a strong signal that something is wrong with the hyperparameters. Also forgot to ask: how many chains, samples per chain? what is your sweep size and number of discarded samples?

[chroli1]

Thanks both!

A bit more detail on the physics parameters of the model:

• I'm using U = t = 1, so deep in the superfluid regime, and no chemical potential, instead fixing the particle number at the level of the Fock space
• I use unit filling, so number of particles = number of sites
• I have tried both open and periodic BCs, seems to make no difference
• I used square lattices, looking at sizes from 2x2 up to 8x8

And on the methodology:

• Generally I kept things simple, so I just used the vmc.py script provided with your paper, didn't change anything except the parameters in the JSON file and commenting out the first part with the bare Jastrow ansatz (see below, I think this could be the issue actually!)
• In particular, I've not changed the approach to choosing the NN initial parameters or the Metropolis rule
• I'm running on a single GPU
• To narrow down the issue, I wrote some fresh code where I just generate some samples from your random initial state and evaluate the expectation value of the Hamiltonian, so no optimisation or much going on. From this, I worked out that the issue seems to be states with vanishing amplitude being sampled, leading to expectation values blowing up
• I have used a wide range of parameters and combinations to try to get things working. Long story short, even using a lot of samples (8192) and varying the number of chains, burn-in size and downsampling over wide ranges, I still find this issue (for context, I varied the number of chains from 1 up to all 8192 samples, and the burn-in and downsampling over orders of magnitude too)
• The only thing there that seems to help is making the downsampling really large (maybe)
• Even for parameters that produce 'good' energy EVs, sometimes the R-hat and \tau statistics have quite poor values

Your suggestions are really helpful. I was aware of some of this, which is reassuring! I think I am not doing this two-stage approach of training a bare Jastrow ansatz first, which might be the issue actually. I will try this out.

Regardless, I think this is a more fundamental methodology issue and not just parameters - as you say, it shouldn't be that difficult.

[Z-Denis]

How do you initialise the chains? It is possible that the default hilbert.random.random_state yields configurations with extremely unbalanced populations that make the sampling of the energy and its forces very inaccurate with MC (large variance). The case of a warm-start from two-body Jastrow might be different should the bare two-body Jastrow be less sensitive to that, thus allowing a thermalisation to a decent configuration support before performing the whole optimisation procedure with the backflow transformation.

[Z-Denis]

Actually, having a look at the new constrained logic in NetKet, it seems to initialise the chains to the right thing (essentially the ideal superfluid):

netket/netket/hilbert/random/fock.py

Lines 77 to 109 in a8675bd

	@partial(jax.jit, static_argnames=("n_particles", "hilb_shape", "shape", "dtype"))
	def _random_states_with_constraint_fock(n_particles, hilb_shape, key, shape, dtype):
	# Distribute hilb.n_particles onto hilb.size sites
	# and put at most hilb.shape-1 particles in every site.
	# Note that this is NOT a uniform distribution over the
	# basis states of the constrained hilbert space.
	assert n_particles is not None
	hilb_size = len(hilb_shape)
	# start with all sites empty
	init = jnp.zeros(shape + (hilb_size,), dtype=dtype)
	# if constrained and uniformly n_max == 2, use a trick to sample quickly
	if set(hilb_shape) == {2}:
	return jax.random.permutation(
	key, init.at[..., :n_particles].set(1), axis=-1, independent=True
	)
	# shape is per site n_max
	n_max = jnp.array(hilb_shape) - 1
	def body_fun(x, key):
	# select all sites which are not yet full
	p = x < n_max
	# uniformly select a site among the not yet full ones
	# and put a particle in it
	carry = x + _choice(key, p)
	return carry, None
	# iterate body_fun above n_particles times
	keys = jax.random.split(key, n_particles)
	return jax.lax.scan(body_fun, init, keys)[0]

It is very similar to the logic I used for the paper, if not identical. Another aspect that one might check is that the Jastrow weight matrix be positive and invertible, otherwise thermalising the chains will lead to exploding occupation variances (<n_i n_j> ~ [1/W]_{ij})

[chroli1]

This is now working!

The issue is that I misunderstood the vmc.py script: I thought the plain Jastrow calculation was independent of the other one for comparison purposes, and commented that out to simplify things. I therefore wasn't doing the plain Jastrow pre-training, and unsurprisingly the optimisation of just the NN backflow part failed. Pretty basic error on my part but there we are!

One other general question: we are interested in looking into disorder in the Bose-Hubbard model. Do you expect that major changes will be needed to my code to do this? Perhaps changes to the sampler are needed?

Thanks so much for your help!

[Z-Denis]

That's great news!

Your project sounds interesting. I do think there are major changes to be considered. I depends a bit on the path you follow, but the model has been designed as to be translationally invariant. If you actually perform VMC for each realisation of the disorder, you will have to (at least) promote the Jastrow factor to a full dense symmetric matrix.

Let me know if you want some data or any further help!

[chroli1]

Thanks! This makes sense - I suppose the translational invariance is baked in to the ResNet/convolutional neural network structure. Perhaps swapping to a different neural network entirely is needed?

[Z-Denis]

Before touching the backflow transformation, I would start breaking the translational invariance at the level of the Jastrow factor and by adding a spatial dependence to the final projection vector of the equivariant backflow features (a_\mu into a_{i\mu}).

[chroli1]

Thanks! This could be an interesting project in itself. If anything interesting comes of it I'll let you know. Thanks again for the help.

[Z-Denis]

To add to my previous comment: The CNN just makes the output equivariant, which is not necessarily bad in a disordered system. If you want to change the architecture of the backflow altogether you might try some vision transformer or whatever stable and scalable attention-based model you fancy.