Mahalanobis OOD on Two Moons

The Mahalanobis out-of-distribution detector fits class-conditional Gaussians with a shared covariance on a feature extractor, then scores each input by its Mahalanobis distance to the nearest class centroid. A single deterministic forward pass yields both a class prediction and an epistemic / out-of-distribution score. This example also calibrates the multi-layer combination weights on a small batch of synthetic out-of-distribution points.

from __future__ import annotations

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

import torch
from torch import nn

from probly.representer import representer
from probly.method.mahalanobis import mahalanobis

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

Setup

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

The transformation strips the classification head to expose penultimate features; the original head is kept for predictions. No spectral normalization is applied – the detector works on the plain feature extractor.

base_model = ResFFN()

mahalanobis_model = mahalanobis(base_model, predictor_type="logit_classifier")

Training

opt = torch.optim.Adam(mahalanobis_model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

mahalanobis_model.train()
for epoch in range(200):
    opt.zero_grad()

    features = mahalanobis_model.encoder(X_tensor)
    logits = mahalanobis_model.classification_head(features)

    loss = criterion(logits, y_tensor)

    loss.backward()
    opt.step()

Fit the Mahalanobis Heads

Estimate the per-class means and the shared covariance from the training features. This only needs to happen once after training.

mahalanobis_model.eval()
mahalanobis_model.fit_mahalanobis_heads(X_tensor, y_tensor)

Calibrate the Combiner

Fit the logistic-regression combination weights on the in-distribution data and a small batch of synthetic out-of-distribution points sampled far from the two moons.

rng = torch.Generator().manual_seed(0)
ood_tensor = torch.rand(500, 2, generator=rng) * 8.0 - 4.0

mahalanobis_model.fit_combiner(X_tensor, ood_tensor)

Uncertainty Evaluation

rep = representer(mahalanobis_model)

plot = plot_example_uncertainty(
    X, y, rep, title="Mahalanobis Out-of-Distribution Score", notion="epistemic", log_scale=True
)

plot.show()