❓ [QUESTION] Training a model to perform Molecular Dynamics simulations on Copper Formate System

[bartdemooij]

Dear nequip developers,

We recently started training a model on a copper formate system. In essence, it was similar to the system you described in the nequip paper (48 Copper atoms and 1 formate molecule), however we used a slightly bigger system comprising 144 copper atoms and 1 formate molecule.

We generated training data similar to the approach you described, we performed nudged elastic band simulations, with 20 images, from which we chose 14 images to perform AIMD of 500 0.5fs steps at 300K. We then started training a model on this, which obtained a low prediction error on both energies and forces on the test set. However, when performing MD simulation (with LAMMPS), the model quickly broke down (we saw the formate molecule entering the copper surface, and the copper surface disintegrate badly).

We figured we needed more frames at higher temperatures to inform the model better, therefore we performed additional MD simulations at 500, 2000 and 4000 Kelvin. Besides the temperature, they were the same as the previous AIMD calculation. So, in total a dataset of 28000 structures were obtained.

These 28k structures were split into 80% training and 20% testing sets. The training set was then also divided in 80% for training and 20% for validation. The reason for this was that smaller training sets did not obtain good MD trajectories.

We used the following configuration of the network (example of yaml file):

root: results/name
run_name: name
seed: 123
dataset_seed: 456                                                                  
append: true                                                             
default_dtype: float32                                                
allow_tf32: false                                                           

# network
r_max: 5.0                                          
num_layers: 5                        

chemical_embedding_irreps_out: 32x0e                                         
feature_irreps_hidden: 32x0o + 32x0e + 32x1o + 32x1e + 32x2o + 32x2e               
irreps_edge_sh: 0e + 1o                                                          
conv_to_output_hidden_irreps_out: 16x0e              

nonlinearity_type: gate                                                  
resnet: false                                              

nonlinearity_scalars:
  e: silu
  o: tanh

nonlinearity_gates:
  e: silu
  o: tanh

# radial network basis
num_basis: 8                                  
BesselBasis_trainable: true          
PolynomialCutoff_p: 6       

# radial network
invariant_layers: 2                                         
invariant_neurons: 64                                       
avg_num_neighbors: null
use_sc: true                                                
compile_model: false                        

dataset: npz                                  
dataset_file_name: directory_to_dataset.npz       
key_mapping:
  z: atomic_numbers
  E: total_energy                                                
  F: forces                                    
  R: pos                                                      
npz_fixed_field_keys:                               
  - atomic_numbers
 
chemical_symbol_to_type:
  H: 0
  C: 1
  O: 2
  Cu: 3

# logging
wandb: true              
wandb_project: project_name                                               
wandb_resume: true                              
verbose: info                                       
log_batch_freq: 1                                      
log_epoch_freq: 1                                                        
save_checkpoint_freq: -1                                                           
save_ema_checkpoint_freq: -1                                                      

# training
n_train: 18000                                                     
n_val: 4400                                                          
learning_rate: 0.005                                                          
batch_size: 5                                                                      
max_epochs: 100000                                                                     
train_val_split: random                                                           
shuffle: false #because shuffle data beforehand                                          
metrics_key: validation_loss                                                       
use_ema: true                                                                      
ema_decay: 0.99                                                                   
ema_use_num_updates: true                                                         
report_init_validation: false

early_stopping_patiences:                                                       
  validation_loss: 50
early_stopping_delta:                                                             
  validation_loss: 0.01
early_stopping_cumulative_delta: false                                           
early_stopping_lower_bounds:                                                       
  LR: 1.0e-6
early_stopping_upper_bounds:                                                       
  wall: 1.0e+100

# loss function
loss_coeffs:                                                                       
  forces: 21904 #i.e. N^2 (=148^2)                                                                       
  total_energy:                                                                    
    - 1
    - MSELoss

metrics_components:
  - - forces                          
    - rmse                               
    - PerSpecies: True                  
      report_per_component: False      
  - - forces
    - mae
    - PerSpecies: True
      report_per_component: False
  - - total_energy
    - mae
    - PerAtom: True             
  - - total_energy
    - mae
    - PerAtom: False

# optimizer
optimizer_name: Adam                                                        
optimizer_amsgrad: true
optimizer_betas: !!python/tuple
  - 0.9
  - 0.999
optimizer_eps: 1.0e-08
optimizer_weight_decay: 0

max_gradient_norm: null

lr_scheduler_name: ReduceLROnPlateau
lr_scheduler_patience: 100
lr_scheduler_factor: 0.5

per_species_rescale_scales_trainable: false
per_species_rescale_shifts_trainable: false
per_species_rescale_shifts: dataset_per_atom_total_energy_mean

global_rescale_shift: null
global_rescale_scale: dataset_forces_rms
global_rescale_shift_trainable: false
global_rescale_scale_trainable: false

Early stopping was triggered after approximately 1350 epochs and obtained the following errors (eV). These are quite okay, only maybe that E_mae could be better but that is due to the loss function (weight ratio F:E = 148^2:1).

            0_f_rmse =  0.213254
            1_f_rmse =  0.098634
            2_f_rmse =  0.097583
            3_f_rmse =  0.046718
          all_f_rmse =  0.114047
             0_f_mae =  0.075106
             1_f_mae =  0.043834
             2_f_mae =  0.042400
             3_f_mae =  0.025226
           all_f_mae =  0.046642
             e/N_mae =  0.035423
               e_mae =  5.242538

Unfortunately, upon performing MD simulations, the system was again acting in a non-physical way.
Below are some screenshots at start, 50fs and 150fs, respectively:

We used the following MD example configuration in LAMMPS:

units real
newton off
read_data structure.data

pair_style nequip
pair_coeff * * model-deployed.pth C Cu H O
mass  1 12.0107
mass  2 63.546
mass  3 1.00794
mass  4 15.999

# Run MD
timestep 1.0
dump 1 all custom 10 traj_nvt.lammpstrj id type x y z ix iy iz
velocity all create 300.0 4928459 

# temp 300K and pressure 1 atm
fix fxnvt all nvt temp 300 300 100.0 tchain 4
fix fxlmom all momentum 10 linear 1 1 1

thermo 10
thermo_style custom step etotal ke temp pe ebond eangle edihed eimp evdwl ecoul elong press vol cella cellb cellc density

run 10000

write_data system_after_nvt.restart
write_data system_after_nvt.data

Do you have any insights into what might be going wrong? Have you tried (and succeeded) to perform MD simulations on "Heterogeneous catalysis of formate dehydrogenation" section of your paper?

Thanks in advance,

Jim Boelrijk and Bart de Mooij

[nw13slx]

Hi Jim and Bart,

Yes. I did manage to run MD simulations with my formate model. It is quite a tricky thing to tune and it involves one unpublished feature that I'm writing about. If you don't mind, you can contact us via email and we can chat more on that.

One quick observation is that your force error is on the high end. The C, H, and O are all a little bit too high. It may be worth trying to the particular species '0' lower than 0.05 eV/A for MAE.

Best,
Lixin

[nw13slx]

And I think Simon did upload our Formate+Cu data, the version used in the paper. Maybe you can try that dataset. @simonbatzner

[simonbatzner]

Hi @bartdemooij, thank you for the detailed issue. As @nw13slx, both her and I were able to run MD simulations (independently of each other) for the system. Also as Lixin said, this is one of the hardest systems we've encountered.

Looking at your LAMMPS input file, it looks like you're using real units, which means energies and forces are expected to be in kcal/mol and kcal/mol/A, respectively. You say that your errors are in eV. That means there is a mismatch and my guess is that this is the reason for the unstable simulations?

[simonbatzner]

A few other comments:

• you're using a time step of 1 fs, that might be doable but is certainly on the high end for a system containing a C-H bond, I'd recommend a maximum of 0.5f fs.
• I would also try to use a bit more of an aggressive thermostat in the first few ps until the temperature has settled, then switch to a more mild one
• regarding thermostats: I think Langevin thermostat is a better fit here for the small molecule
• Yes, the data set is public

[bartdemooij]

Hi Lixin and Simon,

Thank you for the quick response. The MD simulation looks considerably better now without the unit mismatch.

Also, we observed a weird spike in the validation and training loss plots (taken from w&b). Maybe you know an explanation for this as it looks quite weird to me?
Training Loss:

Validation Loss:

To speed up future model generation/training, would you mind telling me how large your training, validation and test sets are to obtain an accurate MD simulation? Is there a rule of thumb of how much training data to use as the equivariance reduces the required training data to get a good model?

Best,

Bart

[simonbatzner]

Hi Bart, that's good to hear.

Re the spikes, this can sometimes happen in Deep Learning (and is not particular to NequIP). What happens is that during training the network encounters a bad batch which leads to large gradients which leads to a bad step in weight space, causing the network to blow up. Usually this recovers to the loss value it was at before and keeps improving. In your case, it looks however like the learning rate scheduler kicked in. If there is no improvement for _ epochs, the LR will be reduced. Due to the spike, there was no improvement and the LR dropped.

We are looking into ways of fixing this adaptively, but I think more generally no perfect solution in Deep Learning exists. Until then, one thing that works is to either reduce the LR (0.001 should also work well) or increase the batch size.

Regarding train/val/test set sizes, that is as you can imagine highly system specific. For many bulk systems, 1,000 is a good start. For heterogeneous catalysis, including bond breaking, part of the challenge is that you have a super diverse energy landscape with many different possible states. Reactions are notoriously tricky to get right with FFs. In our case, 2.5k was enough for good diffusion, but @nw13slx will know more details (or feel free to email us).

[bartdemooij]

Hi Simon,

Thank you for your explanation. That makes sense.

[simonbatzner]

Converting this to a discussion. Feel free to reach out via email if you have more questions.