Skip to content

Beta-NLL (heteroscedastic Gaussian)

← Gaussian | Faithful Gaussian →

β-NLL rescales the per-element diagonal Gaussian negative log-likelihood with a detached function of the predicted variance so that gradients do not collapse variance for the wrong reasons. It is a drop-in sibling of GaussianNLLLoss for the same model outputs \((\mu, \log\sigma^2)\).

When to use which

Situation Prefer
Stable heteroscedastic training; variance collapse or trivial σ² BetaNLLLoss with \(\beta \in (0, 1)\) (often \(0.5\))
Standard NLL baseline; well-tuned optimisation GaussianNLLLoss
Finite-sample coverage guarantees Conformal prediction on top of a calibrated probabilistic model — β-NLL is not conformal

Mathematical definition

Let \(\sigma^2_i\) be the predicted variance for dimension \(i\) (from \(\log\sigma^2_i\)), and let \(\mathcal{L}_{\text{NLL},i}\) be the usual diagonal Gaussian NLL contribution including the \(\log(2\pi)\) term (matching GaussianNLLLoss).

The β-NLL uses a prefactor with stopped variance:

\[ \boxed{\; \mathcal{L}_{\beta\text{-NLL},i} = \big(\max(\sigma_i^2, \varepsilon)\big)^{\beta} \cdot \mathcal{L}_{\text{NLL},i} \;} \]

where \(\big(\max(\sigma_i^2, \varepsilon)\big)^{\beta}\) is computed from detached \(\sigma_i^2\) (no gradient through the prefactor), i.e. Seitzer et al.'s \(\lfloor\sigma^{2\beta}\rfloor\,\mathcal{L}_{\text{NLL}}\): large-variance points are up-weighted relative to plain NLL, counteracting the \(1/\sigma^2\) gradient scaling that lets the mean ignore them. The per-sample loss sums these terms over the feature dimension (1-D [B] inputs are one scalar target per sample). For \(\beta = 0\) the per-element terms are exactly those of GaussianNLLLoss (which averages rather than sums over features under reduction="mean").

API

from torchregress.losses import BetaNLLLoss, beta_nll_loss, GaussianNLLLoss

# Class (same y_pred layouts as GaussianNLLLoss: tuple or concatenated)
loss_fn = BetaNLLLoss(beta=0.5, reduction="mean")
loss = loss_fn((mean, log_var), target, mask=mask, weights=w)

# Functional
loss = beta_nll_loss((mean, log_var), target, beta=0.5, reduction="mean")

create_loss_from_config accepts "type": "beta_nll" with optional "beta".

Runnable example

See the Beta-NLL heteroscedastic demo and the script examples/heteroscedastic_beta_nll_demo.py in the repository.

References

# Reference
1 Seitzer et al. "On the Pitfalls of Heteroscedastic Uncertainty Estimation with Probabilistic Neural Networks" NeurIPS, 2022.

Limitations

  1. Beta selection is heuristic: \(\beta\) controls the variance-rescaling strength but has no closed-form optimum. Typical values are \(\beta = 0.5\) (recommended default from Seitzer et al., 2022). Values near 0 approach standard NLL; values near 1 aggressively downweight variance gradients. Tune on validation NLL.
  2. Not a conformal method: β-NLL stabilises variance training but provides no coverage guarantees. For finite-sample prediction intervals, wrap with SplitConformal or CQR.
  3. Same head as GaussianNLL: β-NLL uses the same \((\mu, \log\sigma^2)\) output as standard NLL. It cannot fix a poorly designed model head — only how variance gradients flow during training.

Recommendations

  • Default: Start with BetaNLLLoss(beta=0.5). This is the recommended value from the Seitzer et al. paper and works well across a broad range of regression problems.
  • When NLL is sufficient: If standard GaussianNLLLoss trains stably without variance collapse, β-NLL offers no benefit. Don't add complexity you don't need.
  • Monitor variance: Track predicted \(\sigma^2\) during training. If it drifts below \(10^{-3}\) or above \(10^2\), the β parameter may need adjustment.
  • Alternative for decoupling: If you specifically want to preserve point prediction accuracy while learning variance, consider FaithfulGaussianLoss instead — it explicitly splits mean and variance objectives.

Next steps