nvalchemi-loss-api

nvalchemi Loss API

Overview

Loss functions are torch.nn.Module subclasses rooted at BaseLossFunction. Each leaf consumes (pred, target, **kwargs) and returns a scalar. ComposedLossFunction routes keyed prediction/target mappings to leaves, applies per-component weights (float or LossWeightSchedule), and returns a ComposedLossOutput TypedDict.

from nvalchemi.training import (
    BaseLossFunction,
    ComposedLossFunction,
    ReductionContext,
    EnergyMSELoss,
    EnergyMAELoss,
    ForceMSELoss,
    ForceL2NormLoss,
    StressMSELoss,
)

---

Built-in losses

Choose losses by the training signal you want:

• EnergyMSELoss: default for smooth energy regression when larger errors should dominate early training; combine with per_atom=True when system sizes vary.
• EnergyMAELoss: more robust to outlier energies and often useful for reporting or late-stage fitting when median absolute accuracy matters.
• EnergyHuberLoss: compromise between MSE and MAE; use when energy labels have occasional noisy outliers but small residuals should remain smooth.
• ForceMSELoss: default force objective; component-wise squared residuals give strong gradients for geometry-sensitive fitting.
• ForceL2NormLoss: use when vector direction/magnitude per atom is the desired error signal rather than independent xyz components.
• ForceHuberLoss: robust force fitting when some force labels are noisy or contain rare large residuals.
• StressMSELoss / StressHuberLoss: add only when stress labels are reliable and the model is configured to produce stresses.

Composition sugar:

loss_fn = 1.0 * EnergyMSELoss() + 10.0 * ForceMSELoss() + 0.1 * StressMSELoss()
out = loss_fn(predictions, targets, step=step, epoch=epoch, batch=batch)
out["total_loss"].backward()

Graph metadata: losses that need graph structure (per_atom=True, normalize_by_atom_count=True, or padded layouts) accept batch= (pulls batch_idx, num_graphs, num_nodes_per_graph automatically) or explicit kwargs.