Unstable NequIP potential using dataset that worked with PaiNN

[gavinwinter]

I am running NVT with the NequIP potential, and it seems like the disorder in the supercell is causing problems, where specific atoms just get locked in a site and things start freezing (after only 10ps), With the same dataset, training a potential with PaiNN seems to give stable MD for nanoseconds, but with the same dataset training a potential with NequIP seems pretty unstable (at least with off-the-shelf hyperparameters). Any thoughts on what the issue may be? Updating the neighbor list every 5 steps seems pretty reasonable, as with PaiNN we can get away with updating it every 20 steps.

I have attached the LAMMPS input file, but I think the parameters for MD are pretty standard.

    # ---------- Initialize simulation ---------------------
    units real
    atom_style atomic
    newton off
    dimension  3
    boundary   p p p
    read_data /home/gridsan/gawinter/projects/labutil/lammps.data
    
    pair_style nequip
    pair_coeff * * /home/gridsan/gawinter/projects/other_nff/results/lgps/example-lgps/lgps-deployed.pth  Li Ge P S
    mass            1 6.941
    mass            2 72.59
    mass            3 30.97376
    mass            4 32.06

    neighbor 1.0 bin
    neigh_modify delay 5 every 1

    velocity  all create 300 87287 dist gaussian
    # ---------- Describe computed properties------------------
    thermo_style custom step pe ke etotal temp
    thermo 100
    dump 1 all custom 20 dump.md.1.* id type xs ys zs
    # ---------- Specify ensemble  ---------------------
    # fix  1 all nve
    fix 1 all nvt temp 300 300 50
    # --------- Run -------------
    timestep 1
    run 100000

[simonbatzner]

Can you share the NequIP config file. I see units real in LAMMPS, are your data in kcal/mol? The comparison to PaiNN doesn't tell you much, this looks like something is just set up wrong or unconverged, PaiNN is a subset of NequIP and we have consistently seen better results.

[gavinwinter]

Realized I never answered your other question. Yes, my dataset units are in kcal/mol and kcal/mol-Angs (to agree with units real).

[simonbatzner]

You're only training for 100 epochs, so your training is probably not converged. Note that comment in configs/example.yaml:

max_epochs: 100                                                                    # stop training after _ number of epochs, we set a small number here to have an example that finished within a few minutes, but in practice we recommend using a very large number, as e.g. 1million and then to just use early stopping and not train the full number of epochs`

Set this to a very high number and then stop your training when the LR has dropped past a certain value, e.g. 1e-5, or the validation loss has not improved for _ epochs, e.g. 50 or 100:

# early stopping based on metrics values.
early_stopping_patiences:                                                          # stop early if a metric value stopped decreasing for n epochs
  validation_loss: 50

early_stopping_lower_bounds:                                                       # stop early if a metric value is lower than the bound
  LR: 1.0e-5

[gavinwinter]

Thanks, Simon. I’ll try that and do a convergence check along with a test evaluation. I’ve just previously seen convergence/plateau within 100 epochs for training PaiNN and SchNet, so I figured it was fine. I’ll check more carefully for convergence this time.

[gavinwinter]

That was part of the problem, yes. In addition, I did not read things carefully enough: I interpreted n_train to be something like the batch size, so I was only training on 100 geometries out of my dataset of ~28,000. (1) I set n_train to 17,000 (so as to have a reasonable train-val-test split) and (2) set the max_epochs to a high number, and everything seems good now. Training did not plateau by the convergence criterion I set, but I stopped training after 1000 epochs, and the potential seems stable for at least 100ps now and is probably good for more. So, the issue seems to be resolved. Thanks.

[simonbatzner]

I see, yeah that'll do it. Just make sure to not use a batch size of 100, we found smaller ones work much better. And good to hear things are looking good now.