Issue with multihead replay fine-tuning MACE-OFF23

[lucasgarayy]

Hello everyone,

I am attempting to fine-tune MACEOFF-23 on 4 different heads corresponding to various DFT theories using multihead replay fine-tuning. However, I'm encountering issues where the pt_head does not train properly. The initial loss remains at 0.00000, and the RMSE_E and RMSE_F values are significantly large throughout training.

Here is the configuration I used:
model: 'MACE'
foundation_model: 'mace/foundation_models/MACE-OFF23_medium.model'
multiheads_finetuning: True
name: 'multi_head_v6_custom'
model_dir: 'mace/models'
log_dir: 'mace/logs'
results_dir: 'mace/models'
device: cuda
heads:
  linear_response_gs:
    train_file: 'datasets/unrattled_datasets/processed_datasets/train_lr-gs.xyz'
    valid_file: 'datasets/unrattled_datasets/processed_datasets/validation_lr-gs.xyz'
    energy_key: 'REF_energy'
    forces_key: 'REF_forces'
    E0s: 'isolated'
  linear_response_es1:
    train_file: 'datasets/unrattled_datasets/processed_datasets/train_lr-es1.xyz'
    valid_file: 'datasets/unrattled_datasets/processed_datasets/validation_lr-es1.xyz'
    energy_key: 'REF_energy'
    forces_key: 'REF_forces'
    E0s: 'isolated'
  spin_flip_es1:
    train_file: 'datasets/unrattled_datasets/processed_datasets/train_sf-es1.xyz'
    valid_file: 'datasets/unrattled_datasets/processed_datasets/validation_sf-es1.xyz'
    energy_key: 'REF_energy'
    forces_key: 'REF_forces'
    E0s: 'isolated'
  spin_flip_es2:
    train_file: 'datasets/unrattled_datasets/processed_datasets/train_sf-es2.xyz'
    valid_file: 'datasets/unrattled_datasets/processed_datasets/validation_sf-es2.xyz'
    energy_key: 'REF_energy'
    forces_key: 'REF_forces'
    E0s: 'isolated'
pt_train_file: 'datasets/mace_datasets/test_large_neut_all.xyz'
num_samples_pt: 10000
max_num_epochs: 2
batch_size: 2
seed: 42

Where the pt_train_file is taken from MACEOFF23 Training Data.

Log Output:

2025-01-23 15:20:39.123 INFO: ===========VERIFYING SETTINGS===========
2025-01-23 15:20:39.125 INFO: MACE version: 0.3.10
2025-01-23 15:20:39.128 INFO: CUDA version: 12.1, CUDA device: 0
2025-01-23 15:20:39.374 INFO: Using foundation model mace/foundation_models/MACE-OFF23_medium.model as initial checkpoint.
2025-01-23 15:20:39.375 INFO: Multihead finetuning mode, setting learning rate to 0.001 and EMA to True. To use a different learning rate, set --force_mh_ft_lr=True.
2025-01-23 15:20:39.375 INFO: Using multiheads finetuning mode, setting learning rate to 0.001 and EMA to True
2025-01-23 15:20:39.376 INFO: ===========LOADING INPUT DATA===========
2025-01-23 15:20:39.376 INFO: Using heads: ['linear_response_gs', 'linear_response_es1', 'spin_flip_es1', 'spin_flip_es2']
2025-01-23 15:20:39.377 INFO: =============    Processing head linear_response_gs     ===========
2025-01-23 15:20:39.495 INFO: Using isolated atom energies from training file
2025-01-23 15:20:39.520 INFO: Training set [339 configs, 339 energy, 14745 forces] loaded from 'datasets/unrattled_datasets/processed_datasets/train_lr-gs.xyz'
2025-01-23 15:20:39.557 INFO: Validation set [89 configs, 89 energy, 3915 forces] loaded from 'datasets/unrattled_datasets/processed_datasets/validation_lr-gs.xyz'
2025-01-23 15:20:39.557 INFO: Total number of configurations: train=339, valid=89, tests=[],
2025-01-23 15:20:39.558 INFO: =============    Processing head linear_response_es1     ===========
2025-01-23 15:20:39.635 INFO: Using isolated atom energies from training file
2025-01-23 15:20:39.646 INFO: Training set [339 configs, 339 energy, 14709 forces] loaded from 'datasets/unrattled_datasets/processed_datasets/train_lr-es1.xyz'
2025-01-23 15:20:39.730 INFO: Validation set [89 configs, 89 energy, 3951 forces] loaded from 'datasets/unrattled_datasets/processed_datasets/validation_lr-es1.xyz'
2025-01-23 15:20:39.732 INFO: Total number of configurations: train=339, valid=89, tests=[],
2025-01-23 15:20:39.735 INFO: =============    Processing head spin_flip_es1     ===========
2025-01-23 15:20:39.914 INFO: Using isolated atom energies from training file
2025-01-23 15:20:39.938 INFO: Training set [339 configs, 339 energy, 14853 forces] loaded from 'datasets/unrattled_datasets/processed_datasets/train_sf-es1.xyz'
2025-01-23 15:20:40.014 INFO: Validation set [89 configs, 89 energy, 3807 forces] loaded from 'datasets/unrattled_datasets/processed_datasets/validation_sf-es1.xyz'
2025-01-23 15:20:40.024 INFO: Total number of configurations: train=339, valid=89, tests=[],
2025-01-23 15:20:40.027 INFO: =============    Processing head spin_flip_es2     ===========
2025-01-23 15:20:40.255 INFO: Using isolated atom energies from training file
2025-01-23 15:20:40.275 INFO: Training set [339 configs, 339 energy, 14769 forces] loaded from 'datasets/unrattled_datasets/processed_datasets/train_sf-es2.xyz'
2025-01-23 15:20:40.308 INFO: Validation set [89 configs, 89 energy, 3891 forces] loaded from 'datasets/unrattled_datasets/processed_datasets/validation_sf-es2.xyz'
2025-01-23 15:20:40.308 INFO: Total number of configurations: train=339, valid=89, tests=[],
2025-01-23 15:20:40.309 INFO: ==================Using multiheads finetuning mode==================
2025-01-23 15:20:40.309 INFO: Using foundation model for multiheads finetuning with datasets/mace_datasets/test_large_neut_all.xyz
2025-01-23 15:21:06.931 INFO: Training set [50195 configs, 0 energy, 4655418 forces] loaded from 'datasets/mace_datasets/test_large_neut_all.xyz'
2025-01-23 15:21:06.943 INFO: Using random 10% of training set for validation with indices saved in: ./valid_indices_42.txt
2025-01-23 15:21:06.957 INFO: Validaton set contains 5019 configurations [0 energy, 467151 forces]
2025-01-23 15:21:06.959 WARNING: Ratio of the number of configurations in the training set and the in the pt_train_file is 0.030015937666017355, increasing the number of configurations in the pt_train_file by a factor of 3
2025-01-23 15:21:06.959 INFO: Total number of configurations in pretraining: train=45176, valid=5019
2025-01-23 15:21:07.191 INFO: Atomic Numbers used: [1, 6, 7, 8, 9, 15, 16, 17, 35, 53]
2025-01-23 15:21:07.193 INFO: Atomic Energies used (z: eV) for head linear_response_gs: {1: -13.568778219076226, 6: -1028.9771505969138, 7: -1483.5761877667198, 8: -2041.370068087562}
2025-01-23 15:21:07.193 INFO: Atomic Energies used (z: eV) for head linear_response_es1: {1: -13.568778219076226, 6: -1028.9771505969138, 7: -1483.5761877667198, 8: -2041.370068087562}
2025-01-23 15:21:07.194 INFO: Atomic Energies used (z: eV) for head spin_flip_es1: {1: -13.568778219076226, 6: -1028.9771505969138, 7: -1483.5761877667198, 8: -2041.370068087562}
2025-01-23 15:21:07.195 INFO: Atomic Energies used (z: eV) for head spin_flip_es2: {1: -13.568778219076226, 6: -1028.9771505969138, 7: -1483.5761877667198, 8: -2041.370068087562}
2025-01-23 15:21:07.195 INFO: Atomic Energies used (z: eV) for head pt_head: {1: -13.571964772646918, 6: -1030.5671648271828, 7: -1486.3750255780376, 8: -2043.933693071156, 9: -2715.318528602957, 15: -9287.407133426237, 16: -10834.4844708122, 17: -12522.649269035726, 35: -70045.28385080204, 53: -8102.524593409054}
2025-01-23 15:21:39.214 INFO: Average number of neighbors: 18.41771125793457
2025-01-23 15:21:39.215 INFO: During training the following quantities will be reported: energy, forces, virials, stress
2025-01-23 15:21:39.216 INFO: ===========MODEL DETAILS===========
2025-01-23 15:22:12.346 WARNING: Standard deviation of the scaling is zero, Changing to no scaling
2025-01-23 15:22:12.569 INFO: Loading FOUNDATION model
2025-01-23 15:22:12.572 INFO: Using filtered elements: [1, 6, 7, 8, 9, 15, 16, 17, 35, 53]
2025-01-23 15:22:12.573 INFO: Model configuration extracted from foundation model
2025-01-23 15:22:12.573 INFO: Using universal loss function for fine-tuning
2025-01-23 15:22:12.574 INFO: Message passing with hidden irreps 128x0e+128x1o)
2025-01-23 15:22:12.574 INFO: 2 layers, each with correlation order: 3 (body order: 4) and spherical harmonics up to: l=3
2025-01-23 15:22:12.575 INFO: Radial cutoff: 5.0 A (total receptive field for each atom: 10.0 A)
2025-01-23 15:22:12.575 INFO: Distance transform for radial basis functions: None
2025-01-23 15:22:13.927 INFO: Total number of parameters: 1437464
2025-01-23 15:22:13.928 INFO: 
2025-01-23 15:22:13.928 INFO: ===========OPTIMIZER INFORMATION===========
2025-01-23 15:22:13.929 INFO: Using ADAM as parameter optimizer
2025-01-23 15:22:13.929 INFO: Batch size: 2
2025-01-23 15:22:13.930 INFO: Using Exponential Moving Average with decay: 0.999
2025-01-23 15:22:13.930 INFO: Number of gradient updates: 50600
2025-01-23 15:22:13.931 INFO: Learning rate: 0.001, weight decay: 5e-07
2025-01-23 15:22:13.931 INFO: UniversalLoss(energy_weight=1.000, forces_weight=100.000, stress_weight=1.000)
2025-01-23 15:22:13.933 INFO: Using gradient clipping with tolerance=10.000
2025-01-23 15:22:13.933 INFO: 
2025-01-23 15:22:13.934 INFO: ===========TRAINING===========
2025-01-23 15:22:13.934 INFO: Started training, reporting errors on validation set
2025-01-23 15:22:13.935 INFO: Loss metrics on validation set
2025-01-23 15:25:33.634 INFO: Initial: head: pt_head, loss=0.00000000, RMSE_E_per_atom=2112625.45 meV, RMSE_F= 1078.86 meV / A, RMSE_stress=    0.00 meV / A^3
2025-01-23 15:25:35.317 INFO: Initial: head: linear_response_gs, loss=0.04175293, RMSE_E_per_atom=  938.64 meV, RMSE_F= 1288.17 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 15:25:37.019 INFO: Initial: head: linear_response_es1, loss=0.07281513, RMSE_E_per_atom=  711.70 meV, RMSE_F= 1684.90 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 15:25:38.761 INFO: Initial: head: spin_flip_es1, loss=0.02945673, RMSE_E_per_atom=  906.29 meV, RMSE_F=  735.57 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 15:25:40.434 INFO: Initial: head: spin_flip_es2, loss=0.06708954, RMSE_E_per_atom=  771.62 meV, RMSE_F= 1587.16 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 16:55:04.632 INFO: Epoch 0: head: pt_head, loss=0.00000000, RMSE_E_per_atom=2112531.72 meV, RMSE_F=  694.19 meV / A, RMSE_stress=    0.00 meV / A^3
2025-01-23 16:55:06.321 INFO: Epoch 0: head: linear_response_gs, loss=0.00980437, RMSE_E_per_atom=  374.92 meV, RMSE_F=  334.31 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 16:55:08.018 INFO: Epoch 0: head: linear_response_es1, loss=0.01700046, RMSE_E_per_atom=  377.06 meV, RMSE_F=  469.07 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 16:55:09.665 INFO: Epoch 0: head: spin_flip_es1, loss=0.00894455, RMSE_E_per_atom=  290.75 meV, RMSE_F=  228.85 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 16:55:11.264 INFO: Epoch 0: head: spin_flip_es2, loss=0.01759908, RMSE_E_per_atom=  407.45 meV, RMSE_F=  542.63 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 18:22:36.435 INFO: Epoch 1: head: pt_head, loss=0.00000000, RMSE_E_per_atom=2112517.82 meV, RMSE_F=  651.63 meV / A, RMSE_stress=    0.00 meV / A^3
2025-01-23 18:22:38.161 INFO: Epoch 1: head: linear_response_gs, loss=0.00757677, RMSE_E_per_atom=  389.37 meV, RMSE_F=  212.86 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 18:22:39.814 INFO: Epoch 1: head: linear_response_es1, loss=0.01554789, RMSE_E_per_atom=  332.18 meV, RMSE_F=  440.03 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 18:22:41.407 INFO: Epoch 1: head: spin_flip_es1, loss=0.00796123, RMSE_E_per_atom=  309.59 meV, RMSE_F=  215.93 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 18:22:43.095 INFO: Epoch 1: head: spin_flip_es2, loss=0.01687295, RMSE_E_per_atom=  378.75 meV, RMSE_F=  525.07 meV / A, RMSE_stress=    0.01 meV / A^3
2025-01-23 18:22:43.305 INFO: Training complete
2025-01-23 18:22:43.305 INFO: 
2025-01-23 18:22:43.306 INFO: ===========RESULTS===========
2025-01-23 18:22:43.306 INFO: Computing metrics for training, validation, and test sets
2025-01-23 18:22:43.310 INFO: Loading checkpoint: ./checkpoints/multi_head_v6_custom_run-42_epoch-12.pt

This run is configured for 2 epochs since processing the entire 1.51GB dataset is time-consuming. However, I have encountered the same issue when working with smaller subsets of the data and increasing the number of epochs.

Has anyone encountered a similar issue with multihead replay fine-tuning? Could this be due to:

• An error in the configuration file?
• Data imbalance between the pretraining dataset (pt_train_file) and individual heads?
• Incorrect learning rate or optimizer settings for pt_head?

Any insights, suggestions, or fixes would be greatly appreciated!

[ilyes319]

I think there might be a problem with the keys in pt head, what happens when you do use the pt_train_file: 'datasets/mace_datasets/test_large_neut_all.xyz' and just trains from scratch?