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

Michał Kozyra; Gesine Reinert

cs.LGFeb 7, 2026

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

Authors: Michał Kozyra, Gesine Reinert

Organizations: Department of Statistics, University of Oxford, United Kingdom.

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 links distribution shift to the input sensitivity of the model. The resulting operators admit a geometric interpretation as a projection of distribution shift onto the sensitivity field of the model, with theoretical guarantees. Beyond detecting a shift, the 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/ImageNet-1K perturbed benchmarks demonstrate that TASTE aligns closely with task degradation while showing competitive performance compared to established baselines, with gains that depend on the shift regime.

Figures & tables

Panel of Figure 1 — Figure 1: Task-aware intuition behind TASTE. Blue arrows show the score-shift field between the training density $p$ and test density $q$ . Without explicit knowledge of $q$ , TASTE measures alignment between distribution shift and the input sensitivity of the model $\nabla f$ : shifts aligned with $\nabla f$ ( magenta ) produce a large response; orthogonal shifts ( brown ) are largely ignored even when large.

Panel of Figure 2 — Figure 2: Directional shift experiment. Prediction error and the task-aware TASTE functional vary strongly with the direction of shift, peaking (in magnitude) when the shift aligns with the sensitive direction $(1,-1)$ . In contrast, the density-based score remains nearly constant for all directions.

Panel of Figure 3 — Figure 3: MNIST under geometric perturbations. Classification accuracy and adjusted TASTE score. Translations (left) leave both accuracy and adjusted TASTE residual unchanged, while rotations (right) induce a monotonic increase in the Stein signal aligned with performance degradation. Shaded regions show $\pm\,2$ standard deviations.

	Adversarial		CIFAR-10-C		CIFAR-10-P		OOD-benchmarks		Overall
Method	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$
MSP	0.3575	0.9656	0.6124	0.9109	0.5978	0.9210	0.8032	0.7878	0.5751 $\pm$ 0.0080	0.8521 $\pm$ 0.0161
ODIN	0.4895	0.9436	0.5499	0.9013	0.5478	0.9222	0.7644	0.7756	0.5776 $\pm$ 0.0084	0.9486 $\pm$ 0.0070
Mahalanobis	0.5819	0.8388	0.5880	0.9181	0.5620	0.9435	0.8188	0.7195	0.6060 $\pm$ 0.0084	0.8664 $\pm$ 0.0116
Energy	0.3550	0.9681	0.6256	0.8900	0.6110	0.9113	0.8520	0.6683	0.5871 $\pm$ 0.0081	0.8421 $\pm$ 0.0125
kNN+	0.4645	0.9591	0.6151	0.8977	0.5956	0.9223	0.8288	0.7023	0.6007 $\pm$ 0.0083	0.8526 $\pm$ 0.0123

Table 1: OOD detection performance on CIFAR-10 across heterogeneous shift regimes.

Method	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	ms/sample
	Adversarial		OOD-benchmarks		Overall		Runtime
MSP	0.2543	0.9551	0.8661	0.4546	0.4582 $\pm$ 0.0019	0.7883 $\pm$ 0.0039	5.95
ODIN	0.6475	0.8212	0.8356	0.4893	0.7102 $\pm$ 0.0026	0.7106 $\pm$ 0.0050	19.78
Mahalanobis	0.6785	0.7452	0.9264	0.1749	0.7611 $\pm$ 0.0024	0.5551 $\pm$ 0.0048	137.69
Energy	0.2242	0.9607	0.8971	0.3634	0.4485 $\pm$ 0.0018	0.7616 $\pm$ 0.0043	6.04
kNN+	0.6099	0.8089	0.9175	0.2501	0.7125 $\pm$ 0.0025	0.6226 $\pm$ 0.0044	95.73

Table 2: OOD detection performance on ImageNet-1K.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Experiment	Level	per_input	Score used	Further aggregation	Reported
Gaussian shifts	Dataset	False	$r_{f}(x)$	$\frac{1}{n}\sum_{i=1}^{n}r_{f}(x_{i})$ per shifted dataset	Shift-response plots
MNIST rotations and translations	Dataset	False	$r_{f}(x)$	$\frac{1}{n}\sum_{i=1}^{n}r_{f}(x_{i})$ per transformation and magnitude	Shift-response plots
CIFAR-10 detection	Datapoint	False	$r_{f}(x)$	None	AUROC, FPR95
ImageNet-1K: predicted class	Datapoint	False	$r_{f}(x)$	None	AUROC, FPR95
ImageNet-1K: TASTE ext	Datapoint	–	Classwise $d_{c}(x)$	$\sum_{c\in\mathcal{E}(x)}d_{c}(x)$ : largest and smallest $5\%$ per input	AUROC, FPR95
MVTec AD localisation	Coordinate	True	$r_{f,i}(x)$	Sum RGB-channel contributions per pixel	Pixel AUROC, AUPRO; maps thresholded at $q_{0.99}$

Appendix

Table 3: Scoring and aggregation protocols. Gaussian experiments use the scalar model output; TASTE ext uses all classwise softmax probabilities; the remaining experiments use the predicted-class softmax probability, holding the selected class fixed during differentiation. All scores and sums shown are signed. The per_input flag refers to Algorithm 1 ; “–” denotes the separate JVP construction in Appendix C.2 .

OOD Dataset	metric	MSP	Energy	ODIN	Mahalanobis	kNN+	TASTE
adversarial_autoattack_linf_2_255	AUROC	0.5019	0.4896	0.5393	0.5445	0.5276	0.4786
	FPR95	0.9611	0.9643	0.93	0.9371	0.9513	0.9571
adversarial_autoattack_linf_4_255	AUROC	0.2841	0.289	0.4851	0.505	0.3839	0.6543
	FPR95	0.9681	0.9711	0.9532	0.9213	0.9711	0.893
adversarial_autoattack_linf_8_255	AUROC	0.1557	0.1675	0.3319	0.7207	0.4341	0.8609
	FPR95	0.9688	0.9713	0.9718	0.513	0.9688	0.602

Appendix

Table 4: OOD detection results: adversarial

OOD Dataset	metric	MSP	Energy	ODIN	Mahalanobis	kNN+	TASTE
cifar10c_brightness	AUROC	0.5321	0.5389	0.4928	0.5356	0.543	0.5373
	FPR95	0.9454	0.9354	0.9562	0.9452	0.939	0.9335
cifar10c_contrast	AUROC	0.6945	0.7206	0.802	0.5841	0.6653	0.6846
	FPR95	0.86	0.781	0.5962	0.929	0.8666	0.7849
cifar10c_defocus_blur	AUROC	0.5801	0.5763	0.5864	0.5667	0.5847	0.5842
	FPR95	0.9308	0.9212	0.8546	0.9318	0.9186	0.8969

Appendix

Table 5: OOD detection results: CIFAR-10-C

OOD Dataset	metric	MSP	Energy	ODIN	Mahalanobis	kNN+	TASTE
cifar10p_brightness	AUROC	0.5497	0.5549	0.499	0.5488	0.5583	0.5557
	FPR95	0.9395	0.9357	0.9585	0.9415	0.934	0.9205
cifar10p_gaussian_noise	AUROC	0.5069	0.5107	0.4618	0.4965	0.5076	0.5107
	FPR95	0.9473	0.9454	0.9629	0.9559	0.9505	0.9456
cifar10p_motion_blur	AUROC	0.6294	0.6368	0.6134	0.6183	0.6446	0.6342
	FPR95	0.9197	0.9203	0.8495	0.9152	0.9095	0.8123

Appendix

Table 6: OOD detection results: CIFAR-10-P

OOD Dataset	metric	MSP	Energy	ODIN	Mahalanobis	kNN+	TASTE
iSUN	AUROC	0.7926	0.8561	0.682	0.8379	0.8203	0.7408
	FPR95	0.7848	0.6455	0.9396	0.6829	0.6995	0.5901
LSUN	AUROC	0.8208	0.8999	0.7047	0.88	0.8677	0.7666
	FPR95	0.7391	0.5161	0.9421	0.5912	0.5945	0.5364
Places365	AUROC	0.8005	0.854	0.7206	0.7943	0.8335	0.7746
	FPR95	0.7878	0.6461	0.86	0.7714	0.6894	0.5609

Appendix

Table 7: OOD detection results: OOD benchmarks

Panel of Figure 4 — Figure 4: Demonstration of per-pixel anomaly heatmaps for the MNIST-based task. Input digits are shown on the left, while the heatmaps are on the right. Note that both original and translated digits do not generate significant anomaly signal.

Panel of Figure 5 — Figure 5: Additional qualitative results. Each panel is a randomly chosen example for one of the MVTec AD categories. Within each panel there are four subplots: the original image (with corruption), detector’s heatmap, prediction overlay at $\alpha=0.01$ , prediction and ground truth (GT) overlay.

Panel of Figure 6 — Figure 6: Additional qualitative results. Each panel is a randomly chosen example for one of the MVTec AD categories. Within each panel there are four subplots: the original image (with corruption), detector’s heatmap, prediction overlay at $\alpha=0.01$ , prediction and ground truth (GT) overlay.

	iSUN		LSUN		Places365		SVHN		Textures		Average
Method	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$
TASTE	0.7603	0.9720	0.7820	0.9380	0.5927	0.9410	0.9477	0.1200	0.6396	0.8890	0.7445	0.7720
TASTE (low res)	0.7120	0.9460	0.7044	0.9480	0.5330	0.9510	0.9206	0.2050	0.5806	0.9490	0.6901	0.8190

Appendix

Table 8: Effect of score model degradation on OOD detection performance on ImageNet-1K. Pre-softmax logit of the predicted class is used as the probe.

	Adversarial		CIFAR-10-C		CIFAR-10-P		OOD-benchmarks		Overall
Method	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$	AUROC $\uparrow$	FPR95 $\downarrow$
TASTE	0.6144	0.865	0.6193	0.8691	0.5954	0.8611	0.7647	0.5677	0.6301	0.8207
TASTE-no-Laplacian	0.6264	0.8711	0.6098	0.8708	0.5902	0.8657	0.7597	0.5751	0.6253	0.8367
Score-only	0.5337	0.9437	0.6364	0.787	0.5789	0.8756	0.4452	0.8623	0.5783	0.8518

Appendix

Table 9: Operator construction ablation study - performance on CIFAR-10 across heterogeneous shift regimes. Results are averaged within each category. The dataset composition is identical to that used in the CIFAR-10 evaluation suite presented in Section 5.3 .

Figure 7: The effect of Laplacian estimation on performance of TASTE. Evaluation on MNIST under rotation-induced shifts using an AlexNet-style CNN shows monotone improvement as the quality of Laplacian estimation increases.

Panel of Figure 8 — Figure 8: Effect of diffusion time step on per-pixel anomaly detection. As the diffusion time step increases, ranking-based metrics (AUROC (Panel a), AUPRO (Panel b) ) decrease, while average precision (Panel c) improves, indicating a shift from global separability to more confident, sparse anomaly responses. Orange: standard area under curve statistics, blue: additionally restricting the False Positive Rate to below 30% and thus promoting early separation.

Figure 9: Model performance and Stein score as a function of test set corruption level. The methodology picks up on the presence of out-of-distribution input data for even very small corruption levels. Shaded regions: $\pm 2$ s.d., as approximate 95% confidence intervals.

Figure 10: Distribution of the per-sample TASTE residuals. The left panel shows evaluation on clean (MNIST) samples, while the right presents the scores generated from corrupted data. The considerable distribution shift of the per-sample TASTE residuals allows for per-input anomaly detection.

Figure 11: Stein score-based test power comparison as a function of test set corruption level $\alpha$ . The left panel presents results for two-sided tests, while a one-sided (arguably more powerful) test is presented on the right. Also, see Appendix C for power analysis based on a first-order operator

Figure 12: Directional evaluation of model sensitivity and OOD scoring under the first-order Stein operator ( 19 ). This figure should be compared with Figure 2 . The signal generated by the (analog of the) TASTE functional for the first-order Stein operator no longer aligns with performance.

Distribution Shift Detection

Distribution shift detection identifies when data or model behavior changes between reference and deployment settings.

See all 62 papers

TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability
May 14, 2026Krish Sharma, Omar Naim, Soumadeep Saha +3Representation Geometry LLM Pruning
ANITI, France · Meta · Apple, USA
STAMP: Predicting Out-of-Distribution Generalization without Target Data
Sep 26, 2026Md Kawsher Mahbub, Milon BiswasOOD Generalization
Department of Computer Science and Engineering NPI University of Bangladesh Manikganj, Bangladesh · Department of Computer & Information Sciences Towson University Towson, MD 21252
TASER: Task-Aware Stein Regularisation for Geometry-Driven Smoothness
May 28, 2026Michał Kozyra, Gesine ReinertDistribution Shift Robustness Neural Network Robustness
Department of Statistics, University of Oxford, United Kingdom

May 14, 2026J/K move · Enter open · S save

TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

Krish Sharma, Omar Naim, Soumadeep Saha, Vinija Jain+2

ANITI, France · Meta · Apple, USA

Representation Geometry LLM Pruning OOD Generalization Neural Network Pruning

Open paper AskPDF arXiv

Recent work has promoted task-aware layer pruning as a way to improve model performance on particular tasks, as shown by TALE. In this paper, we investigate when such improvements occur and why. We show first that, across controlled polynomial regression tasks and large language models, such pruning yields no benefit on in-distribution (ID) data but consistently improves out-of-distribution (OOD) accuracy. We further show empirically that OOD inputs induce layerwise norm and pairwise-distance profiles that deviate from the corresponding ID profiles. This leads to a geometric explanation of task-aware pruning: each task induces a task-adapted geometry, characterized empirically by the representation profiles observed on ID inputs. OOD inputs can introduce a distorted version of the task-adapted geometry. Task-aware pruning identifies layers that create or amplify this distortion; by removing them, it shifts OOD representational norms and pairwise distances toward those observed on the adapted distribution. This realigns OOD inputs with the model's task-adapted geometry and improves performance. We provide causal evidence through controlled distribution shifts and residual-scaling interventions, and demonstrate consistent behavior across model scales.

Abstract

Distribution Shift Detection

Explore similar work

TAPIOCA: Why Task- Aware Pruning Improves OOD model Capability

STAMP: Predicting Out-of-Distribution Generalization without Target Data

TASER: Task-Aware Stein Regularisation for Geometry-Driven Smoothness