4.8. Descriptor DPA4 PyTorch#

Note

Supported backends: PyTorch PyTorch

DPA4 is the DeePMD-kit implementation of the SeZM (Smooth Equivariant Zone-bridging Model) architecture: an SO(3)-equivariant message-passing model for conservative interatomic potentials. The aliases DPA4, SeZM, and sezm all select the same implementation.

model.type: "dpa4" is a convenience scaffold that selects the SeZM descriptor and defaults to the dpa4_ener energy fitting network, so descriptor.type and fitting_net.type may be omitted for energy training. A new energy input then needs only the model type, type_map, and a few descriptor options.

Reference: DPA4 paper in the canonical citation guide.

4.8.1. Quick start#

DPA4 is a PyTorch-only model. Train it with the standard dp --pt workflow:

cd examples/water/dpa4
dp --pt train input.json

examples/water/dpa4/input.json is a complete, compact training input you can copy and adapt. See training energy models for the general training workflow shared by all energy models.

4.8.2. Overview#

DPA4/SeZM predicts atomic energies and obtains forces and virials by differentiating the energy, the same conservative formulation used by standard DeePMD energy models:

\[\mathbf{F}_i = -\frac{\partial E}{\partial \mathbf{r}_i}.\]

Internally the model keeps vector and higher-order angular (SO(3)-equivariant) information while building the descriptor, and only the final descriptor sent to the fitting network is a scalar. This separates geometric representation (equivariant message-passing layers that encode local environments) from energy prediction (the fitting network that maps scalar features to atomic energies). The architecture targets a favorable accuracy–cost trade-off; if you want the design details, see Architecture details at the end of this page.

4.8.3. Configuration#

4.8.3.1. Minimal input#

A minimal DPA4 model only needs the model type, type_map, and the neighbor list (sel, rcut). Every other option uses its documented default.

{
  "model": {
    "type": "dpa4",
    "type_map": [
      "O",
      "H"
    ],
    "descriptor": {
      "rcut": 6.0
    }
  }
}

DPA4/SeZM defaults to float32 (precision); double precision is unnecessary and not recommended (see Hardware selection).

Note

sel behaves differently from classic descriptors. On the conservative energy path it is only an initial neighbor-search capacity that grows on demand, so it never truncates the neighbor list and you do not need to size it to the true maximum neighbor count. The native spin scheme shares this energy path, so it grows on demand too; only the denoising (dens) path and the deepspin spin scheme cap the list at sum(sel). You can also set sel to auto or auto:factor to size it from the training data.

4.8.3.2. Main options#

Every descriptor option, with its default and full description, is listed in the argument reference. The options worth tuning first group into four accuracy–cost levers:

  • Angular width — the primary control. lmax with the per-block pyramid l_schedule (which overrides lmax and n_blocks), and the SO(2) order mmax / m_schedule. Cost grows quickly with lmax; a non-increasing l_schedule is often a good compromise.

  • Depthn_blocks, so2_layers, ffn_blocks.

  • Widthchannels, n_radial.

  • Aggregationn_focus, n_atten_head (0 falls back to a plain envelope-weighted scatter).

The neighbor list is set by rcut and sel, the initial node features by use_env_seed, and the energy head by the fitting neuron list ([0] is a direct projection). The quickest starting point is to copy examples/water/dpa4/input.json and adjust the levers above.

4.8.4. Training#

4.8.4.1. Energy training (default)#

The recommended objective is the standard conservative energy loss. The model predicts energies and forces are obtained by autograd:

{
  "loss": {
    "type": "ener"
  }
}

See training energy models for the general workflow.

4.8.4.2. Property training#

DPA4/SeZM can also train invariant structure properties by selecting the standard property fitting network. Set fitting_net.type to property, provide the property name used by the data file, and use the property loss:

{
  "model": {
    "type": "dpa4",
    "fitting_net": {
      "type": "property",
      "property_name": "band_prop",
      "task_dim": 3,
      "intensive": true
    }
  },
  "loss": {
    "type": "property"
  }
}

The property label follows the usual DeePMD property-data convention: each system provides band_prop.npy when property_name is band_prop. Property fitting is not a water task, so the complete DPA4/SeZM property input lives with the property dataset: see examples/property/train/input_dpa4.json, which trains on the QM9 property subset in examples/property/data.

4.8.4.3. Direct-force denoising (dens, experimental)#

DPA4/SeZM has an experimental direct-force denoising head:

{
  "loss": {
    "type": "dens"
  }
}

Use dens only when the direct-force denoising head is required; it is not the default training path. See examples/water/dpa4/input_dens.json for an example.

4.8.4.4. Spin#

DPA4/SeZM supports the DeePMD-kit spin convention through the standard model.spin block. Two schemes are available, selected by model.spin.scheme:

  • native — the per-atom spin vector enters the descriptor as an equivariant feature, and the magnetic force is the spin gradient of the energy. No virtual atoms are introduced, so the neighbor list and type map keep their real-system sizes.

  • deepspin (default) — the classical DeepSpin representation, in which each magnetic atom is paired with a virtual atom displaced along its spin. It is the default, so a model.spin block without an explicit scheme reproduces the classical behaviour shared with every non-SeZM spin model.

Both schemes train against the conservative ener_spin loss, share the same spin / force_mag data convention, and are not combined with the dens mode. Because the dataset and loss are identical, switching scheme does not require any change to the data or the loss block. Complete inputs are in examples/spin/dpa4/: input.json for the native scheme and input-deepspin.json for the deepspin scheme. See training spin energy models for the general workflow.

use_spin accepts a per-type boolean list or, to avoid enumerating a large type_map, the list of magnetic species as type indices or element symbols (e.g. "use_spin": ["Fe"]), expanded against type_map. For the native scheme, model.spin.allow_missing_label additionally admits training systems that lack a spin data file by filling their spin with zeros; since a zero spin reduces the native descriptor to its spin-free form, a model can be trained on a mixture of spin-labelled and spin-free systems, or pretrained on spin-free data and fine-tuned on spin-labelled data without changing its type map or parameters.

4.8.4.4.1. Native scheme#

{
  "model": {
    "type": "dpa4",
    "type_map": [
      "Ni",
      "O"
    ],
    "spin": {
      "use_spin": [
        true,
        false
      ],
      "scheme": "native"
    },
    "descriptor": {
      "rcut": 6.0
    }
  }
}

use_spin marks which atom types carry spin. Both the conservative force and the magnetic force come from a single energy gradient:

\[\mathbf{F}_i = -\frac{\partial E}{\partial \mathbf{r}_i}, \qquad \mathbf{F}^{\mathrm{mag}}_i = -\frac{\partial E}{\partial \mathbf{s}_i},\]

where \(\mathbf{s}_i\) is the input spin vector. The model does not rescale the spin internally, so \(\mathbf{s}_i\) keeps the dataset’s spin convention and the magnetic force is reported in the matching units of force_mag – there is no virtual_scale factor in this scheme. The magnetic force is reported on the magnetic atom types only, matching the force_mag label. The native scheme relies on the descriptor’s angular degrees to represent the spin direction, so it requires lmax >= 1 (the default).

4.8.4.4.2. DeepSpin virtual-atom scheme (default)#

{
  "model": {
    "type": "dpa4",
    "type_map": [
      "Ni",
      "O"
    ],
    "spin": {
      "use_spin": [
        true,
        false
      ],
      "virtual_scale": [
        0.314
      ],
      "scheme": "deepspin"
    },
    "descriptor": {
      "sel": 120,
      "rcut": 6.0
    }
  }
}

The deepspin scheme augments each magnetic atom with a virtual atom at coord + spin * virtual_scale, which doubles the internal neighbor capacity and type map. virtual_scale is required by this scheme and ignored by the native scheme.

4.8.4.5. Multi-task / shared fitting#

DPA4/SeZM supports shared-fitting multitask training. With case_film_embd: true, the case vector modulates the fitting network instead of being concatenated to the descriptor, which keeps the descriptor case-independent while letting the energy map depend on the task branch. See multi-task training for the workflow and examples/water/dpa4/input_multitask.json for an example.

4.8.4.6. LoRA fine-tuning#

DPA4/SeZM supports LoRA adapters on its SO(3) and SO(2) linear layers, intended for single-task fine-tuning:

{
  "model": {
    "type": "dpa4",
    "lora": {
      "rank": 16,
      "alpha": 16.0
    }
  }
}

Fine-tune from a checkpoint:

dp --pt train lora_ft.json --finetune pretrained.pt

Best checkpoints fold the LoRA deltas back into the base weights, producing plain DPA4/SeZM checkpoints suitable for deployment. See examples/water/dpa4/lora_ft.json.

4.8.5. Zone bridging (ZBL)#

DPA4/SeZM can add an analytical short-range repulsion (typically ZBL) to the learned energy in a protected region:

\[E_i = E_i^{\mathrm{DPA4/SeZM}} + E_i^{\mathrm{ZBL}}.\]

Below bridging_r_inner the distance seen by the descriptor is clamped, with a smooth transition back to the true distance up to bridging_r_outer; a source gate additionally blocks the learned model from leaking information about the frozen short-range pairs. The recommended way to enable it is the concise form, set directly on the dpa4 model:

{
  "model": {
    "type": "dpa4",
    "bridging_method": "zbl",
    "bridging_r_inner": 0.5,
    "bridging_r_outer": 0.8
  }
}

When ZBL bridging is enabled, set training.training_data.min_pair_dist to the same value as bridging_r_inner so frames with shorter atom pairs are excluded from training. See examples/water/dpa4/input-zbl.json for a complete example.

Note

Output-bias statistics and bridging: the model energy is E = E_model + E_bias, and the ZBL term belongs to E_model. The set-by-statistic bias mode (initial statistics, finetune with a random fitting, dp change-bias --mode set) fits E_bias to the raw data labels and by definition ignores E_model — the analytical ZBL contribution included. For a self-consistent calibration of a bridged model use change-by-statistic, which subtracts the complete bridged prediction. See change-bias for the precise definitions.

Internally, a bridged model is a linear composition: the learned model plus the analytical inner_potential term, summed by linear_ener. The concise form above expands to exactly this equivalent explicit form:

{
  "model": {
    "type": "linear_ener",
    "weights": "sum",
    "type_map": [
      "O",
      "H"
    ],
    "models": [
      {
        "type": "dpa4",
        "descriptor": {
          "...": "..."
        },
        "fitting_net": {
          "...": "..."
        }
      },
      {
        "type": "inner_potential",
        "mode": "zbl",
        "r_inner": 0.5,
        "r_outer": 0.8
      }
    ]
  }
}

Both spellings build the same model (one shared normalizer defines the equivalence). The explicit form exposes the composition machinery directly: use it when you combine models beyond the standard bridged pair. In either form, the composition derives the learned descriptor’s clamping window from the analytical term, so the radii are written once.

4.8.6. Performance and precision#

4.8.6.1. Training-time settings#

Three options control training precision and the compiled path:

  • use_amp — bf16 automatic mixed precision on CUDA. Reduces memory and often improves throughput with no expected accuracy loss. Recommended on GPUs with native bf16 (NVIDIA Ampere and newer, e.g. A100 / RTX 30-series); set it off on GPUs without native bf16 to avoid runtime errors and conversion overhead.

  • enable_tf32 — TF32 matmul precision for CUDA training forwards (independent of use_compile, and separate from the inference TF32 control below).

  • use_compile — experimental torch.compile training path. Useful for force-loss training (higher-order derivatives through the energy gradient) and can speed training markedly on supported setups.

4.8.6.2. Inference and deployment settings#

Inference behavior is controlled by environment variables, each with an equivalent input-file option used during training validation:

Environment variable

Input-file option

Default

Effect

DP_COMPILE_INFER

validating.compiled_infer

off

Use the compile path for evaluation/inference. Same torch==2.11 / CUDA ≥ 12.6 requirements as model.use_compile.

DP_TF32_INFER

validating.tf32_infer

0 (highest)

float32 matmul precision for inference: 0 highest, 1 high, 2 medium. Higher values improve throughput but make the potential energy surface less smooth.

DP_AMP_INFER

validating.amp_infer

off

bf16 autocast inside the descriptor interaction blocks for inference, independently of descriptor.use_amp. Training AMP remains controlled by descriptor.use_amp. Usually keeps aggregate MAE similar but can make the potential energy surface less smooth.

DP_TRITON_INFER

0

Triton inference kernel level 0-3 (CUDA eval only, compatible with DP_COMPILE_INFER). 1: universal fused kernels, numerically equivalent to the dense path with full float32 accumulation. 2: adds the table-configured fused SO(2) value path and edge-block backward kernels (still exact float32). 3: additionally runs the SO(2) mixing stack on fp16 tensor cores with split compensation — roughly float32-level accuracy (maximum force deviation about 4e-6 eV/Å on a 4-thousand-atom system) at a substantial speedup; only shapes validated by the tuning sweep are affected. Levels 2 and 3 read launch tables tuned per GPU model (H20 ships built in); on other GPUs the kernels fall back to conservative configurations, and dp --pt freeze tunes the missing entries automatically on the local GPU before exporting (a one-off sweep of a few minutes, baked into the .pt2).

Accepted boolean values for the other switches are 1/true/yes/on and 0/false/no/off; DP_TRITON_INFER accepts only the numeric levels. Shell exports take precedence over the input-file options and over values written in the input; they are read when the model is constructed and changing them afterward has no effect.

For molecular dynamics and other workflows sensitive to the smoothness of the potential energy surface, keep DP_TF32_INFER=0 and DP_AMP_INFER=0. DP_AMP_INFER can coexist with DP_TF32_INFER, but bf16 autocast dominates the eligible operations it covers, so TF32 usually adds little extra throughput there. DP_TRITON_INFER levels 1 and 2 retain full float32 accumulation regardless of the precision policy and are therefore safe for those workflows; level 3 perturbs forces at the 2^-22 rounding scale (three orders of magnitude finer than TF32) and is the recommended fast setting once validated for the target system.

Important

Set these variables before running dp --pt freeze. The exported .pt2 is an AOTInductor artifact, so the SO(2) rotation branch (DP_TRITON_INFER), the matmul precision (DP_TF32_INFER), and inference AMP (DP_AMP_INFER) are captured into the graph at export time and are not re-evaluated when the .pt2 is later loaded by ASE or LAMMPS. A frozen .pt2 runs a forward-only package, so training-time memory-saving switches do not apply to it.

4.8.6.3. Hardware selection#

DPA4/SeZM is designed for fp32 training and inference, so prefer GPUs with high fp32 throughput and native bf16 support rather than strong fp64 performance. Because bf16 AMP substantially reduces the activation memory footprint, very large device memory is usually less important than fp32 FLOPS and bf16 support once the target system and batch size fit.

4.8.7. Export and running in LAMMPS#

4.8.7.1. Freeze to .pt2#

DPA4/SeZM checkpoints use the PyTorch .pt2 (AOTInductor) export path; the ordinary TorchScript freeze path is not used. Run the standard freeze command:

dp --pt freeze -c model.ckpt.pt -o frozen_model

The PyTorch backend detects DPA4/SeZM and writes frozen_model.pt2.

4.8.7.2. Single GPU#

Use the frozen .pt2 with the deepmd pair style. A small example is in examples/water/dpa4/lmp/.

pair_style deepmd frozen_model.pt2
pair_coeff * * O H

4.8.7.3. Multi-GPU (MPI) inference#

Important

Multi-rank support depends on the model, not on a dense/graph choice. DPA4/SeZM reads ghost-neighbour features at every interaction block, so a multi-rank archive must embed a with-comm AOTInductor artifact (model/extra/forward_lower_with_comm.pt2) for the cross-rank exchange. Two different export routes produce one:

  • PT (dp --pt freeze, the default on this page). DPA4/SeZM is detected and exported through freeze_sezm_to_pt2, which uses the compact edge_vec ABI (coord, atype, edge_index, edge_vec, edge_scatter_index, edge_mask), building the neighbour topology on the C++ side. This is not a padded dense neighbour-list export. A plain energy model reports supports_edge_parallel() == True, so the archive carries the with-comm artifact (has_comm_artifact=true) and supports multi-rank LAMMPS out of the box — no extra freeze options.

  • pt_expt (dp --pt_expt freeze). Graph-capable models export through the NeighborGraph ABI (see Graph-native inference route (pt_expt) below), which likewise embeds a with-comm artifact and supports multi-rank LAMMPS.

The two are distinct ABIs consumed by different C++ paths; neither is a dense route, and --lower-kind selects between the pt_expt lowers only.

Which models lose multi-rank, on either route:

  • The deepspin (virtual-atom) spin scheme, which overrides the export ABI to nlist because it expands virtual atoms inside the graph.

ZBL zone bridging is multi-rank capable (including combined with native spin). The Source Freeze Propagation gate folds each node’s full outgoing-edge set, which no single rank observes for ghost owners; the with-comm artifact completes the gate’s per-node [log eta, zero count] partials across ranks with one reverse-accumulate plus one forward-broadcast border exchange of an (N, 2) tensor per forward pass. The payload is narrow, but the cost is not bandwidth alone: both kernels end in an MPI_Barrier, and the force graph differentiates through them, so their transposes add a matching pair of round trips to the force evaluation. On small systems or at high rank counts those synchronizations can dominate; the exchange has not been benchmarked across system sizes and rank counts. Native spin (scheme: "native") and charge/spin conditioning likewise support multi-rank: native spin reuses the edge_vec interface on PT and the NeighborGraph lower on pt_expt. The remainder of this subsection describes the multi-GPU launch recipe.

The exported .pt2 runs across multiple GPUs in LAMMPS using MPI domain decomposition, with the same .pt2 file serving both single- and multi-GPU runs and no extra freeze options needed beyond --lower-kind graph. The launch recipe:

CUDA_VISIBLE_DEVICES=0,1,2,3 mpirun -np 4 lmp -in in.lammps

Each MPI rank uses at most one GPU, so CUDA_VISIBLE_DEVICES must list every GPU the run may use. DPA4/SeZM exchanges neighbor information across the domain boundary, so the LAMMPS atom map must be enabled:

atom_modify map yes
pair_style deepmd frozen_model.pt2
pair_coeff * * O H

Two settings improve multi-GPU runs:

  • For fast GPU-to-GPU exchange, build the C++ interface against a CUDA-Aware MPI library; otherwise the cross-rank exchange falls back to a slower CPU path.

  • Use a non-zero neighbor skin, e.g. neighbor 2.0 bin, to keep per-step GPU memory stable. A zero skin rebuilds the neighbor list every step and can substantially increase memory use.

4.8.8. Graph-native inference route (pt_expt)#

In the pt_expt backend, a plain-energy DPA4/SeZM descriptor – or one configured with spin.scheme: native (see Native spin (magnetic) below) – can be frozen through a NeighborGraph-native inference path instead of the legacy dense neighbor-list path. Frame-level charge/spin conditioning (add_chg_spin_ebd) and ZBL zone bridging are graph-eligible too; only the deepspin virtual-atom spin scheme remains dense-only:

dp --pt_expt freeze -o model.pt2 --lower-kind graph

As with DPA-1’s and DPA-2’s graph paths (see Difference among different backends and the “Graph-native inference route (pt_expt)” section of DPA-2’s documentation), the graph route is a sel-free, carry-all builder: it considers every neighbor within rcut rather than the capacity fixed by sel. For the default conservative-energy path, sel is already only an initial search capacity that grows on demand (see the note in Minimal input), so graph and dense results agree there down to floating-point noise. Where sel does act as a hard cap – the dens denoising path, or the deepspin virtual-atom spin scheme, both of which size the neighbor list to sum(sel) – the graph route’s larger neighbor set can diverge from the dense route, the same deliberate divergence documented for DPA-1/DPA-2.

A DPA4/SeZM descriptor configured with deepspin-scheme spin is not graph-eligible and always runs the dense route regardless of --lower-kind; --lower-kind graph on such a model raises an error at freeze time instead of exporting a silently-dense-only artifact. native-scheme spin and ZBL zone bridging are the opposite case: they have no dense route at all – the analytical bridging term has no dense injection site – so they are always graph-frozen, --lower-kind notwithstanding. Charge/spin conditioning rides the graph lower and constrains neither. Note that a graph-capable model is always frozen to the graph lower in any case, since the dense lower is deprecated in the pt_expt backend. See Native spin (magnetic) below.

Unlike the dense route (see Multi-GPU (MPI) inference above), a graph-frozen .pt2 embeds a with-comm AOTInductor artifact and supports multi-rank LAMMPS: each block’s cross-rank ghost-feature exchange runs through the border_op MPI path once per interaction block, the same mechanism used by DPA-2’s graph route (see the “Graph-native inference route (pt_expt)” section of DPA-2’s documentation). As on DPA-2, multi-rank 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 aborts the run collectively rather than silently desynchronizing the per-block exchange. Pick a domain decomposition that keeps every rank non-empty, or use the dense route, which has no such restriction (but is single-rank only, as noted above).

Native-spin graph .pt2 archives participate too: they carry the with-comm artifact and support multi-rank LAMMPS – see Native spin (magnetic) below.

4.8.8.1. Native spin (magnetic)#

spin.scheme: native (see Spin above for the general convention and the native-scheme JSON example) is graph-only: a native-spin DPA4/SeZM model has no dense (nlist) lower at all, so it is always frozen through the NeighborGraph path regardless of --lower-kind, and always runs through it at inference time. The deepspin virtual-atom scheme is unaffected and keeps using the dense route described in Spin and Multi-GPU (MPI) inference.

  • Native scheme only. deepspin-scheme spin (and the general spin virtual-atom model outside DPA4/SeZM) is dense-only; only scheme: native is graph-eligible. dp --pt_expt freeze --lower-kind graph on a deepspin-scheme model raises an error at freeze time, per the dense/graph eligibility rule above.

  • Graph route only, no dense fallback. A native-spin descriptor has only the graph lower. This is not a spin-specific restriction on the CLI: every graph-capable DPA4/SeZM model, plain-energy included, is frozen to the graph lower. --lower-kind accepts only nlist (the default) and graph, and freeze() overrides any non-graph request to graph whenever the model is graph-lower capable, logging a warning. A dense artifact is therefore not selectable through this entry point for these models.

  • Multi-rank capable. The frozen archive embeds the with-comm artifact (has_comm_artifact is true), so multi-rank LAMMPS works exactly as described in Multi-GPU (MPI) inference: the per-block ghost node features ride border_op, and ghost spins arrive through the LAMMPS sp forward communication – spin itself needs no extra cross-rank exchange.

  • The magnetic force is a second energy gradient. As in the general native-scheme convention (see Spin above), force_mag = -\partial E/\partial\mathbf{s}, computed by pt_expt as a second torch.autograd.grad call alongside the ordinary force = -\partial E/\partial\mathbf{r} call; both are real autograd outputs, not placeholders. The dpmodel (NumPy) backend is energy-only for native-spin models – it has no autograd, so force/force_mag/virial are None placeholders there, exactly as for the plain-energy dpmodel route.

No native-spin combination restrictions remain on the graph route: multi-rank inference, charge-spin FiLM conditioning (add_chg_spin_ebd), and ZBL zone bridging (bridging_method: ZBL) all combine freely with spin.scheme: native.

4.8.9. Embedding extraction#

A trained DPA4/SeZM model can export learned representations for downstream analysis with dp embed. A single forward pass (no force or virial) produces three embeddings per system:

  • descriptor: per-atom local-environment representation, shape (nframes, natoms, dim_descriptor).

  • atomic_feature: per-atom activation after the last fitting hidden layer, shape (nframes, natoms, dim_hidden).

  • structural_feature: whole-structure summary obtained by summing atomic_feature over atoms, shape (nframes, dim_hidden).

dp embed -m model.ckpt.pt -s /path/to/system -o embedding.hdf5

The results are written to a single HDF5 file in which each system is a group holding the three float32 datasets above:

import h5py

with h5py.File("embedding.hdf5", "r") as f:
    type_map = f.attrs["type_map"]
    group = f[next(iter(f.keys()))]
    descriptor = group["descriptor"][:]
    atomic_feature = group["atomic_feature"][:]
    structural_feature = group["structural_feature"][:]

This command operates on the training checkpoint (.pt), not the frozen .pt2, and honors both DP_COMPILE_INFER and DP_TRITON_INFER. See model embeddings for the full description.

4.8.10. Data format#

DPA4/SeZM uses the standard DeePMD-kit data format and also supports the mixed-type data format, which is convenient for datasets that mix many element combinations (and is the usual choice for multi-task training). Keep the type_map order consistent across the dataset, the input file, and any downstream pair_coeff mapping.

4.8.11. Architecture details#

Optional background on how the descriptor works, linking each part to the options that control it. Skip it unless you are tuning those options.

4.8.11.1. Equivariant representation and the l = 0 read-out#

DPA4/SeZM stores intermediate features as SO(3)-equivariant coefficients: a feature block of maximum degree lmax holds all degrees l = 0, …, lmax, each with 2l + 1 angular components (controlled by lmax / l_schedule).

For each frame the model first builds a local neighbor graph within rcut. Each edge stores the displacement vector, smooth cutoff weights, radial basis features, and the rotation between the global frame and an edge-aligned local frame; these are built once and reused by all blocks. One interaction block then (1) gathers source-atom features on each edge, (2) rotates them into the edge-local frame, (3) applies an SO(2)-equivariant convolution on the retained angular orders, (4) rotates the messages back, (5) aggregates them at destination atoms with envelope or attention weights, and (6) updates atom features with an equivariant feed-forward block.

Working in the edge-local frame turns rotations around the edge axis into SO(2) operations, so the cost scales with lmax instead of cubically. The SO(2) convolution retains orders |m| mmax (or m_schedule). After the last block, only the l = 0 scalar channels are read out and passed to the fitting network:

\[\mathcal{D}_i = \mathrm{Scalar}\left(\mathbf{h}_i^{(L)}\right).\]

4.8.11.2. Radial basis and smooth cutoff#

Every edge uses a radial basis (basis_type, with n_radial functions) multiplied by a smooth envelope whose value and first three derivatives vanish at rcut. This smoothness matters for MD because nonsmooth descriptor cutoffs would be inherited by the force derivatives. The two env_exp exponents control the radial-basis envelope and the message-passing edge weights respectively; larger values keep an envelope closer to one for more of the cutoff range.

4.8.11.3. Attention and focus streams#

Messages are aggregated either by envelope-weighted scatter or by attention (n_atten_head > 0). With attention, the cutoff envelope participates in the softmax normalization, so edges near rcut are smoothly suppressed in both the numerator and the denominator. The SO(2) convolution can also use multiple n_focus streams that process the same edge geometry in parallel and are combined by scalar weights, adding capacity while preserving equivariance.

4.8.11.4. Grid nonlinearities#

Several branches can use sphere-grid (S2) or SO(3) Wigner-D grid nonlinearities. The main switches are s2_activation (S2-grid nonlinearity for the SO(2) and/or FFN branch), ffn_so3_grid (SO(3) grid in the block-internal FFN), lebedev_quadrature (Lebedev rules for enabled S2 branches), and grid_mlp / grid_branch (point-wise polynomial MLP or scalar-routed branch mixer per grid path). These trade expressiveness for cost; the final l = 0 output remains a scalar.

4.8.11.5. Environment-seeded initial features#

With use_env_seed enabled, the initial node state is seeded from the local environment: a DeePMD-style environment matrix produces FiLM-like scale and shift values for the first scalar features, and the geometric initial embedding is enabled when non-scalar degrees are present. With it disabled, the initial state contains only atom-local scalar features, which keeps a one-block model closed over the one-hop neighbor shell.

4.8.12. Limitations#

  • DPA4/SeZM is implemented for the PyTorch backend only.

  • Export uses .pt2 (AOTInductor); the TorchScript freeze path is not used.

  • Model compression is not supported.

  • Multi-rank (multi-GPU/MPI) LAMMPS inference works on both export routes: the PT edge_vec archive and the pt_expt NeighborGraph archive each embed a with-comm artifact. ZBL zone bridging (and its native-spin combination) participates: the Source Freeze Propagation gate’s per-node partials are completed across ranks by one reverse-accumulate plus one forward-broadcast border exchange. See Multi-GPU (MPI) inference.

  • The pt_expt graph-native inference route is unavailable only for deepspin-scheme spin, which stays on the dense route. Charge/spin conditioning, ZBL zone bridging and native-scheme spin are all graph-eligible.

  • spin.scheme: native is graph-only (it has no dense route) and supports multi-rank LAMMPS: ghost node features ride border_op per interaction block and ghost spins arrive through the LAMMPS sp forward-comm. Charge-spin FiLM conditioning and ZBL zone bridging both combine with it, multi-rank included. See Native spin (magnetic).

4.8.13. Citation#

If you use DPA4/SeZM, cite the DPA4 paper from the canonical citation guide. Its BibTeX record is maintained in CITATIONS.bib.