torchcell.losses.distributional.energy_score

energy_score(mu: Tensor, sigma: Tensor, v: Tensor | None, target: Tensor, num_samples: int = 128) → Tensor[source]

Monte-Carlo energy score of a low-rank-plus-diagonal Gaussian, per row.

The energy score is the multivariate proper scoring rule that GENERALIZES CRPS to a joint distribution (Gneiting & Raftery 2007):

\[\mathrm{ES}(F, y) = \mathbb E\|X-y\|_2 - \tfrac12\,\mathbb E\|X-X'\|_2, \qquad X, X' \stackrel{iid}{\sim} F .\]

At \(F=1\) the Euclidean norm is the absolute value and this is exactly the energy form of CRPS, so gaussian_crps() is its ORACLE in one dimension.

The predictive family is Gaussian with a low-rank-plus-diagonal covariance,

\[\Sigma = \mathrm{diag}(\sigma^2) + V V^\top,\quad V\in\mathbb R^{F\times k},\]

sampled WITHOUT ever forming \(\Sigma\) (which would be F x F):

\[x = \mu + \sigma\odot\varepsilon + V z,\quad \varepsilon\sim\mathcal N(0, I_F),\ z\sim\mathcal N(0, I_k).\]

The estimator from m samples is

\[\widehat{\mathrm{ES}} = \frac1m\sum_i \|x_i-y\| - \frac{1}{2m(m-1)}\sum_{i\neq j}\|x_i-x_j\| ,\]

which is UNBIASED for every m >= 2: the pairwise term divides by m(m-1), not m^2. Excluding the m zero self-distances is exactly what removes the bias. Dividing by m^2 instead keeps those zeros in the average, shrinking the SUBTRACTED term by a factor (m-1)/m and so inflating the score by \(\tfrac{1}{2m}\mathbb E\|X-X'\|\). That surcharge is proportional to the predictive SPREAD, so it is not a harmless offset: it penalizes dispersion and moves the optimum to an over-sharp (overconfident) distribution, the more so the smaller m is. With the correct denominator m trades variance only. The pairwise distances go through torch.cdist(), which materializes [B, m, m]; the naive x[:, :, None] - x[:, None] would materialize [B, m, m, F] (~790 MB at B=32, m=10, F=6169) and is never formed.

Sampling is reparameterized – eps and z are noise, mu/sigma/V carry the gradient – so the score is differentiable end to end. Nothing is detached.

k = 0 (v is None or v.shape[-1] == 0) is the ABLATION arm and the distinction is worth stating precisely: the SCORE still couples coordinates (the Euclidean norm is taken over all F features jointly), but the predictive FAMILY is diagonal, so it cannot represent any correlation between features. Comparing k = 0 against k > 0 therefore separates “scoring jointly” from “modelling jointly”.

Parameters:

  • mu – Predicted means [B, F].

  • sigma – Predicted marginal scales [B, F] (STRICTLY positive); the diagonal part.

  • v – Global low-rank factor V [F, k], SHARED across the batch (a learned feature-correlation basis, not a per-row output), or None for k = 0.

  • target – Observations [B, F].

  • num_samples – m, the number of predictive samples. Must be >= 2.

Returns:

Per-row energy score [B] (callers reduce over the batch).