DEUP on MNIST

Direct Epistemic Uncertainty Prediction (DEUP) trains a base classifier in phase one, then trains a separate error head in phase two that explicitly predicts per-sample cross-entropy errors using stationarizing features. The predicted error score serves as a direct measure of epistemic uncertainty.

from __future__ import annotations

import numpy as np
import torch
from torch import nn
from torch.utils.data import ConcatDataset, DataLoader, TensorDataset

from probly.method.deup import deup
from probly.quantification import quantify
from probly.representer import representer
from probly_benchmark.data import load_mnist

from examples.utils.model import MLPClassifier
from examples.utils.plotting import plot_mnist_uncertainty

Setup

SEED = 42
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

train_loader, test_loader = load_mnist(batch_size=256)

X_test_batches, y_test_batches = zip(*test_loader)
X_test = torch.cat([x.view(-1, 28 * 28) for x in X_test_batches])
y_test = torch.cat(list(y_test_batches))
images_test = (X_test.view(-1, 28, 28) * 255).byte()

X_train_batches, y_train_batches = zip(*train_loader)
X_train = torch.cat([x.view(-1, 28 * 28) for x in X_train_batches])
y_train = torch.cat(list(y_train_batches))
train_dataset_flat = TensorDataset(X_train, y_train)
train_loader_flat = DataLoader(train_dataset_flat, batch_size=256, shuffle=True)

Model

base_model = MLPClassifier(in_features=28 * 28, hidden_features=256, out_features=10)

deup_model = deup(
    base_model,
    hidden_size=512,
    n_hidden_layers=2,
    stationarizing_features=[
        "log_gmm_density",
        "log_mc_dropout_variance",
    ],
    predictor_type="logit_classifier",
)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
deup_model.to(device)

TorchDEUPPredictor(
  (encoder): MLPClassifier(
    (net): Sequential(
      (0): Linear(in_features=784, out_features=256, bias=True)
      (1): ReLU()
      (2): Linear(in_features=256, out_features=256, bias=True)
      (3): ReLU()
      (4): Identity()
    )
  )
  (classification_head): Linear(in_features=256, out_features=10, bias=True)
  (providers): ModuleList(
    (0): LogGMMDensity(
      (head): GaussianMixtureHead()
    )
    (1): LogMCDropoutVariance()
  )
  (error_head): ErrorPredictionHead(
    (net): Sequential(
      (0): Linear(in_features=2, out_features=512, bias=True)
      (1): ReLU()
      (2): Linear(in_features=512, out_features=512, bias=True)
      (3): ReLU()
      (4): Linear(in_features=512, out_features=1, bias=True)
    )
  )
)

Phase 1 Training: Base Classifier

Train the encoder and classification head with standard cross-entropy.

print("Phase 1: Training base classifier")
optimizer_phase1 = torch.optim.Adam(
    list(deup_model.encoder.parameters()) + list(deup_model.classification_head.parameters()),
    lr=1e-3,
)
criterion = nn.CrossEntropyLoss()

deup_model.train()
for _epoch in range(5):
    correct, total = 0, 0
    for inputs, targets in train_loader_flat:
        inputs, targets = inputs.to(device), targets.to(device)

        features = deup_model.encoder(inputs)
        logits = deup_model.classification_head(features)
        loss = criterion(logits, targets)

        optimizer_phase1.zero_grad()
        loss.backward()
        optimizer_phase1.step()
        correct += (logits.detach().argmax(-1) == targets).sum().item()
        total += len(targets)
    if correct / total >= 0.97:
        break

Phase 1: Training base classifier

Phase 2: Prepare Stationarizing Features & OOD Data

Collecting the DEUP error targets.

1. Freeze the backbone.

2. Fit auxiliary providers (e.g., GMM, Dropout) to compute stationarizing features.

3. Generate synthetic OOD data (uniform noise) to teach the error head about high-error regions.

4. Compute features and target error values (log of BCE loss) for all data.

print("\nPhase 2: Training error head")

for param in deup_model.encoder.parameters():
    param.requires_grad = False
for param in deup_model.classification_head.parameters():
    param.requires_grad = False

deup_model.eval()
providers = list(getattr(deup_model, "providers", []))

for provider in providers:
    provider.to(device)
    provider.fit(deup_model.encoder, deup_model.classification_head, train_loader_flat, device)

_orig_phi = deup_model._compute_stationarizing_features
# clamp stationarizing features to prevent exploding inputs to the error head
deup_model._compute_stationarizing_features = lambda *a: _orig_phi(*a).clamp(-10.0, 10.0)

# Augment the in-distribution data with synthetic uniform-noise OOD so the
# error head sees high-error regions and learns to flag off-manifold inputs.
ood_X = torch.FloatTensor(len(X_train), 28 * 28).uniform_(0, 1)
ood_y = torch.randint(0, 10, (len(X_train),))
phase2_loader = DataLoader(
    ConcatDataset([train_dataset_flat, TensorDataset(ood_X, ood_y)]),
    batch_size=256,
    shuffle=True,
)

bce_criterion = nn.BCELoss(reduction="none")

all_phi = []
all_targets = []

deup_model.eval()
with torch.no_grad():
    for inputs_, targets_ in phase2_loader:
        inputs, targets = inputs_.to(device), targets_.to(device)

        features = deup_model.encoder(inputs)
        logits = deup_model.classification_head(features)

        phi = deup_model._compute_stationarizing_features(features, logits)  # noqa: SLF001

        probs = torch.softmax(logits.float(), dim=-1).detach().cpu()
        one_hot = nn.functional.one_hot(targets_, num_classes=probs.size(-1)).float().detach().cpu()
        per_sample_bce = bce_criterion(probs, one_hot).sum(dim=-1)

        target_val = torch.log10(per_sample_bce.clamp(min=1e-10)).clamp(min=-5.0)

        all_phi.append(phi.detach().cpu())
        all_targets.append(target_val.detach().cpu())

phi_all = torch.cat(all_phi)
targets_all = torch.cat(all_targets)

error_head_dataset = torch.utils.data.TensorDataset(phi_all, targets_all)
error_head_loader = torch.utils.data.DataLoader(error_head_dataset, batch_size=64, shuffle=True, drop_last=False)

Phase 2: Training error head

Fitting DEUP density feature:   0%|          | 0/235 [00:00<?, ?it/s]
Fitting DEUP density feature:  21%|██▏       | 50/235 [00:00<00:00, 495.73it/s]
Fitting DEUP density feature:  43%|████▎     | 101/235 [00:00<00:00, 503.37it/s]
Fitting DEUP density feature:  65%|██████▍   | 152/235 [00:00<00:00, 502.24it/s]
Fitting DEUP density feature:  86%|████████▋ | 203/235 [00:00<00:00, 488.34it/s]
Fitting DEUP density feature: 100%|██████████| 235/235 [00:00<00:00, 488.72it/s]

Fitting LogGMMDensity scaler:   0%|          | 0/235 [00:00<?, ?it/s]
Fitting LogGMMDensity scaler:   7%|▋         | 17/235 [00:00<00:01, 163.60it/s]
Fitting LogGMMDensity scaler:  15%|█▍        | 35/235 [00:00<00:01, 169.61it/s]
Fitting LogGMMDensity scaler:  23%|██▎       | 53/235 [00:00<00:01, 171.46it/s]
Fitting LogGMMDensity scaler:  30%|███       | 71/235 [00:00<00:00, 172.46it/s]
Fitting LogGMMDensity scaler:  38%|███▊      | 89/235 [00:00<00:00, 173.13it/s]
Fitting LogGMMDensity scaler:  46%|████▌     | 107/235 [00:00<00:00, 173.53it/s]
Fitting LogGMMDensity scaler:  53%|█████▎    | 125/235 [00:00<00:00, 173.86it/s]
Fitting LogGMMDensity scaler:  61%|██████    | 143/235 [00:00<00:00, 174.11it/s]
Fitting LogGMMDensity scaler:  69%|██████▊   | 161/235 [00:00<00:00, 174.24it/s]
Fitting LogGMMDensity scaler:  76%|███████▌  | 179/235 [00:01<00:00, 174.37it/s]
Fitting LogGMMDensity scaler:  84%|████████▍ | 197/235 [00:01<00:00, 174.39it/s]
Fitting LogGMMDensity scaler:  91%|█████████▏| 215/235 [00:01<00:00, 174.60it/s]
Fitting LogGMMDensity scaler:  99%|█████████▉| 233/235 [00:01<00:00, 174.57it/s]
Fitting LogGMMDensity scaler: 100%|██████████| 235/235 [00:01<00:00, 173.67it/s]

Fitting LogMCDropoutVariance scaler:   0%|          | 0/235 [00:00<?, ?it/s]
Fitting LogMCDropoutVariance scaler:   2%|▏         | 4/235 [00:00<00:07, 31.03it/s]
Fitting LogMCDropoutVariance scaler:   3%|▎         | 8/235 [00:00<00:07, 31.14it/s]
Fitting LogMCDropoutVariance scaler:   5%|▌         | 12/235 [00:00<00:07, 31.19it/s]
Fitting LogMCDropoutVariance scaler:   7%|▋         | 16/235 [00:00<00:07, 31.21it/s]
Fitting LogMCDropoutVariance scaler:   9%|▊         | 20/235 [00:00<00:06, 31.21it/s]
Fitting LogMCDropoutVariance scaler:  10%|█         | 24/235 [00:00<00:06, 31.21it/s]
Fitting LogMCDropoutVariance scaler:  12%|█▏        | 28/235 [00:00<00:06, 31.23it/s]
Fitting LogMCDropoutVariance scaler:  14%|█▎        | 32/235 [00:01<00:06, 31.19it/s]
Fitting LogMCDropoutVariance scaler:  15%|█▌        | 36/235 [00:01<00:06, 31.17it/s]
Fitting LogMCDropoutVariance scaler:  17%|█▋        | 40/235 [00:01<00:06, 31.15it/s]
Fitting LogMCDropoutVariance scaler:  19%|█▊        | 44/235 [00:01<00:06, 31.15it/s]
Fitting LogMCDropoutVariance scaler:  20%|██        | 48/235 [00:01<00:06, 31.15it/s]
Fitting LogMCDropoutVariance scaler:  22%|██▏       | 52/235 [00:01<00:05, 31.15it/s]
Fitting LogMCDropoutVariance scaler:  24%|██▍       | 56/235 [00:01<00:05, 31.17it/s]
Fitting LogMCDropoutVariance scaler:  26%|██▌       | 60/235 [00:01<00:05, 31.19it/s]
Fitting LogMCDropoutVariance scaler:  27%|██▋       | 64/235 [00:02<00:05, 31.18it/s]
Fitting LogMCDropoutVariance scaler:  29%|██▉       | 68/235 [00:02<00:05, 31.18it/s]
Fitting LogMCDropoutVariance scaler:  31%|███       | 72/235 [00:02<00:05, 31.15it/s]
Fitting LogMCDropoutVariance scaler:  32%|███▏      | 76/235 [00:02<00:05, 31.13it/s]
Fitting LogMCDropoutVariance scaler:  34%|███▍      | 80/235 [00:02<00:04, 31.11it/s]
Fitting LogMCDropoutVariance scaler:  36%|███▌      | 84/235 [00:02<00:04, 31.09it/s]
Fitting LogMCDropoutVariance scaler:  37%|███▋      | 88/235 [00:02<00:04, 31.11it/s]
Fitting LogMCDropoutVariance scaler:  39%|███▉      | 92/235 [00:02<00:04, 31.11it/s]
Fitting LogMCDropoutVariance scaler:  41%|████      | 96/235 [00:03<00:04, 31.08it/s]
Fitting LogMCDropoutVariance scaler:  43%|████▎     | 100/235 [00:03<00:04, 31.12it/s]
Fitting LogMCDropoutVariance scaler:  44%|████▍     | 104/235 [00:03<00:04, 31.11it/s]
Fitting LogMCDropoutVariance scaler:  46%|████▌     | 108/235 [00:03<00:04, 31.10it/s]
Fitting LogMCDropoutVariance scaler:  48%|████▊     | 112/235 [00:03<00:03, 31.10it/s]
Fitting LogMCDropoutVariance scaler:  49%|████▉     | 116/235 [00:03<00:03, 31.12it/s]
Fitting LogMCDropoutVariance scaler:  51%|█████     | 120/235 [00:03<00:03, 31.12it/s]
Fitting LogMCDropoutVariance scaler:  53%|█████▎    | 124/235 [00:03<00:03, 31.13it/s]
Fitting LogMCDropoutVariance scaler:  54%|█████▍    | 128/235 [00:04<00:03, 31.06it/s]
Fitting LogMCDropoutVariance scaler:  56%|█████▌    | 132/235 [00:04<00:03, 31.09it/s]
Fitting LogMCDropoutVariance scaler:  58%|█████▊    | 136/235 [00:04<00:03, 31.10it/s]
Fitting LogMCDropoutVariance scaler:  60%|█████▉    | 140/235 [00:04<00:03, 31.12it/s]
Fitting LogMCDropoutVariance scaler:  61%|██████▏   | 144/235 [00:04<00:02, 31.13it/s]
Fitting LogMCDropoutVariance scaler:  63%|██████▎   | 148/235 [00:04<00:02, 31.13it/s]
Fitting LogMCDropoutVariance scaler:  65%|██████▍   | 152/235 [00:04<00:02, 31.13it/s]
Fitting LogMCDropoutVariance scaler:  66%|██████▋   | 156/235 [00:05<00:02, 31.15it/s]
Fitting LogMCDropoutVariance scaler:  68%|██████▊   | 160/235 [00:05<00:02, 31.15it/s]
Fitting LogMCDropoutVariance scaler:  70%|██████▉   | 164/235 [00:05<00:02, 31.17it/s]
Fitting LogMCDropoutVariance scaler:  71%|███████▏  | 168/235 [00:05<00:02, 31.15it/s]
Fitting LogMCDropoutVariance scaler:  73%|███████▎  | 172/235 [00:05<00:02, 31.18it/s]
Fitting LogMCDropoutVariance scaler:  75%|███████▍  | 176/235 [00:05<00:01, 30.82it/s]
Fitting LogMCDropoutVariance scaler:  77%|███████▋  | 180/235 [00:05<00:01, 30.22it/s]
Fitting LogMCDropoutVariance scaler:  78%|███████▊  | 184/235 [00:05<00:01, 30.53it/s]
Fitting LogMCDropoutVariance scaler:  80%|████████  | 188/235 [00:06<00:01, 30.71it/s]
Fitting LogMCDropoutVariance scaler:  82%|████████▏ | 192/235 [00:06<00:01, 30.84it/s]
Fitting LogMCDropoutVariance scaler:  83%|████████▎ | 196/235 [00:06<00:01, 30.93it/s]
Fitting LogMCDropoutVariance scaler:  85%|████████▌ | 200/235 [00:06<00:01, 31.00it/s]
Fitting LogMCDropoutVariance scaler:  87%|████████▋ | 204/235 [00:06<00:00, 31.05it/s]
Fitting LogMCDropoutVariance scaler:  89%|████████▊ | 208/235 [00:06<00:00, 31.07it/s]
Fitting LogMCDropoutVariance scaler:  90%|█████████ | 212/235 [00:06<00:00, 31.11it/s]
Fitting LogMCDropoutVariance scaler:  92%|█████████▏| 216/235 [00:06<00:00, 31.14it/s]
Fitting LogMCDropoutVariance scaler:  94%|█████████▎| 220/235 [00:07<00:00, 31.12it/s]
Fitting LogMCDropoutVariance scaler:  95%|█████████▌| 224/235 [00:07<00:00, 30.97it/s]
Fitting LogMCDropoutVariance scaler:  97%|█████████▋| 228/235 [00:07<00:00, 31.00it/s]
Fitting LogMCDropoutVariance scaler:  99%|█████████▊| 232/235 [00:07<00:00, 31.04it/s]
Fitting LogMCDropoutVariance scaler: 100%|██████████| 235/235 [00:07<00:00, 31.14it/s]

Phase 3: Training Error Head

Train the error head to predict the target error values from the stationarizing features. This head acts as the final uncertainty estimator.

deup_model.error_head.to(device)
opt_error = torch.optim.SGD(deup_model.error_head.parameters(), lr=0.005, momentum=0.9)
mse_loss_fn = nn.MSELoss()

deup_model.error_head.train()
for epoch in range(20):
    for phi_batch, tgt_batch in error_head_loader:
        phi, tgt = phi_batch.to(device, non_blocking=True), tgt_batch.to(device, non_blocking=True)

        loss = nn.functional.mse_loss(deup_model.error_head(phi), tgt)
        opt_error.zero_grad()
        loss.backward()
        opt_error.step()

Uncertainty Quantification

deup_model.eval()
rep = representer(deup_model)

X_test_dev = X_test.to(device)
with torch.no_grad():
    representation = rep.represent(X_test_dev)

uq = quantify(representation)
_unc = uq.total if hasattr(uq, "total") else (uq.epistemic if hasattr(uq, "epistemic") else uq.aleatoric)
uncertainty = _unc.detach().cpu().numpy() if isinstance(_unc, torch.Tensor) else np.asarray(_unc)
if uncertainty.ndim > 1:
    uncertainty = uncertainty.sum(axis=-1)

Predictions

with torch.no_grad():
    features = deup_model.encoder(X_test_dev)
    mean_probs = deup_model.classification_head(features).softmax(-1).cpu().numpy()

accuracy = (mean_probs.argmax(-1) == y_test.numpy()).mean() * 100
print(f"Test accuracy: {accuracy:.1f}%")

Test accuracy: 97.5%

Visualization

plot = plot_mnist_uncertainty(
    images_test,
    y_test,
    uncertainty,
    mean_probs,
    title="Top-5 Most Uncertain Test Predictions (DEUP)",
    unit="score",
)
plot.show()