Pseudopotential mismatch between pretrained data and fine-tuning data

[dominicvarghese]

Hi everyone,

I have a question about the best practices for multi-head fine-tuning. I am currently trying to fine-tune a foundation model on my own VASP dataset and am seeing good preliminary results, but I have a doubt about the underlying methodology.

Here is my setup:
Foundation Model: I am using MATPES-R2SCAN as my foundation model, which is based on the R2SCAN functional.
Replay Dataset: I am using the standard matpes-r2scan-replay-data.extxyz as my pt_train_file.
My Finetuning Dataset: I have a small, specific dataset for BaO that I generated myself using VASP with the PBEsol functional.

My main concern is the mismatch in the DFT functionals. The foundation model and its replay data are based on the R2SCAN potential energy surface (PES), while my new fine-tuning data is based on the PBEsol PES.

I have a few specific questions about this:
How does this mismatch in functionals affect the fine-tuning process? Will the model's shared weights get "confused" or learn an incorrect average of the two different potential energy surfaces?

I have calculated my own e0s.json file using PBEsol for my train_file. The replay data will use the original R2SCAN E0s. Is the multi-head model designed to handle two different sets of E0s (one for each head)?

For the best results, is it a firm requirement to generate the fine-tuning dataset using the exact same functional (R2SCAN) as the foundation model?

My validation errors on the fine-tuning head look very good, but I want to be sure I am following the procedure for these simulations.

From the current implementation my errors as follows:

+---------------+---------------------+------------------+-------------------+---------------------------------------+
|  config_type  | RMSE E / meV / atom | RMSE F / meV / A | relative F RMSE % | RMSE Stress (Virials) / meV / A (A^3) |
+---------------+---------------------+------------------+-------------------+---------------------------------------+
| train_Default |            2.5      |          3.2     |          0.94     |                     0.7               |
| valid_Default |            2.3      |          3.2     |          0.90     |                     0.7               |
+---------------+---------------------+------------------+-------------------+---------------------------------------+

Any advice or insights on this would be greatly appreciated!

Thanks
Dominic

[ilyes319]

Dear @dominicvarghese,

There is no problem with fine-tuning to a different level of theory. The change in level of theories is handled by two things:

• You specifying a new set of E0s, which now match the new level of theory. This is the main factor of inconsistency between levels of theories/code.
• The multihead model has a new set of readouts (so new weights) that are just for your finetuned set. This means that the model has the flexibility to output different energies/forces/stresses for different levels of theories and therefore it is not outputting an average. Note that the rest of the weights are shared, which makes sure that the model keeps the nice physicality of the foundation models

If you want, you can share with us your log file and we might be able to help you further.

[gabor1]

Just to add to Ilyes' comments. MACE fits the atomisation energy, this is why we need the E0s. This also allows us to cope with different pseudo potentials and functionals better, because the atomisation energy is not as different as the total energy when you change pseudopotentials or functionals.

your errors look good, especially on the forces. I think your energy errors could perhaps be lowered, but I suspect you may be limited there by your k-point sampling in your original DFT data! "standard" settings actually often lead to a few meV/atom error (which many DFT studies don't actually care about because they benefit from the cancellation of errors which occurs when you use the same unit cell, same k-grid when computing energy differences). So you could try to up your k-point density.

[dominicvarghese]

Thanks so much for the clear and helpful replies, @ilyes319 and @gabor1!

@ilyes319 I am attaching the log file here as requested: log.txt

@gabor1 In my command, I set forces_weight=100.0 while energy_weight=1.0, which is likely why the force errors are good. Regarding the k-points, thank you for that insight. I should clarify: my PBEsol data was generated for a 2x2x2 supercell (8 unit cells) using a 4x4x4 k-grid. This should be equivalent to a very well-converged 8x8x8 k-grid for the primitive cell. Given this, do you think improving k-grid would help with the energy error?

This is the full command I used for the fine-tuning.

python -m mace.cli.run_train \
  --name="my_BaO_finetuned_MATPES" \
  --train_file="../full_dataset.xyz" \
  --foundation_model="../MACE-matpes-r2scan-omat-ft.model" \
  --pt_train_file="../matpes-r2scan-replay-data.extxyz" \
  --multiheads_finetuning=True \
  --num_samples_pt=10000 \
  --filter_type_pt="combinations" \
  --subselect_pt="fps" \
  --atomic_numbers="[8, 56]" \
  --energy_weight=1.0 \
  --forces_weight=100.0 \
  --stress_weight=10.0 \
  --E0s="../e0s.json" \
  --energy_key="energy" \
  --forces_key="forces" \
  --stress_key="stress" \
  --max_num_epochs=30 \
  --lr=0.0001 \
  --config_type_weights '{"Default": 10.0}'\
  --weight_pt_head 1.0\
  --ema \
  --ema_decay=0.99999 \
  --device=cuda

[dominicvarghese]

We usually do two-stage training, so after some epochs with high force weight (but often 10:1 use enough!) We then continue with 1:1000 force: energy weight to bring down the energy error, and the force error doesn't go up very much.

There was nothing wrong with the stage 1 model. I was just trying to see whether the two-stage training would lower the energy errors further down.

[gabor1]

before doing any more training there, I would scrutinise the k-point sampling, actually measure the level of convergence in some difficult cases (your most skewed cells, or medium sized cells where the 1x1x1 or 2x2x2 points are the furthest away from the target density. )

[gabor1]

The other possibility that occurs to me, I wonder if Ilyes can corroborate this, is that when fine-tuning with a relatively small number of new configs, the number of energy values is really very small, so putting a very large weight on them can throw the parameter trajectory off. You could try something much milder, e.g. starting from the stage1 model and using weight ratios of 1:1:1, to see if the energy improves.

[ilyes319]

I think one problem is your swa_lr is probably too big. You start with a lr of 1e-4 in the normal phase but the default swa_lr is 1e-3. Try to reduce the swa_lr to 1e-5 or 1e-6.

[ilyes319]

Usually we use ratio more like 10:1:10 (forces/energy/stress).