Question about finetuning on OAM model with MPRelaxSet settings

[Asif-Iqbal-Bhatti]

Hello,

I am new in the ML field, and I want to fine-tune rather than train from scratch for my system. I followed the recipe as stated on the NequI webpage. I am fine-tuning the NequIP-OAM model with my own data that is generated with the MP MPRelaxSet standard using PBE functional. I compare the reference energy, and they are compatible with MP reference energies to within 0.001 eV/atom. Now when I'm fine on NequIP, my energy and stress are offset. Below are the details.

Could you tell me when fine-tuning on MP, what should I take care of? I can think there is a mismatch of reference energies with the OAM model and MP, but NequiP is fine-tuned on MPTraj. What settings need to be changed? Thank you for your inputs or anything I need to to know before doing trail and error analysis.

  model:
    _target_: nequip.model.modify
    modifiers:
      - modifier: modify_PerTypeScaleShift
        shifts:
          F: -0.25125556
          Cl: -0.88440305
          P: -1.884791
          Na: -0.29658376
        scales: ${training_data_stats:per_type_forces_rms}
        shifts_trainable: true
        scales_trainable: true

here is my config file

# ===========
#     RUN
# ===========
run: [train, val, test]

cutoff_radius: 6.0
chemical_symbols: [Cl, S, P, Li]
model_type_names: ${chemical_symbols}

# ============
#     DATA
# ============
data:
  _target_: nequip.data.datamodule.ASEDataModule
  seed: 7521

  #split_dataset:
  #  file_path: ./train_data.extxyz
  #  train: 0.8
  #  val: 0.1
  #  test: 0.1
  train_file_path: ./train_data.extxyz
  val_file_path: ./valid_data.extxyz
  test_file_path: ./test_data.extxyz
  predict_file_path: ./predict_data.extxyz
  
  transforms:
    - _target_: nequip.data.transforms.NeighborListTransform
      r_max: ${cutoff_radius}
    - _target_: nequip.data.transforms.ChemicalSpeciesToAtomTypeMapper
      chemical_symbols: ${chemical_symbols}

  train_dataloader:
    _target_: torch.utils.data.DataLoader
    batch_size: 4
    num_workers: 1
    shuffle: true
  val_dataloader:
    _target_: torch.utils.data.DataLoader
    batch_size: 4
    num_workers: ${data.train_dataloader.num_workers}
  test_dataloader: ${data.val_dataloader}

  stats_manager:
    _target_: nequip.data.CommonDataStatisticsManager
    dataloader_kwargs:
      batch_size: 24
    type_names: ${model_type_names}

# ============
#   TRAINER
# ============
trainer:
  _target_: lightning.Trainer
  accelerator: auto
  enable_checkpointing: true
  max_epochs: 250
  max_time: 03:00:00:00
  log_every_n_steps: 1

  logger:
    #_target_: lightning.pytorch.loggers.WandbLogger
    #project: nequip_finetune
    #name: run_${now:%Y-%m-%d_%H-%M-%S}
    #save_dir: ./wandb_logs
    #log_model: false

    _target_: lightning.pytorch.loggers.csv_logs.CSVLogger
    name: tutorial_log
    save_dir: ./results  # use resolver to place wandb logs in hydra's output directory
    flush_logs_every_n_steps: 50
    
  callbacks:
    - _target_: lightning.pytorch.callbacks.EarlyStopping
      monitor: val0_epoch/weighted_sum
      min_delta: 1e-3
      patience: 10
    - _target_: lightning.pytorch.callbacks.ModelCheckpoint
      monitor: val0_epoch/weighted_sum
      dirpath: ./results
      filename: best
      save_last: true
    - _target_: lightning.pytorch.callbacks.LearningRateMonitor
      logging_interval: epoch

# ============
# TRAINING MODULE + MODEL
# ============
training_module:
  _target_: nequip.train.EMALightningModule
  ema_decay: 0.999

  loss:
    _target_: nequip.train.EnergyForceStressLoss
    per_atom_energy: true
    coeffs:
      total_energy: 1.0
      forces: 1.0
      stress: 1.0

  # REQUIRED: validation metrics
  val_metrics:
    _target_: nequip.train.EnergyForceStressMetrics
    coeffs:
      total_energy_mae: 1.0
      total_energy_rmse: 1.0
      per_atom_energy_mae: 1.0
      per_atom_energy_rmse: 1.0
      forces_mae: 1.0
      forces_rmse: 1.0

  train_metrics: ${training_module.val_metrics}
  test_metrics: ${training_module.val_metrics}

  optimizer:
    _target_: torch.optim.Adam
    lr: 3e-4    # 0.03 was far too high for fine-tuning, lower to avoid divergence

  lr_scheduler:
    scheduler:
      _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
      factor: 0.6
      patience: 5
      threshold: 0.2
      min_lr: 1e-6
    monitor: val0_epoch/weighted_sum
    interval: epoch
    frequency: 1

  model:
    _target_: nequip.model.modify
    modifiers:
      - modifier: modify_PerTypeScaleShift
        shifts:
          F: -0.25125556
          Cl: -0.88440305
          P: -1.884791
          Na: -0.29658376
        scales: ${training_data_stats:per_type_forces_rms}
        shifts_trainable: true
        scales_trainable: true

    model:
      _target_: nequip.model.ModelFromPackage
      package_path: /dodrio/scratch/projects/starting_2025_087/ML_Training_Models/NequIP_training/NequIP-OAM-L-0.1.nequip.zip
pair_potential:
  _target_: nequip.nn.pair_potential.ZBL
  units: metal
  chemical_species: ${chemical_symbols}

[kavanase]

Hi @Asif-Iqbal-Bhatti ,
What do your training curves look like? Does the model learn as expected, or does the error just stay high?

This looks like something seriously wrong. My initial thoughts are that it could be related to the scales choice (I suggest trying without setting scales in your config, and just using the pre-trained scales as the starting point), or potentially similar to recent chemical species mapping issues as I guess you're using a slightly older version – @cw-tan will know for this.

[Asif-Iqbal-Bhatti]

Hi @kavanase,

Thank you for your prompt reply. Regarding fine-tuning, I found the information in the fine-tuning section of the website; however, I could not find the YAML file for fine-tuning, so I assume it is similar to starting from scratch. Here is the training loss plot. I created the plot using the script provided in the tutorial.
. As you can see, it did not go to a large epoch, as there is an early stopping criterion, and I was told for fine-tuning it is okay to fine-tune on a few dozen structures.

I did with and without, and still I face this offset. I see NequIP is trained on OMAT24 and then fine-tuned on MPtraj, and I have done my calculation using MPRelaxSet for the PBE-level theory. I assume my reference energy aligns well with the OAM model. I checked my reference energies using MACE code, and they are in agreement. Could you tell me what factors I should change?
batch size
My training data consists of grain boundaries (all relaxed ionic steps from VASP foare in traiing data or few middle).
I have also varied the " loss:
target: nequip.train.EnergyForceStressLoss
per_atom_energy: true
coeffs:
total_energy: 100.0
forces: 10.0
stress: 5.0"

Could you provide me the tutorial on fine-tuning with all the hyperparameters if we want to fine-tune on MPRelaxSet PBE data?

Regards,
Asif

[cw-tan]

I'll make a new release soon, but there are safety issues with the chemical species when fine-tuning. For example, in v0.15.0 (current stable release), it is wrong to use

chemical_symbols: [Cl, S, P, Li]
model_type_names: ${chemical_symbols}

when fine-tuning the foundation potential -- the chemical symbols should correspond to the exact set of chemical symbols used in the foundation potential, otherwise it's being mapped to the wrong indices and you might not benefit from the pretraining.

Two solutions for now:

1. you can check the training config from https://zenodo.org/records/16980200 and use the exact same order and list of chemical symbols
2. what I recommend is to install nequip from source and use the develop branch, i.e.

git clone https://github.com/mir-group/nequip.git
cd nequip
git checkout develop
pip install -e .

There are some breaking API changes related to the chemical symbols for safety reasons (like the one above). Note specifically here:

nequip/configs/tutorial.yaml

Line 66 in 398a232

- _target_: nequip.data.transforms.ChemicalSpeciesToAtomTypeMapper

You can use these develop fine-tuning docs: https://nequip.readthedocs.io/en/develop/guide/training-techniques/fine_tuning.html, which have more details to account for the atom type mapping. NOTE how this is the develop ones and not the latest ones (which correspond to stable releases). I plan to make an official release before the end of this month anyway, so better to migrate now if possible.

[Asif-Iqbal-Bhatti]

Okay, thanks for the information. What I understood is that it has to respect the order of the chemical symbols? After using model_atom_types, there is still an offset. I am using some common symbols as they are in the pretrained dataset. I see there is a slightly larger reference energy offset for P and S in my fine-tuned dataset than yours. Could that be the case? I think MP PBE uses old POTCARS, and I have that methodology, and OMAT has chosen PBE54. My system has complex grain boundary structures, and it is difficult to train on a single point, as these are hypothetical and the initial forces are large, and I am relaxing with DFT first and then fine-tuning on relaxed structures so as not to have large forces, which I then want to use to study dynamics using NEB.

[cw-tan]

Are you trying the first suggested solution? What I mean is that you'll need ALL the chemical symbols in the same order used for training the foundation potential, which I think is

chemical_symbols = [
    "H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne",
    "Na", "Mg", "Al", "Si", "P", "S", "Cl", "Ar", "K", "Ca",
    "Sc", "Ti", "V", "Cr", "Mn", "Fe", "Co", "Ni", "Cu", "Zn",
    "Ga", "Ge", "As", "Se", "Br", "Kr", "Rb", "Sr", "Y", "Zr",
    "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", "Ag", "Cd", "In", "Sn",
    "Sb", "Te", "I", "Xe", "Cs", "Ba", "La", "Ce", "Pr", "Nd",
    "Pm", "Sm", "Eu", "Gd", "Tb", "Dy", "Ho", "Er", "Tm", "Yb",
    "Lu", "Hf", "Ta", "W", "Re", "Os", "Ir", "Pt", "Au", "Hg",
    "Tl", "Pb", "Bi", "Ac", "Th", "Pa", "U", "Np", "Pu"
]

which might be inconvenient, hence why I suggested the second solution.

Also, if you're fine-tuning, you should be changing the shifts to your dataset, typically computed by doing a DFT calculation on a big box with a single isolated atom.

[Asif-Iqbal-Bhatti]

Okay, thanks. I will try the 2nd option and will let you know. I will do that today.

[cw-tan]

FYI that the recent release (https://github.com/mir-group/nequip/releases/tag/v0.16.0) should make this easier. See latest fine-tuning docs: https://nequip.readthedocs.io/en/latest/guide/training-techniques/fine_tuning.html

[Asif-Iqbal-Bhatti]

Okay, I will use this new version or training.