Metadata-Version: 2.3
Name: edl-losses
Version: 0.6.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"
    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:**

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

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

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

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

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

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