probly.train.evidential.torch.ird_loss

probly.train.evidential.torch.ird_loss(alpha: Tensor, y: Tensor, adversarial_alpha: Tensor | None = None, p: float = 2.0, lam: float = 1.0, gamma: float = 1.0, normalize: bool = True) → Tensor[source]

Information Robust Dirichlet (IRD) loss for predictive uncertainty estimation.

Implements the loss proposed by Tsiligkaridis (2019), combining an Lp calibration term, a trigamma-based regularization term, and an optional entropy-based adversarial regularizer.

Reference:

Tsiligkaridis, “Information Robust Dirichlet Networks for Predictive Uncertainty Estimation”, 2019. https://arxiv.org/abs/1910.04819

Parameters:

  • alpha – Dirichlet concentration parameters, shape (B, K).

  • y – One-hot encoded class labels, shape (B, K).

  • adversarial_alpha – Dirichlet concentration parameters for adversarial inputs, shape (B_a, K).

  • p – Lp norm exponent controlling calibration strength.

  • lam – Weight of the regularization term.

  • gamma – Weight of the entropy regularization term.

  • normalize – Whether to normalize loss terms by batch size.

Returns:

Scalar IRD loss summed over all input examples.