Force Loss Computation in DeePMD-kit#

Overview#

Force loss is computed inside EnergyStdLoss (PyTorch backend, deepmd/pt/loss/ener.py) and its analogues across other backends (deepmd/dpmodel/loss/ener.py, deepmd/pd/loss/ener.py, deepmd/tf/loss/ener.py). The class implements a forward() method that runs the model, computes prefactor-weighted losses for energy, force, virial, atomic energy, and generalized forces, and returns them for gradient-based optimization.

Force Difference and Tensor Shape Handling#

The baseline force loss path computes the element-wise difference and flattens it to a 1-D vector for mean-squared (MSE) or L1 (MAE) reduction 1:

diff_f = (force_label - force_pred).reshape(-1) # shape: (N_atoms * 3,)

relative_f normalization divides this flattened difference by the per-atom label norms, requiring a reshape back to (-1, 3) before the norm and then back to (-1,) 2:

force_label_3 = force_label.reshape(-1, 3)
norm_f = force_label_3.norm(dim=1, keepdim=True) + self.relative_f
diff_f_3 = diff_f.reshape(-1, 3) / norm_f
diff_f = diff_f_3.reshape(-1)

keepdim=True on the per-row norm ensures the division broadcasts correctly across all 3 Cartesian components without a separate expand step.

f_use_norm: Vector-Norm-Based Force Loss#

PR #5294 introduced two companion constructor parameters to EnergyLoss.__init__:

Validation (constructor)#

f_use_norm=True is rejected at construction time unless use_huber=True or loss_func="mae" 3:

if self.f_use_norm and not (self.use_huber or self.loss_func == "mae"):
    raise RuntimeError("f_use_norm can only be True when use_huber or loss_func='mae'.")

This guards against silently applying the norm path to the standard MSE branch where it has no well-defined semantic.

Shape flow when f_use_norm=True#

When enabled, force tensors are reshaped to (-1, 3) so that torch.linalg.vector_norm operates along dim=1 to produce a per-atom scalar norm tensor of shape (-1, 1). The norm is compared against a zero tensor (the target is to drive ‖F_pred − F_label‖₂ → 0) 3:

With loss_func="mae":

force_diff_3 = (force_label - force_pred).reshape(-1, 3)
l1_force_loss = torch.linalg.vector_norm(force_diff_3, ord=2, dim=1, keepdim=True).mean()

With use_huber=True:

force_diff_norm = torch.linalg.vector_norm(
    (force_label - force_pred).reshape(-1, 3), ord=2, dim=1, keepdim=True
)
l_huber_loss = custom_huber_loss(force_diff_norm, torch.zeros_like(force_diff_norm), delta=self.huber_delta)

keepdim=True preserves the (-1, 1) shape so that torch.zeros_like(force_diff_norm) and custom_huber_loss operate element-wise on identically-shaped tensors with no implicit broadcasting issues.

Metric names#

Pre-existing Norm Path: relative_f#

relative_f (a float | None, not a bool) is an older, independent normalization that divides the force error by the magnitude of the label force plus an offset, normalizing each component 2. It is incompatible with Huber loss — that combination raises a RuntimeError at construction 4. Unlike f_use_norm, relative_f operates on the component-level diff after normalization rather than replacing the per-component metric with a per-atom vector norm.

Serialization#

Both loss_func and f_use_norm are persisted in the serialize() output 3 so that checkpointed models can be faithfully restored. Deserialization uses check_version_compatibility 5 to handle version migrations gracefully.

Key Files#