probly.train.evidential.torch.rpn_ng_kl

probly.train.evidential.torch.rpn_ng_kl(mu: Tensor, kappa: Tensor, alpha: Tensor, beta: Tensor, mu0: Tensor, kappa0: Tensor, alpha0: Tensor, beta0: Tensor) → Tensor[source]

KL divergence between two Normal-Gamma distributions.

Computes the KL divergence between a predicted Normal-Gamma distribution and a prior Normal-Gamma distribution, as used in Regression Prior Networks to regularize out-of-distribution predictions.

Reference:

Malinin et al., “Regression Prior Networks”, NeurIPS 2020. https://arxiv.org/abs/2006.11590

Parameters:

  • mu – Predicted mean parameter, shape (B,).

  • kappa – Predicted scaling parameter, shape (B,).

  • alpha – Predicted shape parameter, shape (B,).

  • beta – Predicted scale parameter, shape (B,).

  • mu0 – Prior mean parameter, shape (B,).

  • kappa0 – Prior scaling parameter, shape (B,).

  • alpha0 – Prior shape parameter, shape (B,).

  • beta0 – Prior scale parameter, shape (B,).

Returns:

Scalar KL divergence between predicted and prior Normal-Gamma distributions, averaged over the batch.