xnn.common.finetune.heads#

Multi-head potentials: one shared trunk, one readout and reference per head.

Multi-head replay fine-tuning (Batatia et al., arXiv:2401.00096; Tompa et al., arXiv:2606.12704) trains a pretrained model on the target data and on a replay set of the pretraining distribution at the same time. The two sets share the embedding and message-passing layers and differ only in the readout and in the per-element reference energies, so the replay head keeps the model close to its pretraining behaviour while the target head adapts.

MultiHead gives any xnn potential that declares its readout modules (head_modules) several such heads. Head 0 is the wrapped model’s own readout; every further head is an independent copy of those modules. A structure is routed to its head by the head field of the AtomicGraph (set with label_head()); a batch may mix heads, each sub-batch runs through the trunk once with its own head’s parameters substituted, and the outputs are reassembled in the original order. A deployed model has one head: MultiHead.select() returns the plain potential of a head.

Enable from a config with model.extra["heads"]:

model:
  pretrained: mace-mp-0-small
  heads: [pt_head, Default]          # head 0 keeps the pretrained readout and E0s
  # or, with per-head reference energies:
  heads: {pt_head: null, Default: {atomic_energies: estimated}}

Functions

find_multihead(model)

The MultiHead inside a (possibly wrapped) model, if any.

head_module_names(model)

The attributes of model that make up one readout head.

head_names(spec)

The head names of a model.extra["heads"] entry (a list or a dict).

head_options(spec, head)

The per-head options of head in a model.extra["heads"] mapping ({} for a list).

label_head(structures, head)

Tag structure dicts (or a dataset of them) with a readout head.

Classes

MultiHead(model, heads)

A potential with several readout heads over one shared trunk.
Show Source