Multihead replay converge slowly

[Ssbban]

Hi,

(Highly non-standard)
I want to share that I encountered a similar problem as discussed previously #626 (comment). When I enabled multiheads replay, the RMSE on energy and force converged much slower than training from-scratch method. I tried to check the initial loss and found that the initial loss of the defualt head of multiheads replay method, which was very large.

Here is the training script of multihead replay,

mace_run_train \
    --name="MACE_model" \
    --enable_cueq=True \
    --train_file="train.xyz" \
    --valid_file="valid.xyz" \
    --foundation_model="../../../mace_agnesi_medium.model" \
    --multiheads_finetuning=True \
    --pt_train_file="mp" \
    --energy_weight="10" \
    --forces_weight="1" \
    --stress_weight="0" \
    --E0s="{"1":-12.412976807294708, "8":-2041.8522325643492}" \
    --energy_key='REF_energy_delta' \
    --forces_key='REF_forces_delta' \
    --lr=0.001 \
    --scaling="rms_forces_scaling" \
    --batch_size=1 \
    --valid_batch_size=1 \
    --max_num_epochs=300 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --device=cuda \
    --default_dtype="float64" \
    --seed=3

Here is the script of scratch,

mace_run_train \
    --name="MACE_scratch" \
    --enable_cueq=True \
    --train_file="train.xyz" \
    --valid_file="valid.xyz" \
    --energy_weight="10" \
    --forces_weight="1" \
    --stress_weight=0 \
    --E0s="{"1":-12.412976807294708, "8":-2041.8522325643492}" \
    --energy_key='REF_energy_delta' \
    --forces_key='REF_forces_delta' \
    --lr=0.001 \
    --scaling="rms_forces_scaling" \
    --r_max=6.0 \
    --distance_transform="Agnesi" \
    --batch_size=1 \
    --valid_batch_size=1 \
    --max_num_epochs=300 \
    --ema \
    --ema_decay=0.99 \
    --amsgrad \
    --default_dtype="float64" \
    --device=cuda \
    --seed=3

while the result are,
multihead.txt
scratch.txt

Thank you so much for your help and patience, any suggestions would be greatly appreciated.

[ilyes319]

Hey Can you upload the full log file for both so I can have all the infos.

[ilyes319]

mmmmm that is highly non standard, fitting energy differences and fitting total energy are not the same task, so I would advice you do that carefully with multihead finetuning.

[ilyes319]

Can you tell me precisely what difference you are computing for both E0s and energy/forces?

[Ssbban]

Hi, here is the script I calculated the E0s,

structure_file = "cluster_4.5A.extxyz"
atoms_list = []

calc1 = mace_off('medium', default_dtype='float64')
calc2 = mace_mp('medium', default_dtype='float64')

for atoms in read(structure_file, ':'):
	atoms.info['REF_energy_delta'] = calc1.get_property('energy', atoms) - calc2.get_property('energy', atoms)
	atoms.arrays['REF_forces_delta'] = calc1.get_property('forces', atoms) - calc2.get_property('forces', atoms)
	atoms_list.append(atoms)

file_name = Path(structure_file).stem
write('./' + file_name + '.xyz', atoms_list)

# Calculate E0
from ase import Atoms

atoms_h = Atoms('H')
atoms_h.info['config_type'] = 'IsolatedAtom'
atoms_o = Atoms('O')
atoms_o.info['config_type'] = 'IsolatedAtom'
e0_h = calc1.get_property('energy', atoms_h) - calc2.get_property('energy', atoms_h)
e0_o = calc1.get_property('energy', atoms_o) - calc2.get_property('energy', atoms_o)

print(f"E0 (H): {e0_h} eV")
print(f"E0 (O): {e0_o} eV")

[ilyes319]

with what version of the code are you doing that?

• I am very skeptical about using multihead finetuning for that use case, please when creating discussions mention if you are doing highly non standard things.
• You should be using the agnesi model for computing the E0s difference, not the "medium" model that does not have good referencing. Switch "medium" to "medium-0b".

[Ssbban]

For these results, I use MACE version 0.3.10. I will try this setting, thanks for your help.