which distribution is used in the pretrain phase? about ferminet HOT 3 CLOSED

Nice spot. Sorry, this is a bug in the TF code introduced when porting parallelisation over multiple GPUs from TF-Replicator to Distribution Strategy. The results in the FermiNet paper used TF-Replicator, which has a slightly different API to Distribution Strategy. With TF-Replicator, we are using the concatenated distribution on each GPU. The line you highlighted should read

concat_data = tf.concat([self.hf_data_gen.walkers, self.data_gen.walkers], axis=1)

We made a couple of different design decisions in the JAX version. We are currently only using the FermiNet as the distribution in pretraining. I think the pretraining distribution is not super critical -- the most important thing is to train the network such that it's closer to the ground state wavefunction and that the determinants aren't extremely low rank. Without pretraining, the network can start with an energy in the 1000s of hartrees (positive!) and optimising the network through that energy landscape is ... let's say painful. Certainly we don't see any noticeable difference in training FermiNet after pretraining in the TF or JAX codes.

from ferminet.

Got it. Thanks a lot for the detailed explanation!

from ferminet.

Whilst we don't believe this makes a significant difference, the TF version on master is now updated to sample from both distributions during pretraining.

from ferminet.