Multihead Finetuning result is not satisfactory.

[jungsdao]

mace_run_train \
    --name="MACE" \
    --foundation_model="MACE-matpes-pbe-omat-ft.model" \
	--pt_train_file="selected_configs_key.xyz" \
    --weight_pt_head=0.1 \
    --multiheads_finetuning=True \
    --train_file="dft.xyz" \
    --valid_fraction=0.05 \
    --energy_weight=1.0 \
    --forces_weight=100.0 \
    --swa \
	--start_swa=15 \
    --swa_energy_weight=10.0 \
    --swa_forces_weight=100.0 \
    --E0s="{"1" : -13.12690491252935, "6" : -249.78441787631172, "8" : -571.5831119844246, "29": -2960.1
37964281592}" \
    --atomic_numbers "[1, 6, 8, 29]" \
    --batch_size=12 \
    --max_num_epochs=30 \
    --restart_latest \
	--energy_key="DFT_energy" \
	--forces_key="DFT_forces" \
    --default_dtype="float64" \
    --device=cuda \
	--enable_cueq=True \
    --save_cpu \
    --seed=3

I'm trying multihead finetuning with above keyword but training results looks a bit underwhelming and It looks like pt_head error is decreasing well but default_head error is not decreasing that much (especially energy error is quite stagnant). I wonder what's wrong with my fine-tuning. I remember in other cases, I used to get energy RMSD less than 10 meV/atom.

I have selected 30000 geometries from matpes pretrained data set.

I'll be appreciated to know possible way to get better fine-tuning behavior. Thank you in advance!

Log file is attached below.
stdout.txt

[vue1999]

Hi, could you try running with 10k replay, using these hypers?

mace_run_train \
    --name="MACE" \
    --foundation_model="MACE-matpes-pbe-omat-ft.model" \
	--pt_train_file="selected_configs_key.xyz" \
    --weight_pt_head=0.1 \
    --multiheads_finetuning=True \
    --train_file="dft.xyz" \
    --valid_fraction=0.05 \
    --energy_weight=10.0 \
    --forces_weight=10.0 \
    --E0s="{"1" : -13.12690491252935, "6" : -249.78441787631172, "8" : -571.5831119844246, "29": -2960.1
37964281592}" \
    --atomic_numbers "[1, 6, 8, 29]" \
    --batch_size=8 \
    --max_num_epochs=150 \
    --restart_latest \
	--energy_key="DFT_energy" \
	--forces_key="DFT_forces" \
    --default_dtype="float64" \
    --device=cuda \
	--enable_cueq=True \
    --save_cpu \
    --weight_decay=0 \
    --seed=3