Multihead Replay Fine-tuning of. a foundation model

[spiccinin]

Hello,

I'm experimenting with fine-tuning a foundation model.
I read this section of the documentation:
https://mace-docs.readthedocs.io/en/latest/guide/finetuning.html
in particular the one on multihead fine-tuning, since it "is the recommended way to fine-tune any materials project foundation model". I downloaded the file "mace_agnesi_medium.model" and adapted the command provided in the documentation to my needs:

mace_run_train
--name="MACE"
--foundation_model="../mace_agnesi_medium.model"
--multiheads_finetuning=True
--train_file="../train.xyz"
--valid_fraction=0.05
--test_file="../test.xyz"
--E0s="{1: -12.935, 8: -435.157, 11: -1285.235, 26: -3357.673}"
--energy_key="REF_energy"
--forces_key="REF_forces"
--lr=0.01
--swa_lr=0.005
--max_num_epochs=30
--start_swa=20
--forces_weight=1000
--energy_weight=80
--swa_forces_weight=100
--swa_energy_weight=1000
--eval_interval=5
--swa
--scaling="rms_forces_scaling"
--batch_size=1
--ema
--ema_decay=0.99
--amsgrad
--default_dtype="float64"
--device=cuda
--seed=3

train.xyz and test.xyz contain a very small set of structures for my system of interest, where I use the keys REF_energy and REF_forces.

In the MACE_run-3.log log file I get the following warning:

WARNING: Using multiheads finetuning with a foundation model that is not a Materials Project model, need to provied a path to a pretraining file with --pt_train_file.

The training completes successfully, but I think here I'm doing what the documentation calls "naive fine-tuning". I'm not sure how to proceed. Where do I get the "pretraining file"? Is this the database that has been used to train the full MP_0b foundation model?

Thank you very much!
Simone Piccinin
CNR-IOM, Trieste

[ilyes319]

Hey @spiccinin,

I have changed the finetuning interface and have just updated the documentation (forgot to do it sorry). You need to update the input to add the following line: --pt_train_file="mp" \.

[spiccinin]

Thank you very much. Now it works fine.

[kush2803]

Can we get the --pt_train_file for the omat model as well? I am trying to fine tune the medium-omat-0 model, and need the training set to for it.

[kush2803]

@ilyes319 is this fixed in the latest release? Can really do with the original dataset for fine-tuning omat model.