Stress anti-corelated in the Default and pt_head in Multihead replay fine-tuning

[dominicvarghese]

Hi everyone,

I was doing the Multi head Replay fine tuning for a custom training set.

The pt_head and the Default head shows opposite relations for the stress.

This was the training script used :

python -m mace.cli.run_train \
  --name="STO_Jay-run" \
  --train_file="STO_AIMD_GS_training_data_MLABN.xyz" \
  --foundation_model="../../../MACE-matpes-r2scan-omat-ft.model" \
  --pt_train_file="selected_250_replay.xyz" \
  --multiheads_finetuning=True \
  --atomic_numbers="[38,22,8]" \
  --compute_forces=True \
  --compute_stress=True \
  --E0s="e0s.json" \
  --energy_key="REF_energy" \
  --forces_key="REF_forces" \
  --stress_key="REF_stress" \
  --loss="universal" \
  --energy_weight=10.0 \
  --forces_weight=100.0 \
  --stress_weight=10.0 \
  --max_num_epochs=50 \
  --lr=0.0001 \
  --config_type_weights '{"Default": 10.0}' \
  --weight_pt_head 1.0 \
  --ema \
  --ema_decay=0.99 \
  --device=cuda \
  --valid_fraction=0.05 \
  --amsgrad \
  --default_dtype="float64" \
  --seed=56 \
  --log_level="INFO" \
  --error_table=PerAtomRMSE

The parity plots obtained are attached below:

What does the anti-correlations in stress for the different heads mean? Is the sign convention for the stress from the training set opposite from the pre-training set?

Thanks
Dominic

[gabor1]

this is almost certainly bad data, with the -ve sign of the stress definition going wrong.

[dominicvarghese]

I used the foundational model - in this case MACE-matpes-r2scan-omat-ft.model to predict the stress for random configurations in my training set and those came out pretty close to what I got from DFT (same sign, with the ratio ~1.2 probably because of the difference in the pseudopotential). And when I used the same foundational model to predict the stress from my replay set, those ratio were ~ -1, which suggests the sign flip was in the replay set.

But this replay set was obtained from mace.cli.run_train from the original dataset used to train the model matpes-r2scan-replay-data.extxyz.

I was comparing the dft_stress key mentioned in the model to compare the stress ration and that is -1.
While comparing the mace_mp_stress key gave exactly 1 - So I should be using that as the stress in the replay set to train the model.

Is my understanding correct?

[gabor1]

@ilyes319 can you comment on the sign of the dft_stress label in the matpes-r2scan-replay-data.extxyz ? I think the problem is that there is a tension between keeping the original training set (as produced by others) with its sign convention and flipping the sign when training, with the alternative which is reprocessing the original training set for mace and so then having multiple inconsistent training files floating around. so my suggestion is just follow the numerics: if the replay set you are using shows -1, then flip its sign.

[ilyes319]

I do remember that sign in the matpes training was inverted and needs a sign flip, though I remember re-uploading with the correct sign. We should for sure upload the corrected version to not mislead.

As a side note, in general, I recommend to use pseudolabels when doing multihead finetuning as it works as well and removes these problems

[gabor1]

using the pseudo-label sounds indeed like the clean solution as you originally suggested.

[ilyes319]

@dominicvarghese could you confirm that you use the latest file from https://github.com/ACEsuit/mace-foundations/releases/tag/mace_matpes_0? and if not to check if the two files still agree.

[dominicvarghese]

Yes, I verified the the above link, downloaded the model and the extxyz files and checked again. The mace_mp_stress key is the one that correlates with the model and the dft_stress key is anti-corelated.

[dominicvarghese]

Doing the training with the mace_mp_stress key gives proper parity plots for both Default and pre-train heads.