Transfer Learning (experimental)

Warning

This section of the documentation is only relevant for PET model so far.

Warning

Features described in this section are experimental and not yet extensively tested. Please use them at your own risk and report any issues you encounter to the developers.

This section describes the process of transfer learning, which is a common technique used in machine learning, where a model is pre-trained on the dataset with one level of theory and/or one set of properties and then fine-tuned on a different dataset with a different level of theory and/or different set of properties. This approach to use the learned representations from the pre-trained model and adapt them to the targets, which can be expensive to compute and/or not available in the pre-trained dataset.

In the following sections we assume that the pre-trained model is trained on the conventional DFT dataset with energies, forces and stresses, which are provided as energy targets (and its derivatives) in the options.yaml file.

Fitting to a new level of theory

Training on a new level of theory is a common use case for transfer learning. It requires using a pre-trained model checkpoint with mtt train -c pre-trained-model.ckpt command and setting the new targets corresponding to the new level of theory in the options.yaml file. Let’s assume that the training is done on the dataset computed with the hybrid DFT functional (e.g. PBE0) stored in the new_train_dataset.xyz file, where the corresponsing energies are written in the energy key of the info dictionary of the ase.Atoms object. Then, the options.yaml file should look like this:

training_set:
  systems: "new_train_dataset.xyz"
  targets:
    mtt::energy_pbe0: # name of the new target
    key: "energy" # key of the target in the atoms.info dictionary
    unit: "eV" # unit of the target value

The validation and test sets can be set in the same way. The training process will then create a new composition model and new heads for the target energy_pbe0. The rest of the model weights will be initialized from the pre-trained model checkpoint.

Fitting to a new set of properties

Training on a new set of properties is another common use case for transfer learning. It can be done in a similar way as training on a new level of theory. The only difference is that the new targets need to be properly set in the options.yaml file. More information about fitting the generic targets can be found in the Fitting generic targets section of the documentation.