Reproduce Results of Nequip on MD17 CCSD/CCSD(T) datasets

[lsj2408]

Thanks for your great works!

By the way, can you provide the necessary configs for reproducing the Table III results (molecules in MD17, labeled by CCSD/CCSD_T methods) in the NequIP paper https://arxiv.org/pdf/2101.03164v3.pdf ?

Looking forward to your help~

[simonbatzner]

Hi @lsj2408, of course!

Please find below the parameters I used for the CCSD/CCSD(T) training on the example of the CCSD aspirin case. Please note that this was run on nequip v0.3.3 (our current main branch is on v0.5), you can find the exact git commit and the versions of other codes all in the paper. If you're just interested in getting highly accurate models, I would suggest you run v0.5, there should not be any major differences, but some of the option names have changed. For a full documentation of all possible current options see configs/full.yaml. Let me know if you have questions.

The avg_num_neighbors differs from molecule to molecule, I used:

aspirin: 10
benzene: 9
ethanol: 8
malonaldehyde: 8
toluene: 10

root: ccsd-aspirin-n950-l3-root
run_name: aspirin-ccsd-n950-l3-run_name
workdir: aspirin-ccsd-n950-l3-workdir

requeue: true
seed: 0                                                                           
append: true                                                                  
default_dtype: float32                                                           
allow_tf32: false                                                           

# network
r_max: 4.0                                                       

num_layers: 5                                                                      
chemical_embedding_irreps_out: 64x0e                                           
feature_irreps_hidden: 64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e + 64x3o + 64x3e                 
irreps_edge_sh: 0e + 1o + 2e + 3o                                                
conv_to_output_hidden_irreps_out: 16x0e                         

nonlinearity_type: gate                                                          

nonlinearity_scalars:
  e: silu
  o: tanh
nonlinearity_gates:
  e: silu
  o: tanh

resnet: false                                                                     
num_basis: 8                                                                  
BesselBasis_trainable: true                                                 
PolynomialCutoff_p: 6                                                           

# radial network
invariant_layers: 3                                                                
invariant_neurons: 64                                                        
avg_num_neighbors: 10                                                       
use_sc: true                                                                  
compile_model: false                                            

# data set
dataset: npz                                                                        
dataset_file_name: path/to/data/aspirin_ccsd-train.npz                         
key_mapping:
  z: atomic_numbers                                                  
  E: total_energy                                                       
  F: forces                                                                 
  R: pos                                                                       
npz_fixed_field_keys:                                                             
  - atomic_numbers

# logging
wandb: true                                                                       
wandb_project: ccsd-ccsdt                                                             
wandb_resume: true                                                       

verbose: info                                                                      
log_batch_freq: 1000000                                                                  
log_epoch_freq: 1                                                                  
save_checkpoint_freq: -1                                                           
save_ema_checkpoint_freq: -1                                                        

# training
n_train: 950                                                                      
n_val: 50                                                                      
learning_rate: 0.01                                                                
batch_size: 5                                                                   
max_epochs: 1000000                                                             
train_val_split: random                                                          
shuffle: true                                                                     
metrics_key: loss    
use_ema: true     
ema_decay: 0.99                                                               
ema_use_num_updates: true                                        

early_stopping_patiences:                                                     
  Validation_loss: 1000                                                            
early_stopping_lower_bounds:                                                     
  LR: 1.0e-6                                                                        

# loss function
loss_coeffs:                                                                      
  forces: 1000                                                                   
  total_energy: 1                                                                  

# output metrics
metrics_components:
  - - forces                                
    - rmse                                
    - PerSpecies: True                      
      report_per_component: False          
  - - forces
    - mae
    - PerSpecies: True
      report_per_component: False
  - - total_energy
    - mae
  - - total_energy 
    - rmse

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

lr_scheduler_name: ReduceLROnPlateau
lr_scheduler_patience: 50
lr_scheduler_factor: 0.8

PerSpeciesScaleShift_enable: true
PerSpeciesScaleShift_trainable: true
PerSpeciesScaleShift_shifts: [0.0, 0.0, 0.0]
PerSpeciesScaleShift_scales: [1.0, 1.0, 1.0]


global_rescale_shift: dataset_energy_mean
global_rescale_scale: dataset_force_rms

trainable_global_rescale_shift: false
trainable_global_rescale_scale: false

[lsj2408]

Sorry for the late reply.

The exact nequip commit is a5d160e. My python version is 3.8.12. The torch version is '1.9.0+cu111'. The config file is as following:

root: ccsd-ethanol-n950-l3-root
run_name: ethanol-ccsd-n950-l3-w-valid

requeue: true
seed: 0
append: true
default_dtype: float32
allow_tf32: false
device: cuda:3

# network
r_max: 4.0

num_layers: 5
chemical_embedding_irreps_out: 64x0e
feature_irreps_hidden: 64x0o + 64x0e + 64x1o + 64x1e + 64x2o + 64x2e + 64x3o + 64x3e
irreps_edge_sh: 0e + 1o + 2e + 3o
conv_to_output_hidden_irreps_out: 16x0e

nonlinearity_type: gate

nonlinearity_scalars:
  e: silu
  o: tanh
nonlinearity_gates:
  e: silu
  o: tanh

resnet: false
num_basis: 8
BesselBasis_trainable: true
PolynomialCutoff_p: 6

# radial network
invariant_layers: 3
invariant_neurons: 64
avg_num_neighbors: 8
use_sc: true
compile_model: false

# data set
dataset: npz
dataset_url: http://quantum-machine.org/gdml/data/npz/ethanol_ccsd_t.zip           # url to download the npz. optional
dataset_file_name: ./benchmark_data/ethanol_ccsd_t-train.npz
validation_dataset: npz
validation_dataset_url: http://quantum-machine.org/gdml/data/npz/ethanol_ccsd_t.zip           # url to download the npz. optional
validation_dataset_file_name: ./benchmark_data/ethanol_ccsd_t-test.npz                       # path to data set file
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
# logging
wandb: true
wandb_project: ccsd-ccsdt
wandb_resume: true

verbose: info
log_batch_freq: 1000000
log_epoch_freq: 1
save_checkpoint_freq: -1
save_ema_checkpoint_freq: -1

# training
n_train: 950
n_val: 1000
learning_rate: 0.01
batch_size: 5
max_epochs: 1000000
train_val_split: random
shuffle: true
metrics_key: validation_loss
use_ema: true
ema_decay: 0.99
ema_use_num_updates: true

early_stopping_patiences:
  validation_loss: 1000
early_stopping_lower_bounds:
  LR: 1.0e-6

# loss function
loss_coeffs:
  forces: 1000
  total_energy:
    - 1
    - PerAtomMSELoss

# output metrics
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
    - rmse
    - PerAtom: True

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

lr_scheduler_name: ReduceLROnPlateau
lr_scheduler_patience: 50
lr_scheduler_factor: 0.8

PerSpeciesScaleShift_enable: true
PerSpeciesScaleShift_trainable: true
PerSpeciesScaleShift_shifts: [0.0, 0.0, 0.0]
PerSpeciesScaleShift_scales: [1.0, 1.0, 1.0]


global_rescale_shift: dataset_total_energy_mean
global_rescale_scale: dataset_forces_rms

trainable_global_rescale_shift: false
trainable_global_rescale_scale: false

[lsj2408]

I will also try the v0.3.3 version and update the new results here.

[simonbatzner]

Ah I see what's wrong now, there are two things that have changed between the versions that are not set correctly in the config shared above:

a) the loss
b) the normalization of the targets.

I will explain both in detail below. This will be quite long, apologies for that, we are in the process of writing a detailed documentation and with that all this should be more clear I hope.

I have also run my config with v0.3.3 in the meantime and was able to perfectly reproduce the results. I then also ran one with v0.5 with the correct setup and also got identical results.

Please also note that we select 950 training points and 50 validation points from the ethanol_ccsd_t-train.npz file, and use the test.npz file only for testing, never for validation. It looks from your config like you use the test data for validation, that will of course give you slightly different results.

a) Loss

We have updated the default loss definition between v0.3.3 and v0.5. In the loss definition, we use a joint loss that sums the force-loss and the energy-loss. In v0.3 we simply use the total energy loss and the force loss, weight each and sum them. In v0.5, our default is to use the the PerAtomsMSELoss, which normalizes the energy loss to a per-atom quantity. This is done since forces are a local quantity while the energy is a global quantity. So what PerAtomsMSELoss does is it uses a local per-atom energy loss by diving the global energy loss by N^2, where N comes from the fact that there are N atoms in the system (converting global to local) and the square deals with the fact that we use a MSE loss.

In version 0.3.3, we still used the total energy (i.e. not divided by the number of atoms N) and therefore had different weighting between the two. So if you want to use the same config as I did in the paper you'd have to change this:

loss_coeffs:
  forces: 1000
  total_energy:
    - 1
    - PerAtomMSELoss

to this:

loss_coeffs:
  forces: 1000
  total_energy: 1

This will then not normalize the energy to a per-atom energy and should give you the same loss as what I used in the paper. What you did in your config is you combined PerAtomMSELoss with a high weight on the forces, that is equivalent to setting a force term of N^2 * 1000 = 81,000 in the case of ethanol (9 atoms). In general, this relative weighting is a hyperparameter that is a function of your system. We typically recommend what we have in our example.yaml which is:

loss_coeffs:                                                                   
  forces: 1                                                                       
  total_energy:                                                              
    - 1
    - PerAtomMSELoss

This will do a 1:1 weighting, but with normalized energies, so in the language of the non-normalized weighting it will be force-weight N^2, energy-weight 1 for the reasons outlined above.

b) Normalization of targets

The other thing that is set up differently is the normalization. The options there have changed and what you used above will be a no-option, i.e. it will not have an effect. I have included below what will give you identical behavior to v0.3.3 as well as what we recommend now.

To get back the v0.3.3 scaling, you should set

per_species_rescale_scales_trainable: true
per_species_rescale_shifts_trainable: true
per_species_rescale_shifts: [0.0, 0.0, 0.0]
per_species_rescale_scales: [1.0, 1.0, 1.0]
global_rescale_shift: dataset_total_energy_mean
global_rescale_scale: dataset_forces_rms

per_species_rescale_arguments_in_dataset_units: false

global_rescale_shift_trainable: false
global_rescale_scale_trainable: false

This will shift the energy targets by the mean energy over the data set and scale the force targets by the RMS of force components over the data set. It will also include per-species trainable shifts + scales, initialized to 0 and 1 as explained in the paper. I ran this myself in v0.5 and it perfectly reproduced the training runs from the paper.

Going forward, we recommend using the new v0.5 normalization:

per_species_rescale_scales_trainable: false
per_species_rescale_shifts_trainable: false
per_species_rescale_shifts: dataset_per_atom_total_energy_mean
per_species_rescale_scales: dataset_forces_rms

This will not use a global shift, but will instead initialize the per-species shift to the mean atomic energy, i.e. mean of total energy divided by number of atoms. This is useful for dataset with varying numbers of atoms or if you want to run MD on a system with more atoms than what you trained on. That's why we recommend this. For ethanol, where the number of atoms never changes, this doesn't matter. I also ran this on your system and it gives very similar results (not identical, I didn't set the scales/shifts to trainable here).

---

We are in the process of writing documentation, which will include this of course, I hope that will make it clearer in the future.

Let me know if this works now and if you have any other questions, about the loss function or otherwise.

[lsj2408]

Thanks for your detailed instructions! I've reproduced the results with the correct configuration.

[simonbatzner]

That's great to hear! Let us know if you have any other questions, we're also happy to help with selecting proper hyperparameters or anything, just reach out. Good luck, Simon