torchcell.losses.distributional.DistHead

class DistHead(mode: str, quantiles: Tensor | list[float] | None = None, sigma_floor: float = 0.001, num_features: int | None = None, rank: int = 32, num_samples: int = 128, v_init: float = 0.1)[source]

Bases: Module

Pluggable distributional readout: interpret a head’s raw params, score them, pool.

Stateless except for the quantile grid (quantile mode) and the global low-rank factor V (energy mode). Given a structural head’s output params – [..., F] (point) or [..., F, P] (P=2 gaussian / laplace / nll_gaussian / energy, P=K quantile) – it exposes:

  • point() – the point estimate [..., F] used by the (loss-agnostic) metric and any downstream reporting: identity (point), mu (gaussian / nll_gaussian / energy), the predictive MEDIAN mu (laplace), or the median quantile (quantile);

  • loss() – a masked scalar training loss: MSE (point), mean Gaussian CRPS (gaussian), mean pinball (quantile), mean Laplace CRPS (laplace), mean Gaussian NLL (nll_gaussian), or mean energy score (energy);

  • pit() – per-feature PIT values [B, F] for the calibration diagnostic (every mode but point).

energy is the one mode with learnable state of its own: V is an nn.Parameter of shape [num_features, rank] that is GLOBAL (shared across the batch), because a correlation basis over features is a property of the phenotype space, not of an individual strain. It therefore does NOT widen param_dim – the structural head still emits (mu, sigma), i.e. P=2, and only the head owns V.

property taus: Tensor

The quantile grid (empty for non-quantile modes).

The buffer is stored as _taus and surfaced through this property because nn.Module.__getattr__ is typed Tensor | Module; casting once here keeps every caller (and mypy –strict) working with a plain Tensor.

property v: Tensor | None

The global low-rank factor V [F, k] (energy mode) or None.

None for every non-energy mode and for energy with rank=0 (the diagonal-only ablation). Surfaced through a property for the same reason as taus: nn.Module.__getattr__ is typed Tensor | Module.

property param_dim: int

1 (point), K (quantile), or 2 (every parametric mode).

energy is 2 as well – its extra state is the GLOBAL V, owned by this head, not per-feature output of the structural head.

Type:

Params per feature

point(params: Tensor) → Tensor[source]

Reduce raw params to a point estimate [..., F] (metric / reporting).

Identity (point), the median knot (quantile), or mu for every parametric mode – which is the predictive MEAN for gaussian / nll_gaussian / energy and the predictive MEDIAN for laplace.

samples(params: Tensor, num_samples: int) → Tensor[source]

Draw energy-mode predictive samples [B, m, F] (reparameterized).

Parameters:

  • params – Head params [B, F, 2].

  • num_samples – m, samples per row.

Returns:

Samples [B, m, F] from N(mu, diag(sigma^2) + V V^T).

loss(params: Tensor, target: Tensor, mask: Tensor | None = None, feature_mask: Tensor | None = None) → Tensor[source]

Masked scalar training loss for this head.

Parameters:

  • params – Head params [B, F] (point) or [B, F, P] (every other mode).

  • target – Point targets [B, F] (same F as params).

  • mask – Optional bool [B] selecting supervised rows. When no row is supervised, a graph-connected zero is returned (grads flow as zero).

  • feature_mask – Optional bool [B, F] selecting which FEATURES to score. Used by the masked-label objective to score only the genes still hidden at the current unmasking step – scoring a revealed gene would let the model copy its own input. Row-masked alongside params/target so the two stay aligned.

Returns:

Scalar loss.

pit(params: Tensor, target: Tensor, mask: Tensor | None = None, num_samples: int | None = None) → Tensor[source]

Per-feature PIT values [B, F] for this head’s predictive distribution.

The mode-agnostic calibration diagnostic – see pit_values() for how to read the histogram and coverage() for the interval summary. Not available in point mode (no predictive distribution).

Parameters:

  • params – Head params [B, F] (point) or [B, F, P].

  • target – Point targets [B, F].

  • mask – Optional bool [B] selecting rows, as in loss().

  • num_samples – energy mode only – samples used to build the empirical CDF. Defaults to DEFAULT_ENERGY_EVAL_SAMPLES, NOT the head’s training m: PIT resolution is 1/m, so the (deliberately cheap) training m would quantize the histogram far too coarsely to read. Eval needs no gradients, so the extra samples are nearly free.

Returns:

PIT values [B, F] in [0, 1].