Metadata-Version: 2.3
Name: edl-losses
Version: 0.5.0
Summary: Implementation of various EDL loss functions
Requires-Dist: numpy>=2.4.4
Requires-Dist: torch>=2.11.0
Requires-Python: >=3.13
Description-Content-Type: text/markdown

# 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](http://arxiv.org/abs/1806.01768)
- GEN: [Sensoy et al. (2020) — Uncertainty-Aware Deep Classifiers using Generative Models](http://arxiv.org/abs/2006.04183)
- F-EDL: [Yoon and Kim (2026) — Uncertainty Estimation by Flexible Evidential Deep Learning](http://arxiv.org/abs/2510.18322)

## Installation

Requires Python 3.13+ and PyTorch 2.11+.

Using `uv`:
```bash
uv install edl-losses
```

Using `pip`
```bash
pip install edl-losses
```

Or from source:
```bash
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](examples.ipynb) notebook, which compares the EDL losses in classifying the rotated "1" digit from MNIST.

## API

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

### `EDLLoss`

```python
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, # 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. 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:**

```python
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`

```python
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:**

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

### `GENLoss`

```python
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:**

```python
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`

```python
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:

```python
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:**

```python
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`

```python
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](LICENSE).

## Author

Luca Cotti (<luca.cotti@unibs.it>)