xnn.common.train.losses.weighted_loss#

xnn.common.train.losses.weighted_loss(pred, data, energy_weight, force_weight, stress_weight, *, dipole_weight=0.0, polarizability_weight=0.0, charge_weight=0.0, les_dipole_weight=0.0, les_charge_weight=0.0, huber_delta=0.0, huber_delta_energy=None, huber_delta_forces=None, huber_delta_stress=None, head_weights=None, head_names=None, counts=None)[source]#

Compute a weighted sum of energy, force and stress loss terms.

Each term is included only when the target is present in data and its weight is positive (and, for forces and stress, when the corresponding key is present in pred). The individual terms are:

Energy: per-atom energy error, i.e. the difference between the predicted and target total energy divided by the number of atoms in each structure, then averaged over structures. Dividing by the atom count makes the term size-extensive and well-scaled across structures of different sizes.
Forces: per-atom force error, averaged over every atom and Cartesian component in the batch.
Stress: error between predicted and target stress tensors.
Dipole and polarizability (dipole_weight / polarizability_weight > 0): the mean squared error over the three (nine) components of the predicted and target dipole (polarizability) of each structure, the tensorial targets of PaiNN (and the dipoles of PhysNet and AIMNet2). These terms take the global huber_delta and are not weighted per head.
Charges (charge_weight > 0): the squared error of the predicted charges of every atom against data.charges, averaged over the atoms like the force term (the charge networks of the 3G and 4G HDNNPs, and any model with a "charges" output). It takes the global huber_delta and is not weighted per head.
LES charges and dipole (les_charge_weight / les_dipole_weight > 0): the same two terms for the les_charges and les_dipole outputs of LatentEwald, against the same charges and dipole labels.

Per-structure weighting. When data.weight is set (shape (B,)), every term becomes a weighted mean instead of a plain one: structure b contributes in proportion to weight[b], and for the force term that weight is spread to each of its atoms. Uniform weights reproduce the unweighted loss exactly, so this is a no-op unless the dataset asks for it. It exists because the two terms otherwise disagree about what a sample is: energy counts each structure once while forces count each atom once, so a set mixing small and large structures, or dense scans with sparse sampling, silently allocates the fit. Set weight in the structure dicts (see to_graph()) to allocate it on purpose.

Term and atom weights. data.energy_weight / force_weight / stress_weight (shape (B,)) multiply weight for one term only, and data.atom_force_weight (shape (N,)) multiplies the weight of each atom’s force error. The force term is then sum_i u_i |dF_i|^2 / (3 sum_i u_i) with u_i = weight * force_weight * atom_force_weight, so a weighted atom counts like an atom of an unweighted structure. An atom of weight 0 stays in the graph as environment; only its force label is dropped. A term whose weights are all zero in a batch is skipped.

Partially labelled batches. A batch may mix structures with and without energy, force or stress labels (data.energy_mask / data.forces_mask / data.stress_mask, set by collate()): those terms then average over the labelled structures only, and a batch without any is skipped for that term.

Several heads. When the batch carries data.head with more than one head, the loss is the sum over heads of that head’s weighted terms, each averaged over the head’s own structures (the multi-head replay objective, Tompa et al. eq. 3), and head_weights gives a head its own energy, force and stress weights (eq. 4). A single-head batch reduces to the plain loss above; with head_names given, its logs still carry the head’s prefix, so every batch of a multi-head run logs its terms under the same keys.

Dataset-level metrics. Each term is a weighted mean, and a mean over batches of unequal size (or with unequal numbers of labels) is not the mean of the batch means. Pass a dict as counts to have it filled with {key: (count, term_weight)} for every term key of logs: the weight the term’s mean divides by (structures, atoms or their summed weights) and the term’s loss weight. The mean of a term over many batches is then sum(logs[key] * count) / sum(count), exactly the value one batch holding all the structures would give, and the loss over them is sum(term_weight * mean) over the keys.

Huber tails. With huber_delta > 0 the squared error is replaced by a Huber-like function that is quadratic up to delta and linear beyond, which caps the pull of a few large residuals. The three residuals have very different natural scales (per-atom energies in eV, forces in eV/A), so a single huber_delta is rarely right for all of them; the per-term arguments override it. See _residual_sq() for the exact form and how it compares with MACE’s.

Parameters:

pred (dict of str to Tensor) – Model outputs. Must contain "energy"; may also contain "forces" and "stress".
data (AtomicGraph) – The (possibly batched) target graph. Reads energy, forces, stress, n_atoms (per-structure atom counts), batch (the structure index of every atom), the optional weight and the optional head.
energy_weight (float) – Weight applied to the energy term. The term is skipped when this is not positive or data.energy is None.
force_weight (float) – Weight applied to the force term. The term is skipped when this is not positive, data.forces is None, or pred lacks "forces".
stress_weight (float) – Weight applied to the stress term. The term is skipped when this is not positive, data.stress is None, or pred lacks "stress".
dipole_weight (float, optional) – Weights of the dipole and polarizability terms, skipped the same way (data.dipole / data.polarizability and the "dipole" / "polarizability" predictions). Default 0.0.
polarizability_weight (float, optional) – Weights of the dipole and polarizability terms, skipped the same way (data.dipole / data.polarizability and the "dipole" / "polarizability" predictions). Default 0.0.
charge_weight (float, optional) – Weight of the partial-charge term, skipped the same way (data.charges and the "charges" prediction). Default 0.0.
les_dipole_weight (float, optional) – Weights of the LES dipole and charge terms (data.dipole / data.charges against the "les_dipole" / "les_charges" predictions). Default 0.0.
les_charge_weight (float, optional) – Weights of the LES dipole and charge terms (data.dipole / data.charges against the "les_dipole" / "les_charges" predictions). Default 0.0.
huber_delta (float, optional) – Default crossover from quadratic to linear for every term. 0.0 (the default) means plain squared error throughout.
huber_delta_energy (float, optional) – Per-term overrides of huber_delta. None falls back to it.
huber_delta_forces (float, optional) – Per-term overrides of huber_delta. None falls back to it.
huber_delta_stress (float, optional) – Per-term overrides of huber_delta. None falls back to it.
head_weights (mapping of int to sequence of float, optional) – {head_index: (energy_weight, force_weight, stress_weight)} for the heads whose weights differ from the global ones.
head_names (sequence of str, optional) – Head names for the per-head log keys ("<head>/energy_mse"); indices are used when omitted.
counts (dict or None, optional) – When given, filled with {key: (count, term_weight)} for every term key of logs (see Dataset-level metrics above). Defaults to None.

Returns:

A pair (loss, logs) where loss is the scalar weighted total loss (a 0-dim tensor carrying gradients) and logs maps "loss" to the detached total plus, for each included term, its detached value under "energy_mse", "force_mse", "stress_mse", "dipole_mse", "polarizability_mse", "charge_mse", "les_dipole_mse" and/or "les_charge_mse" (with a "<head>/" prefix per head in a multi-head batch). Those keys keep their names when Huber tails are on, where they hold the Huber value rather than a mean square.

Return type:

tuple of (Tensor, dict of str to float)