feat(stat): isolated-atom energy reference through preset_out_bias and vacuum_ref (#6022)
preset_out_bias documents that the bias of an assigned atom type is
set to the preset value. In change-by-statistic mode (fine-tuning, dp change-bias, change_bias_after_training) the preset was fed to the
least-squares fit as if it were a shift, so an assigned type ended at
pretrained_bias + preset, the residual fitted for the other types
subtracted natoms * preset instead of natoms * (preset - pretrained_bias), and every further call accumulated the preset again.
The four copies of the mode branch (the pt, pd and dpmodel atomic models
and the pt model wrapper) shared the defect.
This PR makes the preset the bias of an assigned output in every mode,
adds bundled tables of isolated-atom energies so that a table is named
instead of typed, and adds the fitting option vacuum_ref, which
references the network output of every atom to the output the same
network gives an isolated atom of its type. With both options the energy
of an atom without neighbors is exactly its preset bias, whatever the
network parameters are, so the dissociation limit of an energy model is
pinned to the isolated-atom energies of the reference calculation.
Changes
preset_out_bias

An output with at least one assigned type is fixed by the preset
alone. Every element that occurs in the training data must be assigned
(a missing element is an error), the assigned types take the preset
value in both set-by-statistic and change-by-statistic, a type
absent from the data gets a zero bias at initialization and keeps its
stored bias at fine-tuning, and no statistics are computed for, read
from or written to the cache for such an output; its output std keeps
its stored value. The types excluded by atom_exclude_types (the
virtual types of a spin model, for example) need no preset, and elements
outside the type_map are ignored. Outputs without a preset are fitted
as before.
Four forms are accepted: a dict keyed by element symbol, the name of a
bundled table, the path of a JSON file holding such a dict, and a list
with one entry per type of the type_map (null leaves a type
unassigned; the form for tensor outputs). A table is resolved once when
the input is processed, in single-task configurations, in every branch
of a multi-task configuration and for a preset written next to
model_dict, and its values are stored in the model, so a trained model
depends neither on the file nor on the bundled data.
Bundled tables live in deepmd/utils/preset_out_bias_tables.json, one
entry per name with its source and its values keyed by element
symbol, so further tables can be added without code changes. The entries
are the isolated-atom reference energies of the UMA training tasks
published with fairchem
(configs/uma/training_release/element_refs/iso_atom_elem_refs.yaml,
MIT license), in eV: omat24 (89 elements), omol25 (83, neutral
atoms), omc25 (94), odac25 (94) and oc20 (97). A bundled name
takes precedence over a file of the same name.
All preset handling lives in deepmd.utils.preset_out_bias
(normalization, table resolution, remapping, validation, per-type rows),
shared by the pt, pd and dpmodel atomic models; the per-backend
change_out_bias bodies reduce to the shared helper plus the fit of the
remaining outputs, and _store_out_stat writes bias rows without
touching the std. change_type_map remaps the preset together with the
stored bias; an output the model does not produce, a preset on a fitting
whose statistics do not distinguish atom types, and an assigned preset
on a dipole model are rejected.
The dpmodel factory (pt_expt, jax) and the pt linear and ZBL builders
forward the option; the pt_expt DPA4/SeZM builder accepts it.

vacuum_ref (fitting option, default off)

The atomic energy becomes E_i = bias(t_i) + f(x_i; c_i) - f(x_vac(t_i); c_i), where x_vac(t) is the descriptor of an isolated
atom of type t computed by the same descriptor with the current
parameters and c_i is the atom’s own conditioning (frame parameters,
atomic parameters, case embedding). Without frame or atomic parameters
the reference is one row per type; otherwise the reference rows follow
the atoms. The reference is evaluated at every training step, so it
follows the parameters; forces and virials are unchanged.
The reference atom is the neutral atom in its ground state: with
charge/spin conditioning it carries zero charge and the ground-state
multiplicity of the element, and with native spin a spin vector of one
Bohr magneton per unpaired electron (deepmd.utils.vacuum_reference,
from the electron-configuration table); charged, excited or differently
magnetized isolated atoms keep the deviation the network learns. The
condition tables are built once per type map at the atomic model (pt:
buffers of SeZMAtomicModel; dpmodel: numpy attributes of
DPAtomicModel, non-persistent buffers on pt_expt).
Routes: on pt_expt every graph-native model (DPA1, DPA2, DPA4, DPA4C)
carries one isolated node per type through the same forward as
single-node frames appended to the flat node axis
(append_isolated_frames); on the PyTorch backend the SeZM edge route
appends the reference nodes in forward_with_edges. The dense route
evaluates single-atom frames.
Freezing removes the reference atoms from every exported model: the
vacuum descriptor of every type is evaluated once on the export object
and folded into the fitting bias, or, for a fitting with frame or atomic
parameters whose reference varies between atoms, stored in the fitting
as a deployment constant (vacuum_table, kept out of checkpoints and
serialization) from which the fitting evaluates its references. The
pt_expt export folds inside _trace_and_export_impl, so dp freeze,
.pte/.pt2 conversion and change-bias on frozen artifacts all
resolve the reference on the object they export, and the archive keeps
the live model. A SeZM checkpoint in dens mode is rejected by the
freeze: the DeNS head serves training alone and is not exported.
The fused pt_expt fitting operator and the fused energy/force route
take the reference through the per-type bias (graph_fitting(fit, descriptor, atype, atom_bias)); a fitting whose reference varies
between atoms is served by the autograd route.
Scope: vacuum_ref is available for DPA4/SeZM on the PyTorch backend
and for the graph-native models on pt_expt. The TensorFlow and Paddle
fittings reject a serialized vacuum_ref: true model with
NotImplementedError instead of ignoring the option. atom_ener is
unchanged on every backend; a fitting rejects the two options together.
Documentation: the section “Isolated-atom energy reference” of
doc/model/train-energy.md, the argcheck entries preset_out_bias and
vacuum_ref, and the examples examples/water/dpa4/input_e0.json
(single-task) and examples/water/dpa4/input_multitask_e0.json
(per-branch presets).

Breaking changes

Reusing output-statistics caches requires a per-output preset record
matching the current constraints. Existing caches without that record
are recomputed in update mode; read-only users must regenerate the
output statistics once with stat_file_mode: update. This cache
validation change does not alter model or checkpoint serialization.
A partial preset is no longer fitted around: if an output assigns any
type, every element observed in the data must be assigned, otherwise the
statistics raise. Previously the unassigned observed types were fitted
with least squares around the preset.
With set-by-statistic, the types of an assigned output that do not
occur in the data get a zero bias instead of a fitted one, and the
output std of such an output keeps its initial value; preset outputs
never read or write the output-statistics cache.
DescrptSeZM.forward_with_edges takes vacuum_conditions: dict | None (the reference-atom inputs) and returns (descriptor, latent, vacuum); deepmd.pt_expt.kernels.graph_fitting.graph_fitting takes the
per-type bias explicitly.
The fitting serialization of every backend includes vacuum_ref, and
its @version is bumped (general fitting 4 to 5, polarizability 5 to 6,
property 6 to 7, population 4 to 5); older dictionaries without the key
load with vacuum_ref: false, and the TensorFlow and Paddle fittings
emit false and reject true behind the version check.
compute_stats_do_not_distinguish_types loses its unused
assigned_bias parameter.

Validation

Preset semantics: source/tests/common/test_preset_out_bias.py
(forms, bundled tables, resolution in single-task and multi-task
configurations, rows, excluded and unknown types), the pt, dpmodel and
pd test_atomic_model_global_stat.py (set → change → change with
an assigned output, the rejection of an unassigned observed element, the
output std under a full preset, a spin model with a dict preset, presets
that need no data, change_type_map), the jax end-to-end test and the
argcheck example inputs.
Reference identity E_i = bias + f(x_i) - f(x_vac) and “an isolated
atom gives exactly its bias”, in float64 to 1e-10:
source/tests/pt/model/test_fitting_vacuum_ref.py (the full (fparam, aparam, case_embd, aparam_as_mask) × mixed_types product,
torch.jit.script, default_fparam, fold and stored table, state-dict
round trip, change_type_map),
source/tests/common/dpmodel/test_fitting_invar_fitting.py,
test_fitting_call_graph.py, test_vacuum_ref_model.py (se_e2_a dense
and DPA1 graph routes),
source/tests/pt_expt/fitting/test_dpa4_ener.py.
source/tests/tf/test_fitting_vacuum_ref.py checks that the TensorFlow
fittings reject a serialized vacuum_ref and load a dictionary of the
previous version; the Paddle, dpmodel and PyTorch serialization tests do
the same.
Reference rows through the descriptors:
source/tests/pt/model/test_sezm_vacuum_ref.py (plain, charge/spin
conditioning with native spin, DeNS, forces unchanged, fold, fused
training kernels under AMP),
source/tests/pt/model/test_sezm_parallel.py (reference rows under the
LAMMPS-style communication path),
source/tests/pt_expt/model/test_dpa4_vacuum_ref.py (graph route with
charge/spin and native spin, a padded node axis, dense and graph routes
agree to 1e-12, .pt2 and .pte freezes with and without frame
parameters, the archive keeps the live model),
source/tests/pt_expt/model/test_fused_vacuum_ref.py (fused
energy/force route of DPA1 and compressed DPA4C equals autograd exactly,
autograd fallback with frame parameters),
source/tests/pt/model/test_sezm_vacuum_freeze.py (.pt2 freezes on
CPU and CUDA targets, with and without frame parameters, a dens
checkpoint is rejected).
End to end: the two example inputs train (single-task, multi-task,
two-process DDP); a model frozen with dp --pt freeze gives an isolated
O and H exactly their preset energies. Paddle is not installed here and
the TensorFlow operators are not built; the pd and tf tests mirror pt
and are verified by CI.



Summary by CodeRabbit

New Features

Added flexible per-element output-bias presets, including bundled
tables, JSON files, and element-mapped values.

Added isolated-atom (vacuum_ref) energy references for supported
energy models, including folding references during export.

Added vacuum-reference support for DPA4/SeZM and native-spin
workflows.

Added support for varying graph frame sizes and isolated reference
frames.

Bug Fixes

Preset biases now validate observed types and remain fixed during
statistical adjustments.

State-dependent statistics no longer incorrectly reuse cached results.

Improved model serialization and export consistency for bias and
vacuum-reference settings.