Losses — Training Terms#

Composable, tensor-first loss functions for MLIP training.

See also

User guide: Losses — conceptual overview, usage patterns, and how to write your own loss term.

A typical training loss is a composition of tensor-first leaf losses. The composition routes prediction/target mappings into each leaf, applies the configured weights, and returns a structured output whose total_loss is the scalar used for backpropagation:

from nvalchemi.training import ComposedLossFunction, EnergyMSELoss, ForceMSELoss

loss_fn = ComposedLossFunction(
    components=[
        EnergyMSELoss(),
        ForceMSELoss(),
    ],
    weights=[1.0, 10.0],
)

out = loss_fn(predictions, targets, batch=batch, step=step, epoch=epoch)
out["total_loss"].backward()

for name, value in out["per_component_unweighted"].items():
    logger.info("%s raw loss: %s", name, value.detach())

Dtype alignment#

Leaf losses default to dtype_policy="strict", which preserves prediction and target tensors and raises on dtype mismatch during validation. Built-in leaves also accept "prediction_to_target" and "target_to_prediction" to cast one tensor before validation. ComposedLossFunction(dtype_policy=...) provides the same policy as a call-time default for strict leaves without mutating reusable component instances. For compositions built with operator sugar, set loss_fn.dtype_policy after construction.

Leaf and composition#

Leaf losses subclass BaseLossFunction; compositions use ComposedLossFunction and return a ComposedLossOutput.

`BaseLossFunction`	Abstract `torch.nn.Module` base for ALCHEMI loss functions.
`ReductionContext`	Lightweight metadata bag flowing through the loss template pipeline.
`ComposedLossFunction`	Weighted sum of `BaseLossFunction` components.
`ComposedLossOutput`	Output returned by `ComposedLossFunction`.
`LossWeightSchedule`	Runtime-checkable protocol for loss-weight schedules.

Concrete losses#

Built-in leaf losses for common quantum-chemistry targets.

`EnergyMSELoss`	Mean-squared-error loss on per-graph total energy.
`EnergyMAELoss`	Mean-absolute-error loss for per-graph energy targets.
`EnergyHuberLoss`	Huber loss on total energy or energy per atom.
`ForceMSELoss`	Mean-squared-error loss on per-atom forces.
`ForceHuberLoss`	Huber loss on per-component force residuals.
`ForceL2NormLoss`	Mean per-atom force-vector L2 loss.
`StressMSELoss`	Mean-squared-error loss on the per-graph stress tensor.
`StressHuberLoss`	Huber loss on per-graph stress tensors.

Weight schedules#

Pydantic frozen models satisfying LossWeightSchedule. Custom schedules may also satisfy the protocol directly. For strategy checkpoint round-trips, implement to_spec() returning a BaseSpec. The built-in Pydantic schedule base provides this method from model_dump().

`ConstantWeight`	Time-invariant loss weight that returns `value` at every update.
`LinearWeight`	Loss weight that ramps linearly from `start` to `end`.
`CosineWeight`	Loss weight that eases from `start` to `end` on a half-cosine curve.
`PiecewiseWeight`	Step-function loss weight that switches value at fixed boundaries.

Reduction helpers#

Per-graph reduction helpers — scatter reductions (V ... → B ...) and matrix reductions (B ... m n → B ...) — importable for use in custom losses.

`per_graph_sum`	Sum per-node values into per-graph values via `scatter_add_`.
`per_graph_mean`	Mean of per-node values across each graph.
`frobenius_mse`	Per-graph Frobenius MSE over the trailing two matrix dims.