xnn.common.finetune#

Fine-tuning pretrained potentials: multi-head replay, LoRA, references.

The strategies of the MACE fine-tuning protocols (Batatia et al., arXiv:2401.00096; Tompa et al., arXiv:2606.12704; Wang et al., ELoRA, ICML 2025), available to every xnn model that supports them:

  • naive fine-tuning: model: {pretrained: <any hub model>} in a training config continues training every weight;

  • layer freezing: freeze_parameters(), or optim.freeze / optim.train_only patterns;

  • LoRA: inject_lora() / merge_lora(), or model.extra["lora"];

  • multi-head replay: MultiHead, select_replay(), pseudolabel(), or model.extra["heads"] with data.replay_path;

  • reference energies: estimate_atomic_energies() (model-aware reestimation) and average_atomic_energies(), or atomic_energies: estimated / average in the config.

See the how-to guide on fine-tuning for the recipes.

Modules

freeze

Freezing parameters by name pattern (layer-freezing fine-tuning).

heads

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

lora

Low-rank adaptation (LoRA) of the linear layers of any xnn potential.

reference

Per-element reference energies (E0s) for fine-tuning.

replay

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