Parameters for Multihead Replay Finetuning

[raulloiscuns]

Hi everyone,

I'm relatively new to MACE and I'm working on fine-tuning the mh-1 foundation model using the multi-head replay approach. My fine tuning dataset is around 500-750 configurations of ~400 atoms each.

I've put together the following workflow based on the documentation, but I'm not entirely sure if the parameters make sense for a dataset of this size.

Here is my current setup:

# Processing the dataset
python -m mace.cli.fine_tuning_select \
  --configs_pt datasets/replay_data_mh_1_omat_pbe.xyz \
  --configs_ft datasets/finetuning_dataset.xyz \
  --num_samples 10000 \
  --subselect fps \
  --model mh-1 \
  --output datasets/selected_configs.xyz \
  --filtering_type combinations \
  --head_pt omat_pbe \
  --head_ft Default

# Train the model
python -m mace.cli.run_train \
  --name modified_mh-1 \
  --foundation_model mh-1 \
  --foundation_head omat_pbe \
  --E0s estimated \
  --pt_train_file datasets/selected_configs.xyz \
  --train_file datasets/finetuning_dataset.xyz \
  --config_type_weights '{"Default":10.0}' \
  --valid_fraction 0.05 \
  --energy_key REF_energy \
  --forces_key REF_forces \
  --energy_weight 1 \
  --forces_weight 10 \
  --stress_weight 1 \
  --swa \
  --swa_energy_weight 10.0 \
  --swa_forces_weight 100.0 \
  --device cuda \
  --default_dtype float64 \
  --enable_cueq True \
  --seed 42

Currently, I am extracting 10,000 configurations for the pre-training set. Is this number of configurations safe for my system? Should I modify it without risking catastrophic forgetting?

Also, I want to make sure the model focuses heavily on my fine tuning data. To do this, I included --config_type_weights '{"Default":10.0}'in the training command. Does this parameter actually work the way I expect in a multi-head setup? I'm curious if this is the standard way to achieve this.

Finally, I'd appreciate any rule of thumb or tips on other parameters you usually tweak in these cases. Thanks in advance for the help!

[ilyes319]

hey @raulloiscuns that looks good,! if you want examples of finetuning, I uploaded all the finetuning examples of the MP-0 paper here https://drive.google.com/drive/u/1/folders/1gyjFBWjcoh-NJfoGNMnHWAq3jUqkxknO you can have a look at the logs to see what i used and how it goes. I used 100K configs, but i think 10K is the best to use.

[raulloiscuns]

hi @ilyes319, thanks a lot for the feedback and for sharing those folders. I’ve started looking through the logs and they are incredibly helpful.

Regarding the dataset selection, I noticed the command fine_tuning_select generates a selected_configs_combined.xyz file where you can set --weight_pt 1.0 and --weight_ft 10.0. However, since the run_train command asks for the files separately via --pt_train_file and --train_file, I was wondering if I should just ignore that combined file? I want to ensure the model really focuses on the new chemistry, so I'm curious if using --config_type_weights '{"Default":10.0}' in the training command is the standard way to achieve that prioritization in a multi-head setup.

I’m also curious about your experience with SWA. Do you find it particularly useful for smaller fine-tuning datasets like mine (~750 configs)?

Finally, since the documentation shows several different setups, if you could also share a "rule of thumb" for standard parameters like patience, learning rate, or the scheduler in these cases, that would be ideal. Having a solid starting point based would be a huge help to get things running properly.

Thanks again for the help!

[lazerpo]

@raulloiscuns @ilyes319 hi, I am joining the discussion because I have also noticed this and I have a similar question. Since in run_train there is no --weight_ft option like in the fine_tuning_select script, what's the correct way to prioritize the ft dataset? Can we just set --weight_pt < 1? Or should we use --config_type_weights?

On a different note, I also have questions about SWA when doing multihead fine-tuning. Since very early stopping is recommended, does stage two still make sense?

Thank you.