TASTE: Task-Aware Out-of-Distribution Detection via Stein Operators

TL;DR

TASTE introduces a task-aware out-of-distribution detection method using Stein operators, linking distribution shifts to model sensitivity, providing interpretability, localization, and outperforming existing baselines on various benchmarks.

Contribution

The paper presents a novel Stein operator-based framework that connects distribution shifts to model sensitivity, enabling detection, localization, and interpretability in OOD detection.

Findings

Aligns closely with task degradation

Outperforms established baselines

Provides interpretable per-pixel diagnostics

Abstract

Out-of-distribution detection methods are often either data-centric, detecting deviations from the training input distribution irrespective of their effect on a trained model, or model-centric, relying on classifier outputs without explicit reference to data geometry. We propose TASTE (Task-Aware STEin operators): a task-aware framework based on so-called Stein operators, which allows us to link distribution shift to the input sensitivity of the model. We show that the resulting operator admits a clear geometric interpretation as a projection of distribution shift onto the sensitivity field of the model, yielding theoretical guarantees. Beyond detecting the presence of a shift, the same construction enables its localisation through a coordinate-wise decomposition, and for image data-provides interpretable per-pixel diagnostics. Experiments on controlled Gaussian shifts, MNIST under geometric perturbations, and CIFAR-10 perturbed benchmarks demonstrate that the proposed method aligns closely with task degradation while outperforming established baselines.