DEUP on Two Moons

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

from sklearn.datasets import make_moons
import torch
from torch import nn
from torch.utils.data import ConcatDataset, DataLoader, TensorDataset

from probly.method.deup import deup
from probly.representer import representer

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

Setup

# Set random seed for reproducibility.
SEED = 42
torch.manual_seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

X, y = make_moons(n_samples=500, noise=0.05, random_state=0)
X_tensor = torch.from_numpy(X).float()
y_tensor = torch.from_numpy(y).long()

dataset = TensorDataset(X_tensor, y_tensor)
train_loader = DataLoader(dataset, batch_size=64, shuffle=True)

Model

base_model = MLPClassifier()

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=2, out_features=64, bias=True)
      (1): ReLU()
      (2): Linear(in_features=64, out_features=64, bias=True)
      (3): ReLU()
      (4): Identity()
    )
  )
  (classification_head): Linear(in_features=64, out_features=2, 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 loss.

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(50):
    for inputs, targets in train_loader:
        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()

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.

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, 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(500, 2).uniform_(-3, 3)
ood_y = torch.randint(0, 2, (500,))
phase2_loader = DataLoader(
    ConcatDataset([dataset, TensorDataset(ood_X, ood_y)]),
    batch_size=64,
    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)

Fitting DEUP density feature:   0%|          | 0/8 [00:00<?, ?it/s]
Fitting DEUP density feature: 100%|██████████| 8/8 [00:00<00:00, 2830.40it/s]

Fitting LogGMMDensity scaler:   0%|          | 0/8 [00:00<?, ?it/s]
Fitting LogGMMDensity scaler: 100%|██████████| 8/8 [00:00<00:00, 1565.77it/s]

Fitting LogMCDropoutVariance scaler:   0%|          | 0/8 [00:00<?, ?it/s]
Fitting LogMCDropoutVariance scaler: 100%|██████████| 8/8 [00:00<00:00, 298.44it/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(200):
    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()

Evaluation

deup_model.error_head.eval()

rep = representer(deup_model)
plot = plot_example_uncertainty(X, y, rep, title="DEUP Predictive Uncertainty", notion="total")
plot.show()