Help with MACE fine-tuning. Forces fits nicely, but energies don't

[PakkaDP]

Hello, I am new to MACE and have been trying to fine-tune the MACE-matpes-pbe-omat model with multi-head fine-tuning to my dataset obtained from AIMD using DFT with PBE functionals. I have tried fine-tuning the model with 100 and 200 configs and it achieved relatively good performance on the forces, but severely bad errors on the energies. I am unsure of what to do, should I increase the number of epochs or could this issue be solved just by increasing the amount of data? Right now, I set the max_num_epochs to 30 because the documentation mentioned that it should converge after about 10-30 epochs. The energy errors interestingly didn't seem to improve by much even after doubling the amount of data, so I was wondering if there was something fundamentally wrong with how I'm doing training. I've also tried recalculating the E0s, first with estimated E0s and second by recalculating them on VASP. The initial errors improved with the VASP recalculated E0s, but otherwise the final predicted energies are still just as bad.

Here is my yaml input file:
foundation_model: "path/to/foundation/model/MACE-matpes-pbe-omat-ft.model" \ multiheads_finetuning: True \ name: "pbe_finetuned" \ model_dir: "MACE_models/MACE-matpes-pbe-omat/pbe_finetuned" \ log_dir: "MACE_models/MACE-matpes-pbe-omat/pbe_finetuned" \ checkpoints_dir: "MACE_models/MACE-matpes-pbe-omat/pbe_finetuned" \ results_dir: "MACE_models/MACE-matpes-pbe-omat/pbe_finetuned" \ num_workers: 4 \ train_file: "system_train.xyz" \ valid_fraction: 0.10 \ test_file: "system_test.xyz" \ energy_key: "energy_aimd" \ forces_key: "forces_aimd" \ device: "cuda" \ atomic_numbers: "[...]" \ subselect_pt: "fps" \ batch_size: 4 \ max_num_epochs: 30 \ pt_train_file: "matpes_pbe" \ filter_type_pt: "combinations" \ num_samples_pt: 10000 \ default_dtype: "float64" \ swa: True \ enable_cueq: True \ seed: 123

Here's the results for 100 configs (RMSE E/atom = 2553 eV/atom):

Here's for 200 stage one weights (RMSE E/atom = 2406 eV/atom):

Here's for 200 stage two weights (it's as if there was a trade-off between forces and energies):

[vue1999]

Hey, could you please share the log files for the training run(s)?
From the input file and the plots, a couple thoughts:

• First, the test set seems to have a significantly different energy per atom distribution compared to the train/valid, are the E0s same for the test set? (if not, you need to adjust the E0s of the model to the test set before evaluating energies)
• In the test set there is an outlier with >60eV/atom energy, is that correct?
• While sometimes useful for from-scratch model training, I wouldn't use SWA for finetuning foundation models. Using a single stage training, with balanced energy and force weights in the loss function (or 1:10 energy:force, matching the loss used in the foundation model training) yields better results.

[PakkaDP]

Here is my log file for the 200 configs run:
pbe_finetuned_run-123.log

Responding to your thoughts:

• The E0s should be the same since I used data from the same AIMD run to train, validate, and test the model.
• I checked the data again, and there were a few outliers, I will try cleaning the data and training the model again.
• I see, I tried following a paper that used SWA (the weights changed in the final 20% of training) since they yielded good results, but they did train for longer epochs than I did, and I haven't considered the fact that they might've used from-scratch training.

Some questions I'd like to ask.

• Regarding your first point, are you suggesting I should include E0s data within the test set? I didn't do that since the T01 tutorial only included the E0s in the train set and not the test set, but that would explain the very strange energy errors that won't seem to budge.
• From your pointers, are you suspecting the data quality to be the source of the issue? How do you suggest I evaluate the quality of the data, because I am using already available data.
• Using the data currently available to me, what steps would you suggest I do?

Thank you very much for your pointers @vue1999.

[vue1999]

Thanks for the log file! Looking at this in detail, there are several problems with the current training setup:

1. Initial checkpoint is not the foundation model

Your log shows 'Using foundation model pbe_100_multi_stagetwo.model as initial checkpoint'. This is probably the output of a previous 100-config run, not the original foundation model. You should always fine-tune from the original foundation model, not from a previously fine-tuned one.

2. pt_head E0s are mostly zero

Related to the first point, the pt_head E0s in your log are {1: -0.064, ..., 6: -0.157, 8: -0.160, 28: -0.225} with all other elements being 0.0. These come from the model you're loading. The initial pt_head RMSE of 2130 meV/atom also show that the pt_head's E0s are incorrect. This means the model spends its capacity trying to compensate for fundamentally wrong energy references on the replay data, which corrupts the fine-tuning head as well (the two heads share model parameters apart from the readouts). This is likely caused by using a previously fine-tuned model as foundation.

3. Wrong replay data

Based on the logfile you are using the MPtraj replay data (mp_traj_combined.xyz) instead of the MatPES PBE replay data. For multihead finetuning the matpes model, either enable pseudolabeling, or you can manually download the correct replay data from https://github.com/ACEsuit/mace-foundations/releases/download/mace_matpes_0/matpes-pbe-replay-data.xyz and pass it as pt_train_file: "path/to/matpes-pbe-replay-data.xyz".

4. Loss weights

The default energy:force weights are 1:100, as shown in your log (UniversalLoss(energy_weight=1.000, forces_weight=100.000, ...)). This heavily prioritizes forces over energies. Try using 1:1 or at most 1:10 (energy:force).

To answer your specific questions:
Including E0s in the test set: No, you don't need to include E0 configurations in the test set. My previous comment was about a scenario where train and test data use different energy references (e.g., different DFT codes or settings). If your train, valid, and test data all come from the same AIMD run with the same settings, the E0s are consistent and this is fine.
Data quality: The data quality is not the primary issue. The only data concern is the +60 eV/atom outlier in the test set, which you should check (and remove).
There also seems to be a distribution shift (test energies range from about -5 to 0 eV/atom while training seems more tightly centred around -5 eV/atom), but that's secondary and difficult to confirm exactly from the plot because of the overlap.

My recommendation:

First, try simple naive fine-tuning:

foundation_model: "mace-matpes-pbe-0"
multiheads_finetuning: False
foundation_model_elements: True  # to keep all elements
energy_weight: 1
forces_weight: 1
swa: False
lr: 0.001
ema: True
ema_decay: 0.999
clip_grad: 1.0
weight_decay: 0
batch_size: 4
max_num_epochs: 200  # start with around 200 epochs (* 180 configs / 4 batch_size ≈ 9000 gradient updates)
E0s: if you have proper isolated atom configurations (single atoms in vacuum), use those. Otherwise use E0s: "estimated". The important thing is that isolated atoms are truly isolated, not extracted from multi-atom structures via averaging.

Once naive fine-tuning works well, try multihead

For multihead fine-tuning, the important things are:

Use the original foundation model (not a previously fine-tuned one)
Manually download the correct MatPES PBE replay data (or use pseudolabeling, in practice there's very little difference)
Keep energy:force weights at 1:1 (or 1:10)
Keep swa: False
Reduce epochs and learning rate compared to naive FT (e.g., lr: 0.0001, ema_decay: 0.9999, max_num_epochs: 30)
Set stress_weight: 0 if your fine-tuning data has no stresses
Also, you don't need to manually specify atomic_numbers - the code determines this automatically from the data. If you want to preserve all foundation model elements during naive fine-tuning, use foundation_model_elements: True instead.

[PakkaDP]

Thank you, I will try your suggestions.

Do you have any rules of thumb or references for further improving the model's training?

[PakkaDP]

The reason turned out to be the data. For some reason, my VASP output had different ionic step energies from the energies obtained in each final electronic step. After parsing the energies from the final electronic steps, the model fitted well. Thank you.

[gabor1]

It's always the data... :)