edl-losses 0.6.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"
    beta: float | Literal["anneal"] = "anneal",
    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.
) -> Tensor # scalar

Implements Equations 3–5 of Sensoy et al. 2018. ReLU is applied internally to produce evidence, but you may want to apply a softplus or exp to logits to avoid zeroing out negative logits.

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 beta is "anneal", the KL regularization term is weighted by min(1, epoch / annealing_epochs), which gradually increases the influence of the KL term over the first annealing_epochs epochs. This helps prevent early underfitting when evidence is still low. When beta is set to a fixed value (e.g. 1.0) applies a constant weight to the KL term throughout training. Setting beta=0 disables the KL term entirely, which may lead to faster convergence but worse uncertainty estimates.

Example:

model = MyClassifier()  # output layer has no activation
optimizer = torch.optim.Adam(model.parameters())
loss_fn = EDLLoss(loss_type="sse", beta="anneal", annealing_epochs=10)

for epoch in range(1, num_epochs + 1):
    for x, y in dataloader:
        optimizer.zero_grad()
        loss = loss_fn(model(x), y, epoch)
        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(
    beta: float | Literal["auto", "anneal"] = "auto", # KL weight; "auto" uses expected misclassification prob, "anneal" uses linear annealing
    anneal_epochs: int = 10, # number of epochs for KL annealing if `beta` is "anneal"
    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
    epoch: int | None = None, # current training epoch for KL annealing if `beta` is set to "anneal"
) -> 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. When beta="anneal", the KL term is weighted by min(1, epoch / anneal_epochs), which gradually increases the influence of the KL term over the first anneal_epochs epochs. Setting beta to a fixed float applies a constant weight to the KL term throughout training. Setting beta=0 disables the KL term, which may lead to faster convergence but worse uncertainty estimates.

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

Example:

loss_fn = GENLoss(beta="auto")  # or beta="anneal" for linear KL annealing

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

gen_inference

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

This is the same as edl_inference, but exp() is used instead of relu() to comput evidence.

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:

loss_fn = FEDLoss()

for x, y in dataloader:
    optimizer.zero_grad()
    alpha, p, tau = model(x)
    loss = loss_fn(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])