load params and continue optimization

[dambuck]

Hello,

I am trying to optimize a vqs in multiple steps by decreasing the learning rate and increasing the sample size each step. I have a class which I use to initialize the whole thing.

class my_model():
    _init__(self):
    - build graph, sampler, machine, optimizer
    run(self):
          driver.run(....)

So my skript is basically:
VMC=model(lr=0.1, n_samples=1000)
VMC.run(out='prelim')

VMC=model(lr=0.01, n_samples=2*1000, load_state='prelim')
VMC.run(out='mid')

...

After every optimization I initialize a new class where I load the parameters of the prior optimization via:

with open('prelim.mpack', 'rb') as: 
    VMC.vqs.variables = flax.serialization.from_bytes(VMC.vqs.variables, file.read())

So I am not saving anything myself, but I rely on the standard json logger which saves the parameters only.
However the optimization never continues close to the last energy. The Thermalization through n_discard is never changed. This is also affecting the case when load some parameter and want to only sample some observables without doing optimization . What am I missing?

Thanks

[PhilipVinc]

I don't understand what is your question and what the plots represent. Could you be more specific?

However let me mention two things:

• Why are you recreating the VMC every time? You can simply change the number of samples by doing vqs.n_samples = 12345 to change the number of samples.

• Why do you rebuild the VMC driver every time? You can simply change the optimizer by doing VMC.optimizer = optax.adam(lr=0.02) for example...

You also say that

However the optimization never continues close to the last energy.

This might be because you are not sufficiently well thermalised in your chains. Does this also happen if you severely increase n_discard? By default, our MCMC chains are initialised when you create the MCState/change n_samples from a (uniformly) random state, and then as you sample along the chain slowly thermalises, so even if the samples at step 1 are not thermalised, at step 10 they most likely are. When you recreate MCState or change n_samples, it's possible that your chains don't thermalise well.

A way to avoid this is to sample a very large number of samples after creating MCState or changing n_samples by doing vqs.sample(chain_length=10000), and then discarding them.

[dambuck]

Thanks for the answer, the thermalization issue was my question actually. The second plot shows the energy when I continue the optimization with the parameters of the vqs in the lines before. I expected it to start close to the energy of the first plot.

• When using vqs.sample(chain_length=10000), compared to vqs.sample(n_samples=10000) . What is the difference here? The doc string says chain_length is the MC chain length, is that the chain length on every rank? (i use mpi)
  Is something like:

vqs.sample(chain_length=10000)
vqs.n_samples=5000
vqs.expect(O)

nonesense? Because changing n_samples gets rid of the thermalization I did the line before. Also does the expect function compute the Observable based on the samples sampled in the first line?

• when caling vqs.sample(chain_length=10000), are the samples saved?

About your mentions, thats just me being stupid.

[PhilipVinc]

Ah sorry, I confused myself. Changin n_samples does not reset the state or the sampler, so it does not reset the thermalisation.

In general n_samples= Chain_length*n_ranks*n_chains. And if you set n_samples we compute chain length accordingly.