Ctrl+K

nvalchemi.training.ForceMSELoss#

class nvalchemi.training.ForceMSELoss(*, target_key='forces', prediction_key='predicted_forces', normalize_by_atom_count=True, ignore_nonfinite=False, dtype_policy='strict')[source]#

Mean-squared-error loss on per-atom forces.

Forces enter this loss as per-atom vector quantities, unlike energy totals. Writing \(r_{ia\alpha} = \hat{F}_{ia\alpha} - F_{ia\alpha}\) for the residual of Cartesian component \(\alpha\) of atom \(a\) in graph \(i\), the normalize_by_atom_count flag selects how the squared component residuals are reduced across a mixed-size batch:

  • normalize_by_atom_count=True (default): each graph’s mean squared component error is averaged over graphs, so every graph has equal weight (a graph-balanced reduction):

    \[L = \frac{1}{B} \sum_{i=1}^{B} \frac{1}{3 N_i} \sum_{a=1}^{N_i} \sum_{\alpha=1}^{3} r_{ia\alpha}^2.\]

  • normalize_by_atom_count=False: one global mean over every valid force component, so a graph’s weight is proportional to its atom count:

    \[L = \frac{1}{3V} \sum_{i=1}^{B} \sum_{a=1}^{N_i} \sum_{\alpha=1}^{3} r_{ia\alpha}^2, \qquad 3V = \sum_{i=1}^{B} 3 N_i.\]

Dense force tensors use shape (V, 3). Padded force tensors use shape (B, V_max, 3) and ignore padding entries according to num_nodes_per_graph supplied either as (B,) counts or a (B, V_max) node mask. A hat denotes the prediction; see the module docstring for the shared notation.

Tensor Contract#

pred, targetForces | Float[torch.Tensor, “B V_max 3”]

Dense per-node forces of shape (V, 3) or padded per-graph forces of shape (B, V_max, 3). Shape validation requires exact equality.

batch_idxBatchIndices, optional

Required for dense (V, 3) forces when normalize_by_atom_count=True. Ignored for padded forces.

num_nodes_per_graphInteger[torch.Tensor, “B”] | Bool[torch.Tensor, “B V_max”], optional

Required for padded (B, V_max, 3) forces. May be explicit per-graph counts or a padded node-validity mask.

param target_key:

Target container key for the target tensor.

type target_key:

str, default “forces”

param prediction_key:

Prediction container key for the model output.

type prediction_key:

str, default “predicted_forces”

param normalize_by_atom_count:

Control the batch reduction for already-per-atom force residuals. True computes a graph-balanced mean by dividing each graph’s force-error sum by its valid component count before averaging over graphs. False computes one global elementwise mean over all valid force components.

type normalize_by_atom_count:

bool, default True

param ignore_nonfinite:

When True, target force components that are NaN or infinite are excluded from both loss value and gradient using torch.isfinite(). Intended for batches where some atoms/graphs lack force labels. Implemented with branch-free tensor ops for torch.compile compatibility. A graph whose entire force tensor is non-finite contributes 0.0 to the loss.

type ignore_nonfinite:

bool, default False

param dtype_policy:

How to handle prediction/target dtype mismatches before validation. "strict" raises; the other policies cast one tensor to match the other.

type dtype_policy:

{“strict”, “prediction_to_target”, “target_to_prediction”}, default “strict”

compute_residual(pred, target, valid)[source]#

Return squared force-component residuals, zeroing invalid entries.

Parameters:
Return type:

Tensor

extra_repr()[source]#

Human-readable hyperparameter summary for nn.Module’s repr.

Return type:

str

mask(pred, target, ctx, **kwargs)[source]#

Return component-level validity mask for dense or padded forces.

Parameters:
Return type:

Tensor

reduce(residual, valid, ctx, **kwargs)[source]#

Reduce force-component residuals to a scalar loss.

Parameters:
Return type:

Tensor

Parameters:

  • target_key (str)

  • prediction_key (str)

  • normalize_by_atom_count (bool)

  • ignore_nonfinite (bool)

  • dtype_policy (DTypePolicy)

On this page
Show Source