Reducing the RMSE for my MACE models

[vijaymocherla]

Hi,

I am trying to train 3 different MACE models (16e0,16e0+16o1,32e0) for different functionals on relatively small datasets (~900-1000 configrations). For each functional, my data is split into multiple heads contributing about 250-350 configurations and there are about 37 atoms per configurations in each of the .xyz files.

For example, the config file for one of my models looks as follows:

# Model params
model: "ScaleShiftMACE"
num_channels: 16
max_L: 0
hidden_irreps: '16x0e'
r_max: 5.0
correlation: 2 
num_interactions: 2
forces_weight: 10
energy_weight: 1

# traning params
device: cuda
default_dtype: float64
save_cpu: yes
batch_size: 16
amsgrad: yes
max_num_epochs: 1200
swa: True
seed: 314159265
swa: yes
start_swa: 900
# swa_energy_weight: 1000
# swa_forces_weight: 1
ema: yes
ema_decay: 0.99
error_table: PerAtomRMSE

# data params
energy_key: "energy"
forces_key: "forces"
E0s: "average"
name: "pbe_16e0"
model_dir: "pbe_16e0"
log_dir: "pbe_16e0"
checkpoints_dir: "pbe_16e0"
results_dir: "pbe_16e0"

# model descriptors
heads:
  DFT:
    train_file: ["./data/train_gs_pbe_3x3x4.xyz", "./data/train_rnd_ads_pbe_3x3x4.xyz", "./data/train_thermal_slab_pbe_3x3x4.xyz"]
    valid_fraction: 0.10
    test_file: ["./data/test_gs_pbe_3x3x4.xyz", "./data/train_rnd_ads_pbe_3x3x4.xyz", ./data/test_thermal_slab_pbe_3x3x4.xyz]

I've tried to optimize the training of my models with different batchsizes (2,4,8,16,32). However, I am not able to improve the performance of models beyond ~1.5 RMSE E / meV/atom. Given the performance of MACE models for similar systems in arXiv:2403.15334, I was expecting the error in energy to be sub meV/atom (<30 meV for a given configuration).

Questions

1. Are the traning datasets too small to achieve lower RMSE for energies and forces?
2. Can I improve my traning with different configuration weights for the data from multiple heads?
3. Do I have to use larger batchsizes to improve my performance?

Here's are the final error-tables for the above model:
+-------------+---------------------+------------------+-------------------+
| config_type | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % |
+-------------+---------------------+------------------+-------------------+
| train_DFT | 1.2 | 22.4 | 0.89 |
| valid_DFT | 1.2 | 21.5 | 0.94 |
+-------------+---------------------+------------------+-------------------+

+-------------+---------------------+------------------+-------------------+
| config_type | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % |
+-------------+---------------------+------------------+-------------------+
| Default_DFT | 1.5 | 28.6 | 1.00 |
+-------------+---------------------+------------------+-------------------+

A plot from training the above config file looks like this:

[ilyes319]

Hey @vijaymocherla,

• First as a general statement, whether a model is accurate enough, should always be based on downstream observables. Raw mae can be an indication, but what is "good" can be widely dependent on the kind of system you study and what observable you want to compute. Moreover depending on the convergence of your DFT, you might reach the actual noise of your DFT. In general, i would say that around 1 meV/atom is good for most applications.

• The models you are using are extremely small. if you look at our papers; we usually use as our smallest models "128x0e" which we call "small". Our medium models are "128x0e + 128x1o". Moreover, i never encourage to use anything else than correlation: 3. I would encourage you to use a more appropriate model size. It should provide a large accuracy improvement. In case you need more speed, i recommend you to use cuequivariance, see: https://mace-docs.readthedocs.io/en/latest/guide/cuda_acceleration.html

[vijaymocherla]

Hey @ilyes319,

Thanks for the quick reply!

I started with the above models based on the information from the paper that benchmarked various other models including MACE. Here's a snippet from arXiv:2403.15334:

Cutoff distance of 5 ̊A, correlation order of 2 and 2 interactions were used in the final MACE models. Initial loss function energy and force weights of 5 and 95, respectively, were used in the training process. After 1200 epochs the weights were changed to 1000 and 1 for energies and forces, respectively. Model size (hidden_irreps) in the final model was set to 16x0e, so equivariant message passing was found to be not necessary to obtain high accuracy results in the
present case

In any case, would you suggest using larger models (say 128x0e) along with different different weights for different heads?

[ilyes319]

that paper used these size because it was trained on CPU. I recommend you use "128x0e + 128x1o".

[vijaymocherla]

Ok, I didn't see any improvement or differences across 16x0e, 32x0e and 16x0e+16x1o. But, then I guess it just that my models are too small. Again, thanks a lot for the help!

[ilyes319]

that is system dependent. If you dont see an improvement, you have probably hit the DFT accuracy already.

[remorses]

[vijaymocherla]

Are you an agent spamming this thread!?

[remorses]

Sorry I used the gh cli on a completely wrong thread by mistake