Source code for metatrain.experimental.dpa3.documentation
"""
DPA3 (Experimental)
======================
This is an interface to the DPA3 (Deep Potential Attention 3) architecture
:footcite:p:`dpa3_2025` implemented in `deepmd-kit
<https://github.com/deepmodeling/deepmd-kit>`_.
DPA3 extends the DPA series with a Line Graph representation and the RepFlow
framework, enabling richer many-body interactions through joint edge-angle
message passing. See the `paper <https://arxiv.org/abs/2506.01686>`_ and the
`deepmd-kit documentation <https://docs.deepmodeling.com/projects/deepmd/>`_
for further details.
.. note::
The ``type_map`` required by deepmd-kit is derived automatically from the
atomic numbers present in the dataset; it is *not* a user-facing
hyperparameter.
{{SECTION_INSTALLATION}}
{{SECTION_DEFAULT_HYPERS}}
Tuning hyperparameters
----------------------
The most impactful hyperparameters (roughly in decreasing order of importance):
.. container:: mtt-hypers-remove-classname
.. autoattribute:: {{model_hypers_path}}.descriptor
:no-index:
.. autoattribute:: {{trainer_hypers_path}}.learning_rate
:no-index:
.. autoattribute:: {{trainer_hypers_path}}.batch_size
:no-index:
Increasing ``descriptor.repflow.nlayers`` typically improves accuracy at the
cost of training time. ``descriptor.repflow.e_rcut`` controls the interaction
range and should be chosen based on the physical system. Reduce ``e_sel`` and
``a_sel`` for faster iteration on small systems.
Exporting a foundation DPA3 model
---------------------------------
As it is now, exporting a foundation DPA3 model from one of the `provided model
files <https://huggingface.co/deepmodelingcommunity>`_ involves using ``mtt train`` with
0 epochs.
Point ``dpa3_model`` at the checkpoint exactly as you downloaded it. These
checkpoints are trained on several datasets at once, and ``dpa3_model_branch``
selects the task to export; there is no need to run ``dp --pt freeze`` first,
and a frozen model is in fact rejected. If you do not know which tasks a
checkpoint contains, leave ``dpa3_model_branch`` out: the resulting error lists
every branch the checkpoint provides.
Use the following ``options.yaml`` file:
.. code-block:: yaml
architecture:
name: experimental.dpa3
model:
dpa3_model: path/to/deepmd-model.pt
dpa3_model_branch: OMol25
training:
num_epochs: 0
batch_size: 1
training_set: ./dummy_dataset.xyz
validation_set: ./dummy_dataset.xyz
with ``dummy_dataset.xyz`` being any dataset containing at least one structure with just
the ``energy`` property. For example, you can use:
.. code-block::
2
Properties=species:S:1:pos:R:3:forces:R:3 energy=-2.1
H 0.0 0.0 0.0 0.0 0.0 0.0
H 1.0 0.0 0.0 0.0 0.0 0.0
{{SECTION_MODEL_HYPERS}}
with the following definitions needed to fully understand some of the parameters:
.. autoclass:: {{architecture_path}}.documentation.DescriptorHypers
:members:
:undoc-members:
.. autoclass:: {{architecture_path}}.documentation.RepflowHypers
:members:
:undoc-members:
.. autoclass:: {{architecture_path}}.documentation.FittingNetHypers
:members:
:undoc-members:
"""
from typing import Literal, Optional
from typing_extensions import NotRequired, TypedDict
from metatrain.composition.documentation import FixedCompositionWeights
from metatrain.utils.hypers import init_with_defaults
from metatrain.utils.loss import LossSpecification
###########################
# MODEL HYPERPARAMETERS #
###########################
[docs]
class RepflowHypers(TypedDict):
"""RepFlow descriptor block parameters."""
n_dim: int = 128
"""Node feature dimension."""
e_dim: int = 64
"""Edge feature dimension."""
a_dim: int = 32
"""Angle feature dimension."""
nlayers: int = 6
"""Number of RepFlow interaction layers."""
e_rcut: float = 6.0
"""Edge (pair) cutoff radius in length units."""
e_rcut_smth: float = 5.3
"""Start of cosine smoothing for the edge cutoff."""
e_sel: int = 1200
"""Maximum number of edge neighbors per atom."""
a_rcut: float = 4.0
"""Angle (triplet) cutoff radius in length units."""
a_rcut_smth: float = 3.5
"""Start of cosine smoothing for the angle cutoff."""
a_sel: int = 300
"""Maximum number of angle neighbors per atom."""
axis_neuron: int = 4
"""Number of axis neurons in the embedding network."""
skip_stat: bool = True
"""Skip statistics computation (use pretrained stats)."""
a_compress_rate: int = 1
"""Compression rate for angle features."""
a_compress_e_rate: int = 2
"""Compression rate for angle-edge features."""
a_compress_use_split: bool = True
"""Use split compression for angle features."""
update_angle: bool = True
"""Update angle features at each layer."""
update_style: str = "res_residual"
"""Residual update style. Options: ``"res_residual"``, ``"res_avg"``."""
update_residual: float = 0.1
"""Residual scaling factor for updates."""
update_residual_init: str = "const"
"""Initialisation for the residual scaling. Options: ``"const"``,
``"norm"``."""
smooth_edge_update: bool = True
"""Apply smooth cutoff function to edge updates."""
use_dynamic_sel: bool = True
"""Dynamically adjust neighbor selection at runtime."""
sel_reduce_factor: float = 10.0
"""Reduction factor for dynamic neighbor selection."""
[docs]
class DescriptorHypers(TypedDict):
"""Descriptor hyperparameters wrapping the RepFlow block."""
type: str = "dpa3"
"""Descriptor type identifier used by deepmd-kit."""
repflow: RepflowHypers = init_with_defaults(RepflowHypers)
"""RepFlow block parameters."""
activation_function: str = "custom_silu:10.0"
"""Activation function. Format: ``"name"`` or ``"name:param"``.
Supported names include ``"tanh"``, ``"gelu"``, ``"custom_silu"``."""
use_tebd_bias: bool = False
"""Add bias to the type embedding."""
precision: int = 32
"""Floating-point precision for the descriptor (``32`` or ``64``).
This controls the internal precision of deepmd-kit's descriptor
computation. For mixed-precision training, set this independently
of ``fitting_net.precision``; for uniform precision, set both to the
same value and match ``base_precision`` accordingly."""
concat_output_tebd: bool = False
"""Concatenate type embedding to descriptor output."""
[docs]
class FittingNetHypers(TypedDict):
"""Fitting network hyperparameters."""
neuron: list[int] = [240, 240, 240]
"""Hidden layer sizes for the fitting network."""
resnet_dt: bool = True
"""Use a ResNet-style time step in each hidden layer."""
seed: int = 1
"""Random seed for weight initialisation."""
precision: int = 32
"""Floating-point precision for the fitting network (``32`` or ``64``).
Can differ from ``descriptor.precision`` for mixed-precision training."""
activation_function: str = "custom_silu:10.0"
"""Activation function (same format as the descriptor)."""
type: str = "ener"
"""Fitting type. ``"ener"`` for energy fitting."""
numb_fparam: int = 0
"""Number of frame-level parameters."""
numb_aparam: int = 0
"""Number of atom-level parameters."""
dim_case_embd: int = 0
"""Dimension of the case embedding (multi-task)."""
trainable: bool = True
"""Whether fitting network weights are trainable."""
rcond: Optional[float] = None
"""Cutoff for pseudo-inverse in linear fitting."""
atom_ener: list[float] = []
"""Per-type atomic energy offsets."""
use_aparam_as_mask: bool = False
"""Treat atom-level parameters as a mask."""
class ModelHypers(TypedDict):
"""Hyperparameters for the DPA3 model.
The ``type_map`` needed by deepmd-kit is derived automatically from the
dataset's atomic numbers and should not be set manually.
"""
dpa3_model: Optional[str] = None
"""Path to a deepmd-kit training checkpoint, used exactly as distributed
(for example ``DPA-3.3-1M.pt``). A frozen model written by
``dp --pt freeze`` is not accepted, since the branch selection and the
model configuration metatrain needs both live in the checkpoint. When
provided, the model weights are loaded from this file instead of being
initialised from scratch. Energy biases and standard deviations stored in
the deepmd-kit model are extracted and handed to metatrain's
``CompositionModel`` and ``Scaler`` so that fine-tuning starts from the
pretrained values."""
dpa3_model_branch: Optional[str] = None
"""Name of the task to load from a multi-task pretrained DPA3 model. If
the model has several branches and this is not set, the resulting error
lists the branches the checkpoint provides."""
descriptor: DescriptorHypers = init_with_defaults(DescriptorHypers)
"""Descriptor configuration (RepFlow block and related settings)."""
fitting_net: FittingNetHypers = init_with_defaults(FittingNetHypers)
"""Fitting network configuration."""
##############################
# TRAINER HYPERPARAMETERS #
##############################
class TrainerHypers(TypedDict):
"""Hyperparameters for training DPA3 models."""
distributed: NotRequired[bool]
"""Whether to use distributed training. When not set, distributed training
is enabled automatically when running under more than one SLURM task.
Setting this option explicitly is deprecated."""
distributed_port: int = 39591
"""Port for DDP communication."""
batch_size: int = 8
"""The number of samples to use in each batch of training. This
hyperparameter controls the tradeoff between training speed and memory usage. In
general, larger batch sizes will lead to faster training, but might require more
memory."""
max_atoms_per_batch: Optional[int] = None
"""If set, use greedy atom-count packing instead of fixed ``batch_size``.
Structures are accumulated into each batch until adding another would exceed this
limit, producing variable numbers of structures per batch. Supported with any
dataset type. When set, ``batch_size`` is ignored for constructing training
and validation batches (it is still used internally for composition model and
scaler fitting)."""
min_atoms_per_batch: int = 0
"""Minimum total number of atoms required to keep a batch when
``max_atoms_per_batch`` is set. Batches whose total atom count falls below this
threshold are discarded during packing. Defaults to ``0`` (no minimum)."""
num_epochs: int = 100
"""Number of epochs."""
learning_rate: float = 0.001
"""Learning rate."""
scheduler_patience: int = 100
"""Number of epochs with no improvement before reducing the learning rate."""
scheduler_factor: float = 0.8
"""Factor by which the learning rate is reduced on plateau."""
log_interval: int = 1
"""Interval to log metrics."""
checkpoint_interval: int = 100
"""Interval to save checkpoints."""
scale_targets: bool = True
"""
Normalize targets to unit std during training.
If true, a single scale is computed for each target, given by the uncentered
standard deviation across all values in the dataset for that target.
For targets with more than one property (i.e. > 1 block or >= 1 block with > 1
property), per-property scales are also computed, and used to re-scale model
predictions.
See also :ref:`scale-targets`.
"""
fixed_composition_weights: FixedCompositionWeights = {}
"""Weights for atomic contributions.
This is passed to the ``fixed_weights`` argument of
:meth:`CompositionModel.train_model
<metatrain.composition.CompositionModel.train_model>`,
see its documentation to understand exactly what to pass here.
"""
per_structure_targets: list[str] = []
"""Targets to calculate per-structure losses."""
log_mae: bool = False
"""Log MAE alongside RMSE."""
log_separate_blocks: bool = False
"""Log per-block error."""
best_model_metric: Literal["rmse_prod", "mae_prod", "loss"] = "rmse_prod"
"""Metric used to select best checkpoint (e.g., ``rmse_prod``)."""
loss: str | dict[str, LossSpecification] = "mse"
"""This section describes the loss function to be used. See the
:ref:`loss-functions` for more details."""