Multihead finetuning the foundation model with with training set calculated with different functional (RPBE)

[riyazpu]

I’m observing that during training, there are two heads being trained in parallel: default and pt_head. While the default head seems to be converging well (low RMSE and loss), the pt_head is diverging, with significantly higher errors that keep increasing slightly with each epoch.

Here are some sample logs:
"2025-04-30 07:32:38.093 INFO: Epoch 339: head: default, loss=0.00485631, RMSE_E_per_atom= 6.95 meV, RMSE_F= 61.55 meV / A, RMSE_stress= 3.14 meV / A^3
2025-04-30 07:36:05.052 INFO: Epoch 340: head: pt_head, loss=0.01610629, RMSE_E_per_atom= 88.83 meV, RMSE_F= 137.34 meV / A, RMSE_stress= 15.39 meV / A^3
2025-04-30 07:36:10.008 INFO: Epoch 340: head: default, loss=0.00485394, RMSE_E_per_atom= 6.95 meV, RMSE_F= 61.55 meV / A, RMSE_stress= 3.14 meV / A^3
2025-04-30 07:39:37.136 INFO: Epoch 341: head: pt_head, loss=0.01611283, RMSE_E_per_atom= 88.84 meV, RMSE_F= 137.40 meV / A, RMSE_stress= 15.39 meV / A^3
2025-04-30 07:39:42.099 INFO: Epoch 341: head: default, loss=0.00485162, RMSE_E_per_atom= 6.94 meV, RMSE_F= 61.54 meV / A, RMSE_stress= 3.14 meV / A^3
2025-04-30 07:43:09.225 INFO: Epoch 342: head: pt_head, loss=0.01612016, RMSE_E_per_atom= 88.84 meV, RMSE_F= 137.45 meV / A, RMSE_stress= 15.40 meV / A^3
2025-04-30 07:43:14.197 INFO: Epoch 342: head: default, loss=0.00484925, RMSE_E_per_atom= 6.94 meV, RMSE_F= 61.53 meV / A, RMSE_stress= 3.14 meV / A^3
2025-04-30 07:46:40.817 INFO: Epoch 343: head: pt_head, loss=0.01612688, RMSE_E_per_atom= 88.84 meV, RMSE_F= 137.49 meV / A, RMSE_stress= 15.41 meV / A^3
"
My questions:

Is it important for both heads (default and pt_head) to converge during training?

Since the final model output seems to be based only on the default head, can the pt_head be safely ignored if it diverges?

What might cause high RMSEs for pt_head?

Any insights or recommendations would be appreciated!

[ilyes319]

Hello @riyazpu,
The default head corresponds to the performance on your actual finetuning dataset.
The pt_head tracks the performance on the replay head, that prevents too much forgetting of the original model. It is normal that the pt_head gets worse, as some amount of forgetting is unavoidable. So the behavior you are observing is normal.