edl-losses 0.4.0

Implementation of various EDL loss functions

edl-losses

A research library implementing Evidential Deep Learning (EDL) loss functions for uncertainty-aware classification in PyTorch.

EDL methods replace the standard softmax output with a distribution over class probabilities (typically a Dirichlet) allowing the model to express not just which class is most likely, but how confident it is in that prediction. This enables detection of out-of-distribution inputs, misclassification detection, and calibrated uncertainty estimates, all in a single forward pass with no sampling required.

The following papers are currently implemented:

• EDL: Sensoy et al. (2018) — Evidential Deep Learning to Quantify Classification Uncertainty
• GEN: Sensoy et al. (2020) — Uncertainty-Aware Deep Classifiers using Generative Models
• F-EDL: Yoon and Kim (2026) — Uncertainty Estimation by Flexible Evidential Deep Learning

Installation

Requires Python 3.13+ and PyTorch 2.11+.

Using uv:

uv install edl-losses

Using pip

pip install edl-losses

Or from source:

git clone https://github.com/LucaCtt/edl-losses
cd edl_losses
pip install -e .

Quick Background

All three methods share the same core idea: instead of having the network output a point estimate of class probabilities via softmax, the network outputs the parameters of a Dirichlet distribution over class probabilities. The mean of this distribution is used for classification, while its concentration (or variance) quantifies uncertainty.

EDL (Sensoy et al. 2018)

The network outputs raw logits which are passed through ReLU to produce non-negative evidence for each class. The Dirichlet parameters are α = evidence + 1. A custom loss (SSE, CE, or Type II MLE) is minimized along with a KL divergence term that penalizes evidence generated for incorrect classes. Uncertainty is K / S where S = Σαₖ.

GEN (Sensoy et al. 2020)

Extends EDL by incorporating out-of-distribution (OOD) samples during training. The network is trained with a Bernoulli noise-contrastive estimation (NCE) loss that discriminates in-distribution from OOD samples per class. Evidence is derived via exp() rather than ReLU. OOD samples can be generated by a VAE+GAN or any other perturbation strategy.

F-EDL (Yoon & Kim 2026)

Replaces the Dirichlet with a Flexible Dirichlet (FD) distribution, which is a mixture of Dirichlets parameterized by concentration α, allocation probabilities p, and dispersion τ. This allows multimodal beliefs over class probabilities, better handling of ambiguous inputs. The model outputs three parameter vectors from three separate heads. Uncertainty decomposes into total (TU), epistemic (EU), and aleatoric (AU) components in closed form.

Example Usage

See the examples.ipynb notebook, which compares the EDL losses in classifying the rotated "1" digit from MNIST.

API

from edl_losses import (
    EDLLoss,
    edl_inference,
    GENLoss,
    FEDLLoss,
    fedl_inference,
)

EDLLoss

edl_loss = EDLLoss(
    loss_type: str = "sse", # "sse" | "ce" | "mse"
    kl_reg: bool = True,
    annealing_epochs: int = 10
)
edl_loss(
    logits: Tensor, # (B, K) raw network output, before any activation
    labels: Tensor, # (B,) ground truth class indices
    epoch: int | None, # current training epoch, used for KL annealing. If None, it will be incremented automatically at each call.
) -> Tensor # scalar

Implements Equations 3–5 of Sensoy et al. 2018. ReLU is applied internally to produce evidence. Three base losses are available:

• "sse" — Sum of squares Bayes risk (recommended by the paper, most stable)
• "ce" — Cross-entropy Bayes risk
• "mse" — Type II Maximum Likelihood

When kl_reg=True, a KL divergence term penalizes evidence assigned to incorrect classes, annealed from 0 to 1 over the first annealing_epochs epochs.

Note: The original paper uses ReLU for evidence. This can slow convergence compared to GEN/F-EDL which use exp(). If accuracy is lower than expected, consider clamping logits and using exp() before passing to this loss.

Example:

model = MyClassifier()  # output layer has no activation
optimizer = torch.optim.Adam(model.parameters())

for epoch in range(1, num_epochs + 1):
    for x, y in dataloader:
        optimizer.zero_grad()
        loss = EDLLoss()(model(x), y, epoch=epoch, loss_type="sse")
        loss.backward()
        optimizer.step()

edl_inference

edl_inference(
    logits: Tensor, # (B, K) raw network output
) -> tuple[Tensor, Tensor, Tensor] # (predicted_classes (B,), uncertainty (B,), class_probs (B, K))

Uncertainty is K / S where S = Σαₖ. Values close to 1 indicate maximum uncertainty ("I don't know"); values close to 0 indicate high confidence.

Example:

with torch.no_grad():
    pred, uncertainty, probs = edl_inference(model(x)) # uncertainty ∈ (0, 1] — high means uncertain

GENLoss

gen_loss = GENLoss(
    eps:  float = 1e-8,
)
gen_loss(
    logits_in: Tensor, # (B, K) network output on in-distribution samples
    logits_out: Tensor, # (B, K) network output on OOD samples
    labels: Tensor, # (B,) ground truth class indices
    beta: float | str = "auto", # KL weight; "auto" uses expected misclassification prob
) -> Tensor # scalar

Implements Equations 4–6 of Sensoy et al. 2020. Requires OOD samples at training time. The loss has two components:

• L1 — Bernoulli NCE loss: trains each output fₖ as a binary classifier distinguishing class-k samples from OOD samples.
• L2 — KL regularizer: pushes the conditional Dirichlet over non-true classes toward uniform, weighted by beta.

When beta="auto", the weight is set to (1 - p̂ₖ) per sample, i.e. the expected misclassification probability, which implements learned loss attenuation.

Note: Clamp logits to a reasonable range (e.g. [-10, 10]) before passing to this loss to avoid numerical instability from the internal exp().

Example:

for x, y in dataloader:
    x_ood = generate_ood_samples(x)  # your OOD generator
    optimizer.zero_grad()
    loss = GENLoss()(model(x), model(x_ood), y)
    loss.backward()
    optimizer.step()

GEN uses the same edl_inference function for inference, since the output head is identical to EDL.

FEDLLoss

fedl_loss = FEDLoss(
    eps:    float = 1e-8,
)
fedl_loss(
    alpha: Tensor, # (B, K) concentration parameters, from exp() head
    p: Tensor, # (B, K) allocation probabilities, from softmax() head
    tau: Tensor, # (B,) or (B, 1) dispersion, from softplus() head
    labels: Tensor, # (B,)
) -> Tensor # scalar

Implements the objective from Section 3.2 and Appendix A.1 of Yoon & Kim 2026. The loss has two components:

• L_MSE — Expected MSE over the FD distribution, computed in closed form via FD moments.
• L_reg — Brier score on p, promoting well-calibrated allocation probabilities.

No KL term or annealing schedule is required. The model must expose three separate output heads:

class MyFEDLModel(nn.Module):
    def forward(self, x):
        z = self.backbone(x)
        alpha = torch.exp(self.head_alpha(z)) # evidence, > 0
        p = torch.softmax(self.head_p(z), dim=-1) # allocation probs, sums to 1
        tau = F.softplus(self.head_tau(z)).squeeze(-1) # dispersion, > 0
        return alpha, p, tau

Example:

for x, y in dataloader:
    optimizer.zero_grad()
    alpha, p, tau = model(x)
    loss = FEDLoss()(alpha, p, tau, y)
    loss.backward()
    optimizer.step()

fedl_inference

fedl_inference(
    alpha: Tensor, # (B, K)
    p: Tensor, # (B, K)
    tau: Tensor, # (B,) or (B, 1)
    eps: float = 1e-8,
) -> tuple[Tensor, Tensor, Tensor] # (predicted_classes (B,), total_uncertainty (B,), class_probs (B, K))

Returns predicted classes, total uncertainty (TU), and expected class probabilities E[π]. TU is defined as 1 - Σ E[πₖ]² and lies in (0, 1].

License

MIT. See LICENSE.

Author

Luca Cotti ([e-mail])