xnn.common.finetune.replay

Replay data for multi-head fine-tuning: selection and pseudolabels.

Multi-head replay trains a second head on structures from the pretraining distribution so the shared trunk keeps the foundation model’s breadth (Tompa et al., arXiv:2606.12704, Sec. 2.5). The replay set is usually a subsample of the pretraining data, element-matched to the fine-tuning set, with its original labels; with pseudolabels (the pretrained model’s own predictions) any structurally diverse set serves, and the replay head is asked to reproduce the pretrained model rather than its reference data.

Module Attributes

REPLAY_FILTERS

Element filters of select_replay(), relative to the fine-tuning elements: keep every structure (none); only those made of fine-tuning elements (subset, the default and the paper's protocol); exactly the fine-tuning element set (exact); those containing all fine-tuning elements (superset).

Functions

element_filter(structures, species[, mode])	Keep the structures whose elements relate to species as mode says.
pseudolabel(model, structures[, batch_size, ...])	Label structures with a model's own energies, forces (and stress).
select_replay(structures[, species, n, ...])	Assemble a replay set: element filter, then a random subsample.