Fine-tuning MH-1

[matr1x-1]

I am trying to fine-tune a MH-1 foundation model with the omat_pbe head on the GAP-Carbon-20 dataset. In order to extract the necessary E0 value, I started a fresh MACE model and used the value there. The problem is that when I run the fine-tuning using replay the inital errors are very high - far above the 0.4 eV/atom mentioned in the documentation as a warning signal that E0 might be wrong.
This is not the only case where this has happened. I have other datasets that I created myself with CP2K and where I calculated E0 directly using DFT. The problem was exactly the same.
This is the way I call the training:
 

srun --cpu-bin=socket ./mps-wrapper.sh python -m mace.cli.run_train   \
        --name="mymodel_finetuned"   \
        --train_file="carbon_MACE_READY.xyz" \
        --valid_fraction=0.05 \
        --foundation_model "mh-1" \
        --downloads_dir "../mace_downloads" \
        --foundation_head="omat_pbe"   \
        --pt_train_file="replay-small-2000_train.xyz" \
        --pt_valid_file="replay-small-2000_valid.xyz" \
        --weight_pt_head=5.0 \
        --multiheads_finetuning=True   \
        --device=cuda     \
        --atomic_numbers="[6]" \
        --energy_key="REF_energy" \
        --batch_size=1 \
        --E0s "{6: -6.554891353818097}" \
        --num_workers=1 \
        --default_dtype="float64" \
        --forces_key="REF_forces" \
        --lr 1e-4 \
        --stress_weight=0.0 \
        --loss="ef" \
        --compute_stress=False


Is there anything that I might have missed? Thank you!

[ilyes319]

Hey @matr1x-1, sorry for the delay. Would you be able to share the log file? Note that the easiest to check if it is an E0 problem, is to do the E0 re evaluation trick. Here is a python script you can use:

reestimate_e0s.py

  python reestimate_e0s.py \
    --model mh-1 \
    --head omat_pbe \
    --configs path_to_data.xyz \
    --energy-key energy

[matr1x-1]

@ilyes319 Very interesting! I ran your script on my dataset to extract the E0s and now it actually works. The errors are way lower than before. Thanks a lot!