Descriptor DPA-2

4.6. Descriptor DPA-2 TensorFlow PyTorch JAX Paddle DP#

Note

Supported backends: TensorFlow 2 TensorFlow, PyTorch-TorchScript and PyTorch-Exportable PyTorch, JAX JAX, Paddle Paddle, DP DP

The DPA-2 model implementation. See DPA-2 paper for more details.

Training example: examples/water/dpa2/input_torch_medium.json, see README for inputs in different levels.

4.6.1. Theory#

DPA-2 is an attention-based descriptor architecture proposed for large atomic models (LAMs); see the DPA-2 paper.

At a high level, DPA-2 builds local representations with three coupled channels (paper notation):

  • Single-atom channel \(\mathbf{f}_\alpha\)

  • Rotationally invariant pair channel \(\mathbf{g}_{\alpha\beta}\)

  • Rotationally equivariant pair channel \(\mathbf{h}_{\alpha\beta}\)

for neighbors \(\beta\in\mathcal{N}(\alpha)\) within cutoffs.

4.6.1.1. Descriptor pipeline#

The descriptor follows two main stages:

  1. repinit (representation initializer)

    • Initializes and fuses type and geometry information from local environments.

  2. repformer (representation transformer)

    • Stacked message-passing layers that update \(\mathbf{f}\), \(\mathbf{g}\), and per-atom representations \(\mathbf{h}\) through convolution/symmetrization/MLP and attention-style interactions.

The final descriptor is formed from learned single-atom representations and then passed to downstream fitting/model components.

4.6.1.2. Message-passing intuition#

DPA-2 updates local features layer-by-layer with residual connections. Conceptually, each layer performs neighborhood aggregation using geometry-conditioned interactions:

\[\mathbf{h}_\alpha^{(l+1)} = \mathbf{h}_\alpha^{(l)} + \mathrm{MP}^{(l)}\left(\mathbf{h}_\alpha^{(l)}, \{\mathbf{h}_\beta^{(l)}\}_{\beta\in\mathcal{N}(\alpha)}, \{\mathbf{g}_{\alpha\beta}\}_{\beta\in\mathcal{N}(\alpha)}\right)\]

where \(\mathrm{MP}^{(l)}\) denotes the layer-specific message-passing operator.

4.6.1.3. Physical properties#

Consistent with the DPA-2 design goals in the paper, the model family is built to satisfy:

  1. Translational invariance (depends on relative coordinates)

  2. Rotational and permutational symmetry requirements

  3. Conservative formulation when used in energy models (forces/virials from energy gradients)

  4. Smoothness up to second-order derivatives

4.6.1.4. Multi-task training context#

DPA-2 is designed for multi-task pre-training with a shared descriptor and task-specific downstream heads. See Multi-task training for workflow details.

4.6.2. Requirements of installation PyTorch#

If one wants to run the DPA-2 model on LAMMPS, the customized OP library for the Python interface must be installed when freezing the model.

The customized OP library for the Python interface is installed by default when building DeePMD-kit from source.

If one runs LAMMPS with MPI, the customized OP library for the C++ interface should be compiled against the same MPI library as the runtime MPI. If one runs LAMMPS with MPI and CUDA devices, it is recommended to compile the customized OP library for the C++ interface with a CUDA-Aware MPI library and CUDA, otherwise the communication between GPU cards falls back to the slower CPU implementation.

4.6.3. Limitations of the JAX backend with LAMMPS JAX#

When using the JAX backend, 2 or more MPI ranks are not supported. One must set map to yes using the atom_modify command.

atom_modify map yes

See the example examples/water/lmp/jax_dpa.lammps.

4.6.4. Data format#

DPA-2 supports both the standard data format and the mixed type data format.

4.6.5. Type embedding#

Type embedding is within this descriptor with the tebd_dim argument.

4.6.6. Model compression#

Model compression is supported when repinit/tebd_input_mode is strip.

  • If repinit/attn_layer is 0, both the type embedding and geometric parts inside repinit are compressed.

  • If repinit/attn_layer is not 0, only the type embedding tables are compressed and the geometric attention layers remain as neural networks.

An example is given in examples/water/dpa2/input_torch_compressible.json. The performance improvement will be limited if other parts are more expensive.

4.6.7. Graph-native inference route (pt_expt) PyTorch#

In the pt_expt backend, a graph-eligible DPA-2 descriptor (repinit/use_three_body false – the three-body sub-block is not graph-eligible – and not compressed) can be frozen through a NeighborGraph-native inference path instead of the legacy dense neighbor-list path:

dp --pt-expt freeze -o model.pt2 --lower-kind graph

As with DPA-1’s graph path (see Difference among different backends), the graph route considers all neighbors within the cutoff rather than a fixed, padded selection, so its numeric result can differ slightly (down to the AOTInductor floating-point noise floor at non-binding sel, larger if sel is binding) from the dense/nlist path.

Note

Default route change in pt_expt (eager & training). For a graph-eligible DPA-2 descriptor, the pt_expt backend now defaults to the carry-all graph route not only for --lower-kind graph freezing but also in eager inference/evaluation and in (compiled) training (neighbor_graph_method=None resolves to the graph). This changes the numerical behavior of existing pt_expt configurations relative to the dense neighbor-list route (by the amounts described above — negligible at non-binding sel). The other backends (dpmodel/PyTorch-TorchScript/Paddle/TensorFlow/JAX) are unaffected: they keep the dense route as their only path.

To retain the legacy dense route on pt_expt:

  • Inference / evaluation: pass neighbor_graph_method="legacy" to forward_common / call_common (forces the dense neighbor-list path).

  • Training: call model.atomic_model.descriptor.disable_graph_lower() on the constructed model before training. This flips uses_graph_lower() to False, which both the eager forward and the compiled-training lower honor, so both run the dense route consistently (the same mechanism the spin model uses to stay on the dense path). A first-class training-config knob for this is planned as a follow-up.

Note

Smoothness at the cutoff. The graph route is exactly smooth at the cutoff, like the dense path. The non-attention channels (environment matrix, switch envelope, convolution, drrd/grrg, g1g1, symmetrization) are smooth by construction. The repformer attention channels (update_g1_has_attn, update_g2_has_attn) additionally use a fixed-phantom-count softmax: the dense smooth-attention denominator keeps exactly sel − n_real padding terms at \(e^{-\mathrm{attnw\_shift}}\) (a geometry-independent count); the graph kernels reproduce this by excluding masked pairs from the softmax and adding \(\max(\mathrm{sel} - n_\mathrm{real}, 0)\) phantom denominator terms per center. An edge entering the cutoff sphere does so at logit \(-\mathrm{attnw\_shift}\) exactly while the phantom count drops by one, so the swap is value-preserving and the energy/force are continuous (verified at the float64 noise floor, \(\lesssim 10^{-13}\)). This also makes the carry-all graph attention agree with the dense attention term-for-term at non-binding sel. The only residual \(e^{-20}\)-scale discontinuity remains for a center with sel or more real neighbors within the block cutoff — a regime where the dense path itself suffers a far larger discontinuity from truncating a real neighbor, i.e. where sel is misconfigured.

DPA-2’s repformer block performs message passing (per-layer neighbor feature aggregation), so a graph-frozen .pt2 archive additionally embeds a with-comm AOTInductor artifact. Multi-rank LAMMPS runs dispatch to this artifact and drive an MPI ghost-atom exchange (border_op) once per repformer layer, instead of folding ghosts onto local owners as the non-message-passing (e.g. DPA-1) graph path does. .pt2 archives frozen with --lower-kind graph before this artifact was introduced do not carry it and must be re-frozen to support multi-rank inference.

Per-atom virial on the graph route uses a different (but equally valid) decomposition than the dense path: each edge’s full bond-virial contribution is assigned to its source (neighbor) atom, whereas the dense path’s autograd-based per-atom decomposition distributes message-passing contributions across atoms differently. For non-message-passing descriptors the two conventions coincide; for DPA-2 they differ elementwise. Only the total virial (summed over all atoms) is convention-independent between the two paths; per-atom virial values from the graph and dense routes should not be compared directly.

Multi-rank message-passing inference on the graph route requires every MPI rank to own or ghost at least one atom – a rank with zero atoms in both categories raises an error rather than silently desynchronizing the collective per-layer ghost exchange. For a graph-eligible DPA-2 descriptor, deepmd freezes through the graph route only: --lower-kind nlist is overridden to graph for such models, so a dense artifact is not selectable for the same model. Pick a domain decomposition that keeps every rank non-empty, or configure the model so it is not graph-eligible if you need the dense route.