The RMSE of pt_head is extremely anomalous, when using the Multihead Replay Finetuning method

[ContiMdc]

Dear Developer:
When utilizing the Multihead Replay Finetuning method, the RMSE for my pt_head is extremely anomalous. I generated the pt_train_file using "Method 1: Preprocessing via CLI" as described in the documentation. After applying this method, my dataset was replaced by the 10,000 data points contained within the pt_train_file; perceiving this behavior as abnormal, I decided to proceed with fine-tuning using my original dataset for the Default head (covering the testing and validation phases), while reserving the merged dataset specifically for training the pt_head. Provided below are the parameters I used to generate the pt_train_file and to fine-tune the model. Could you please tell me what might have gone wrong?

Merge Dataset:

r"python.exe",

"-m", "mace.cli.fine_tuning_select",

"--configs_pt", r"replay-data-mh-1-omat-pbe_no_ref.xyz",

"--configs_ft", r"npg_train.extxyz",

"--num_samples", "10000",
"--subselect", "fps",

"--model", r"mace-mh-1.model",

"--output", r"npg_train.extxyz",

"--filtering_type", "combinations",
"--device", "cpu",
"--default_dtype", "float64",
"--head_pt", "omat_pbe",
"--head_ft", "omat_pbe",
"--weight_pt", "1.0",
"--weight_ft", "10.0",

**Fine-tuning Parameters:**

    "python",
r"mace/cli/run_train.py",

"--name", "MACE",

"--foundation_model", "mace-mh-1.model",
"--foundation_head", "omat_pbe",
"--E0s", "foundation",


"--pt_train_file", "npg_train_combined.extxyz",
"--train_file", "npg_train.extxyz",
"--valid_fraction", "0.05",
"--test_file", "npg_test.extxyz",

"--multiheads_finetuning", "True",

"--atomic_numbers", "[8,6,1]",

"--energy_weight", "1.0",
"--forces_weight", "100.0",

"--scaling", "rms_forces_scaling",
"--batch_size", "2",
"--max_num_epochs", "300",

"--swa",
"--start_swa", "200",
"--swa_energy_weight", "10.0",
"--swa_forces_weight", "100.0",

"--ema",
"--ema_decay", "0.99",
"--amsgrad",

"--default_dtype", "float64",
"--device", "cuda",
"--seed", "123",

"--compute_stress", "True",
"--stress_key", "REF_stress"

**run.log:**

2026-03-31 10:38:53.831 INFO: ===========TRAINING===========

2026-03-31 10:38:53.831 INFO: Started training, reporting errors on validation set
2026-03-31 10:38:53.831 INFO: Loss metrics on validation set
2026-03-31 10:39:00.396 INFO: Initial: head: pt_head, loss=0.00561533, RMSE_E_per_atom= 6036.40 meV, RMSE_F=41702.56 meV / A, RMSE_stress= 13.47 meV / A^3
2026-03-31 10:39:01.569 INFO: Initial: head: Default, loss=0.00279222, RMSE_E_per_atom= 53.62 meV, RMSE_F= 68.58 meV / A, RMSE_stress= 3.60 meV / A^3
2026-03-31 10:50:14.871 INFO: Epoch 0: head: pt_head, loss=0.00467794, RMSE_E_per_atom= 6049.30 meV, RMSE_F=38833.37 meV / A, RMSE_stress= 12.96 meV / A^3
2026-03-31 10:50:16.037 INFO: Epoch 0: head: Default, loss=0.00159546, RMSE_E_per_atom= 6.39 meV, RMSE_F= 54.78 meV / A, RMSE_stress= 1.44 meV / A^3
2026-03-31 11:01:38.841 INFO: Epoch 1: head: pt_head, loss=0.00352917, RMSE_E_per_atom= 6041.99 meV, RMSE_F=22949.08 meV / A, RMSE_stress= 6.84 meV / A^3
2026-03-31 11:01:40.005 INFO: Epoch 1: head: Default, loss=0.00149091, RMSE_E_per_atom= 5.74 meV, RMSE_F= 51.71 meV / A, RMSE_stress= 1.11 meV / A^3
2026-03-31 11:12:53.559 INFO: Epoch 2: head: pt_head, loss=0.00219206, RMSE_E_per_atom= 6035.34 meV, RMSE_F= 9314.58 meV / A, RMSE_stress= 4.37 meV / A^3
2026-03-31 11:12:54.723 INFO: Epoch 2: head: Default, loss=0.00142840, RMSE_E_per_atom= 5.05 meV, RMSE_F= 50.05 meV / A, RMSE_stress= 0.98 meV / A^3
2026-03-31 11:24:30.529 INFO: Epoch 3: head: pt_head, loss=0.00312265, RMSE_E_per_atom= 6031.66 meV, RMSE_F=25744.99 meV / A, RMSE_stress= 8.17 meV / A^3
2026-03-31 11:24:31.704 INFO: Epoch 3: head: Default, loss=0.00139027, RMSE_E_per_atom= 4.78 meV, RMSE_F= 48.90 meV / A, RMSE_stress= 0.81 meV / A^3
2026-03-31 11:35:37.781 INFO: Epoch 4: head: pt_head, loss=0.00176727, RMSE_E_per_atom= 6030.35 meV, RMSE_F= 8393.25 meV / A, RMSE_stress= 3.13 meV / A^3
2026-03-31 11:35:38.941 INFO: Epoch 4: head: Default, loss=0.00133849, RMSE_E_per_atom= 4.59 meV, RMSE_F= 47.05 meV / A, RMSE_stress= 0.83 meV / A^3
2026-03-31 11:46:52.281 INFO: Epoch 5: head: pt_head, loss=0.00264193, RMSE_E_per_atom= 6028.58 meV, RMSE_F=18489.89 meV / A, RMSE_stress= 5.75 meV / A^3
2026-03-31 11:46:53.433 INFO: Epoch 5: head: Default, loss=0.00130285, RMSE_E_per_atom= 4.27 meV, RMSE_F= 45.94 meV / A, RMSE_stress= 0.78 meV / A^3
2026-03-31 11:57:33.715 INFO: Epoch 6: head: pt_head, loss=0.00094856, RMSE_E_per_atom= 6024.22 meV, RMSE_F= 756.01 meV / A, RMSE_stress= 0.83 meV / A^3
2026-03-31 11:57:34.874 INFO: Epoch 6: head: Default, loss=0.00129234, RMSE_E_per_atom= 4.23 meV, RMSE_F= 45.64 meV / A, RMSE_stress= 0.79 meV / A^3
2026-03-31 12:08:22.550 INFO: Epoch 7: head: pt_head, loss=0.00140388, RMSE_E_per_atom= 6020.65 meV, RMSE_F= 6807.95 meV / A, RMSE_stress= 2.02 meV / A^3
2026-03-31 12:08:23.720 INFO: Epoch 7: head: Default, loss=0.00127134, RMSE_E_per_atom= 4.13 meV, RMSE_F= 44.85 meV / A, RMSE_stress= 0.71 meV / A^3
2026-03-31 12:19:15.353 INFO: Epoch 8: head: pt_head, loss=0.00095438, RMSE_E_per_atom= 6017.21 meV, RMSE_F= 2449.15 meV / A, RMSE_stress= 1.02 meV / A^3
2026-03-31 12:19:16.522 INFO: Epoch 8: head: Default, loss=0.00124814, RMSE_E_per_atom= 3.98 meV, RMSE_F= 44.00 meV / A, RMSE_stress= 0.77 meV / A^3
2026-03-31 12:30:15.794 INFO: Epoch 9: head: pt_head, loss=0.00090762, RMSE_E_per_atom= 6014.96 meV, RMSE_F= 1399.91 meV / A, RMSE_stress= 0.81 meV / A^3
2026-03-31 12:30:16.960 INFO: Epoch 9: head: Default, loss=0.00122825, RMSE_E_per_atom= 3.81 meV, RMSE_F= 43.32 meV / A, RMSE_stress= 0.76 meV / A^3