nan when training with 'adam'

[lzx0106]

I tried to obtain the ground state of BeH with BeH.py:

import sys

from absl import logging from ferminet.utils import system from ferminet import base_config from ferminet import train logging.get_absl_handler().python_handler.stream = sys.stdout logging.set_verbosity(logging.INFO) cfg = base_config.default() cfg.system.electrons = (3,2) # (alpha electrons, beta electrons) cfg.system.molecule = [system.Atom('Be', (1.3269, 0.0, 0.0)), system.Atom('H', (0.0, 0.0, 0.0))] cfg.batch_size = 256 cfg.pretrain.iterations = 100 cfg.optim.optimizer = 'adam' train.train(cfg)

After 3454 steps, I got a nan as below:

INFO:absl:Step 03447: -14.9133 E_h, variance=0.0626 E_h^2, pmove=0.52 INFO:absl:Step 03448: -14.9394 E_h, variance=0.1356 E_h^2, pmove=0.54 INFO:absl:Step 03449: -14.9157 E_h, variance=0.0708 E_h^2, pmove=0.54 INFO:absl:Step 03450: -14.9189 E_h, variance=0.0347 E_h^2, pmove=0.53 INFO:absl:Step 03451: -14.9058 E_h, variance=0.0451 E_h^2, pmove=0.53 INFO:absl:Step 03452: -14.9168 E_h, variance=0.0460 E_h^2, pmove=0.52 INFO:absl:Step 03453: -14.9193 E_h, variance=0.1345 E_h^2, pmove=0.52 INFO:absl:Step 03454: -14.9189 E_h, variance=0.0443 E_h^2, pmove=0.51 INFO:absl:Step 03455: nan E_h, variance=nan E_h^2, pmove=0.51 INFO:absl:Step 03456: nan E_h, variance=nan E_h^2, pmove=0.00 INFO:absl:Step 03457: nan E_h, variance=nan E_h^2, pmove=0.00 INFO:absl:Step 03458: nan E_h, variance=nan E_h^2, pmove=0.00 INFO:absl:Step 03459: nan E_h, variance=nan E_h^2, pmove=0.00 INFO:absl:Step 03460: nan E_h, variance=nan E_h^2, pmove=0.00

[jsspencer]

Your script does not set the optimizer to adam. Either way, some combination of:

• check you're using latest versions of JAX and KFAC, if not upgrade
• use more pretrain iterations
• use a larger batch size
• adjust the learning rate and norm constraint [if using KFAC]
• adjust the learning rate [if using ADAM] -- note that the learning rate set in base_config.default() is tuned for larger systems with KFAC and is probably not suitable for ADAM

should help. EIther way, your script runs for 50k iterations for me without a problem.