Composition¶

Maintained by @ppegolo.

The composition model computes per-species contributions to invariant targets by solving a deterministic least-squares problem. It is typically used as an additive baseline within other architectures (e.g. PET, SOAP-BPNN, PhACE) to capture compositional offsets before training the main model, but it can also be trained on its own. The main use case for standalone training is to produce a checkpoint that initializes the composition baseline of another architecture, avoiding a re-fit of the weights.

How the model works¶

The composition model predicts a target as a sum of per-species weights \(w_t\), one for each atomic type \(t\), fitted by least squares on the training data. Only invariant quantities are fitted (energy being the typical example). For spherical targets, these are blocks with \(\lambda = 0, \sigma = 1\). The fit can be bypassed for chosen targets and atomic types with the atomic_baseline trainer hyperparameter, which sets the corresponding weights to fixed, user-supplied values.

The exact form of the fit depends on the sample kind of the target.

For per-structure targets like energy (or anything with sample_kind set to "system"), the prediction for a structure \(A\) is

\[\hat{y}_A = \sum_{i \in A} w_{t_i} = \sum_t n^A_t \, w_t,\]

where \(t_i\) is the type of atom \(i\) and \(n^A_t\) is the number of atoms of type \(t\) in \(A\). The weights minimize \(\sum_A \lVert y_A - \sum_t n^A_t w_t \rVert^2\), i.e. they solve the normal equations \(X^\top X \, W = X^\top Y\), where \(X_{At} = n^A_t\), \(Y\) stacks the target values \(y_A\), and \(W\) stacks the weights \(w_t\) (a small ridge term is added to the diagonal for numerical stability).

For per-atom targets (i.e. anything with sample_kind set to "atom"), the prediction for atom \(i\) is simply \(\hat{y}_i = w_{t_i}\). In this case \(X\) is a one-hot encoding of the atomic types, so \(X^\top X\) is diagonal (the per-type atom counts \(N_t\)) and the fit reduces to

\[w_t = \frac{1}{N_t} \sum_{i \,:\, t_i = t} y_i,\]

that is, the mean of the target over all atoms of type \(t\). Types that never appear in the training data are assigned a zero weight.

As a simple example, take a per-atom target where the hydrogen atoms carry the values 1.0 and 1.5 and a single oxygen atom carries 2.0: the fit gives \(w_\mathrm{H} = 1.25\) and \(w_\mathrm{O} = 2.0\), so every hydrogen is predicted as 1.25 and every oxygen as 2.0. For a per-structure energy, a water molecule would instead be predicted as \(2 w_\mathrm{H} + w_\mathrm{O}\).

Installation¶

To install this architecture along with the metatrain package, run:

pip install metatrain[composition]

where the square brackets indicate that you want to install the optional dependencies required for composition.

Default Hyperparameters¶

The description of all the hyperparameters used in composition is provided further down this page. However, here we provide you with a yaml file containing all the default hyperparameters, which might be convenient as a starting point to create your own hyperparameter files:

architecture:
  name: composition
  model: {}
  training:
    atomic_baseline: {}
    batch_size: null

Model hyperparameters¶

This architecture has no model hyperparameters: there is nothing to set under the architecture.model section of the config file.

Trainer hyperparameters¶

The parameters that go under the architecture.trainer section of the config file are the following:

TrainerHypers.atomic_baseline: Dict[int, float]]]], FieldInfo(annotation=NoneType, required=False, default={})] = {}¶

Fixed per-species baselines, overriding the least-squares fit for the targets/atomic types they cover. A dict mapping each target name to either a single weight for all atomic types, or a dict mapping atomic types to weights. Unlike the identically-named hyperparameter of other architectures, a path to a composition checkpoint is not accepted here.

TrainerHypers.batch_size: Annotated[NotRequired[NotRequired[int | None]], FieldInfo(annotation=NoneType, required=False, default=None)] = None¶

Number of structures to accumulate at a time when building the least-squares problem. This only affects memory usage, not the fitted weights, since the composition model is a deterministic fit rather than an iterative optimization. Defaults to the size of the smallest training dataset.