Adding a new MLIP-only architecture

For MLIP-only models (models that only predict energies and forces), metatrain provides base classes metatrain.utils.mlip.MLIPModel and metatrain.utils.mlip.MLIPTrainer that implement most of the boilerplate code. See MLIP Base Classes for more details.

Example: Creating an MLIP-only architecture

To demonstrate how easy it is to add a new MLIP-only architecture using the base classes, let’s look at the mlip_example architecture in metatrain. This minimal architecture always predicts zero energy, serving as a simple template for MLIP development.

The model (model.py) only needs to implement the compute_energy method:

from metatrain.utils.mlip import MLIPModel

class ZeroModel(MLIPModel):
    """A minimal example MLIP model that always predicts zero energy."""

    __checkpoint_version__ = 1

    def __init__(self, hypers, dataset_info):
        super().__init__(hypers, dataset_info)
        # Request a neighbor list with the cutoff from hyperparameters
        cutoff = hypers["cutoff"]
        self.request_neighbor_list(cutoff)

    def compute_energy(
        self,
        edge_vectors,
        species,
        centers,
        neighbors,
        system_indices,
    ):
        # Get the number of systems and return zeros
        n_systems = system_indices.max().item() + 1
        return torch.zeros(n_systems, device=edge_vectors.device)

The trainer (trainer.py) only needs to specify whether to use rotational augmentation:

from metatrain.utils.mlip import MLIPTrainer

class ZeroTrainer(MLIPTrainer):
    """Trainer for the ZeroModel."""

    __checkpoint_version__ = 1

    def use_rotational_augmentation(self):
        return False  # No rotational augmentation for this example

That’s it! The base classes handle all the training loop, data loading, composition weights, scaling, checkpointing, and export functionality. This allows you to focus on implementing the core physics of your model in the compute_energy method.

The complete example architecture can be found in src/metatrain/mlip_example/.