.. _transfer-learning:

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`` 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:

.. code-block:: yaml

  architecture:
     training:
       finetune:
        read_from: path/to/checkpoint.ckpt

  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 :ref:`Fitting generic targets <fitting-generic-targets>`
section of the documentation.