Prototype-based Optimal Transport for Out-of-Distribution Detection

1 Introduction

Deep neural networks (DNNs) deployed in real-world scenarios often encounter out-of-distribution (OOD) inputs, such as inputs not belonging to one of the DNN's known classes. Ideally, reliable DNNs should be aware of what they do not know. However, they typically make overconfident predictions on OOD data [Nalisnick et al., 2018]. This notorious behavior undermines the credibility of DNNs and could pose risks to involved users, particularly in safety-critical applications like autonomous driving [Filos et al., 2020] and biometric authentication [Wang and Deng, 2021]. This gives rise to the importance of OOD detection, which identifies whether an input is OOD and enables conservative rejection or transferring decision-making to humans for better handling.

The representations of ID data within the same class tend to be gathered together after the training process, as shown in Figure 1(a). In contrast, the representations of OOD data are relatively far away from ID data, as they are not involved in the training process. In other words, the distributions of ID and OOD representations in the latent space exhibit a distinct separation. Therefore, we can expect that the distribution discrepancy between the representations of test inputs (i.e., a mixture of ID and OOD data) and pure ID data is primarily caused by the presence of OOD data. Such a distribution discrepancy motivates us to differentiate OOD data from test inputs by quantifying the individual contribution of each test input to the overall distribution discrepancy.

In this way, a critical question arises: how to measure the distribution discrepancy between test inputs and ID data, while quantifying the contribution of each test input? To this end, we utilize optimal transport (OT), a principled approach with rich geometric awareness for measuring the discrepancy between distributions. Concretely, OT aims to minimize the total transport cost between two distributions to measure the distribution discrepancy based on a predefined cost function (typically the geometric distance between samples). The smaller the total cost is, the closer the two distributions are. Since the total cost comprises the sum of transport costs between sample pairs, OT facilitates a fine-grained assessment of individual sample contributions to the overall discrepancy, making it particularly well-suited for OOD detection. Furthermore, the transport cost captures the geometric differences between ID and OOD representations in the latent space, providing a geometrically meaningful interpretation.

Based on the above intuition, in this paper, we propose a novel OOD detection method called POT that utilizes the Prototype-based Optimal Transport. Concretely, we first construct ID prototypes with the class-wise average of training sample representations to represent the distribution of ID data. We then apply OT between the representations of test inputs and the ID prototypes, obtaining a transport cost matrix where each entry indicates the transport cost between the corresponding pair of test input and prototype. As illustrated in Figure 1(b), the total transport cost, calculated by summing all matrix entries, reflects the overall distribution discrepancy. The transport cost from each test sample to all ID prototypes (i.e., the row sum), serves as a measure of individual contribution to the overall discrepancy, indicating the likelihood of being OOD. However, the task of OOD detection remains inadequately addressed due to the presence of OOD data with smaller distribution shifts. These OOD data lie closer to ID data in the latent space, rendering the transport costs to ID prototypes insufficient for detecting them. To tackle this issue, we propose generating virtual outliers to approximate the OOD region, particularly the areas near ID prototypes, using linear extrapolation between ID prototypes and the average representation of test inputs. By integrating the transport costs from test inputs to ID prototypes with the cost to virtual outliers, the detection of OOD samples with smaller distribution shifts could be emphasized, thereby enhancing the overall distinction between ID and OOD data.

Our key contributions are as follows: (1) We present a novel perspective for OOD detection by measuring distribution discrepancy and propose an effective detection method using prototype-based optimal transport. (2) Extensive experiments on various benchmark datasets demonstrate that our proposed method achieves state-of-the-art (SOTA) performance, outperforming 21 previous OOD detection methods. Moreover, in the scenarios where training data is unavailable, our method consistently beats the robust competitors by a margin of 22.5% in FPR95 on the CIFAR-100 benchmark.

2 RelatedWork

2.1 OODDetection

OOD detection has attracted growing research attention in recent years. Existing approaches can generally be categorized into two major lines:

(1) One line of work utilizes the outputs from pretrained models to design scoring functions for differentiating OOD samples. These post-hoc methods can be further divided into three subcategories. 1) The confidence-based methods [Sun et al., 2021; Song et al., 2022; Hendrycks et al., 2022; Wang et al., 2022b; Liu et al., 2023] adjusts model outputs to obtain the desired confidence, including maximum softmax probability [Hendrycks and Gimpel, 2017], energy [Liu et al., 2020], and generalized entropy [Liu et al., 2023]. 2) The density-based methods [Hendrycks et al., 2022; Sun and Li, 2022; Zhang et al., 2023c; Liu et al., 2024] identifies certain properties or patterns of ID data, such as neuron coverage [Liu et al., 2024], by learning the corresponding density functions, and the OOD samples that deviate from these properties or patterns tend to reside in low-density regions. 3) The distance-based methods [Lee et al., 2018; Ren et al., 2021; Sehwag et al., 2021; Sun et al., 2022] adopts distance metrics, e.g., Mahalanobis distance, between test input and ID samples or centroids differentiate OOD samples, under the assumption that OOD data lie relatively far away from ID data. Different from these works, we introduce a novel perspective for OOD detection on measuring distribution discrepancy. The most closely related work is dual divergence estimation (DDE) [Garg et al., 2023], which estimates the dual KL-Divergence between the test samples and ID samples. However, while DDE estimates divergence in a dual space optimized by DNNs and relies heavily on the quality of transformed sample representations, POT enables direct measurement of distribution discrepancy in the latent space.

(2) Another line of work focuses on altering models with training-time regularization to amplify the differences between OOD and ID samples [Hendrycks et al., 2019a; Ming et al., 2022; Zhang et al., 2023a; Wang et al., 2023; Lu et al., 2023]. For example, by incorporating a suitable loss, models are encouraged to produce predictions with uniform distributions [Hendrycks et al., 2019a] or higher energy [Liu et al., 2020] for OOD data. Building on this, some approaches investigate refining or synthesizing outliers to improve the model performance. For instance, [Ming et al., 2022] utilizes a posterior sampling-based technique to select the most informative OOD samples from a large outlier set, while [Wang et al., 2023] implicitly expands the auxiliary outliers by perturbing model parameters. However, the cost of model retraining can be prohibitive, especially when the parameter scale is large. Additionally, modifying models may also have side effects of degrading model performance on the original task. In contrast, this paper focus on post-hoc methods, which are easy to implement and generally applicable across different models. Such properties are highly practical for adopting OOD detection methods in real-world applications.

2.2 Optimal Transport

As a mathematical tool for comparing distributions, optimal transport (OT) has been successfully employed in diverse machine learning tasks, including domain adaptation [Turrisi et al., 2022], generative adversarial training [Arjovsky et al., 2017], object detection [Ge et al., 2021], and partial-label learning [Wang et al., 2022a]. The most related one to our work is [Lu et al., 2023], which also applies OT for the OOD detection problem. By assuming both an unlabeled and a labeled training set, [Lu et al., 2023] uses OT to guide the clustering of samples for label assignment to the unlabeled samples, thereby augmenting the training data and facilitating the model retraining for OOD detection. In contrast, our work operates in a post-hoc manner without the requirement of extra training data or model retraining.

3 The ProposedMethod

3.1 Problem Setting

In the context of supervised multi-class classification, we denote the data space as X and the corresponding label space as Y = {1, 2, · · · , C}. The training dataset with in-distribution (ID) samples Dtr = {(xi, yi)}ni=1 is sampled from the joint distribution PXY. The marginal distribution on X is denoted as P in X . The model trained on training data typically consists of a feature encoder g : X → Rd , mapping the input x ∈ X to a d-dimensional representation, and a linear classification layer f : Rd → RC , producing a logit vector containing classification confidence for each class. Given the test inputs Dte = {xj}mj=1 , the goal of OOD detection is to identify whether xj is out-of-distribution w.r.t P in X .

3.2 Prototype-based Optimal Transport for OOD Detection

Constructing Class Prototypes The key idea of our method is to use OT to measure the distribution discrepancy between test inputs and ID data while quantifying the individual contribution of each test input. A straightforward approach involves applying OT between test inputs and the training data. However, the standard OT is essentially a linear programming problem, suffering from cubic time complexity and incurring prohibitive computational cost when applied to largescale training sets [Peyr´e et al., 2019]. An alternative is to sample a smaller subset of the training set for efficiency, but this can result in missing classes, leading to mismatches in ID inputs. Such mismatches can exaggerate the contribution of ID inputs to the distribution discrepancy, thus incorrectly identifying them as OOD. To overcome this, we propose to characterize each class with a prototype and align test inputs to these prototypes. Specifically, given the training dataset Dtr = {(xi, yi)}ni=1 , we construct prototype for each class as the average representation for that class extracted from the feature encoder g:

where Nc is the number of samples in class c. The utilization of prototypes offers two key benefits: it reduces computational overhead by reducing the data scale and ensures all classes are represented, mitigating the risk of mismatches caused by missing classes.

OOD Detection Using Prototype-based Optimal Transport By formalizing an optimal transport problem between the representations of test inputs {zj = g(xj)}mj=1 and ID prototypes {ηi}Ci=1 , we search for the minimal transport cost that represents the distribution discrepancy, while subject to the mass conservation constraint:

where ⟨·, ·⟩F stands for the Frobenius dot-product. µ and ν are the probability simplexes of ID prototypes and test samples:

where pi = Ni ∑ni=1 Ni , qj = 1m , and δηi is the Dirac at position ηi. γ is a transport plan, essentially a joint probability matrix, with an entry γij describing the amount of probability mass transported from prototype ηi to test input zj. All feasible transport plans constitute the transportation polytope Π(µ, ν) [Cuturi, 2013]. E is the ground cost matrix, where the entry Eij denotes the point-to-point moving cost between ηi and zj, which is defined with the Euclidean distance as Eij = ||ηi - zj ||2.

To efficiently resolve Equation 2, we introduce the entropic regularization term H(γ) and write the optimization problem as:

where λ > 0 and H(γ) = ∑i,j γij(log γij - 1) [Peyr´e et al., 2019]. In this way, the optimal transport plan γ can be written as:

Here, a ∈ RC and b ∈ Rm are known as scaling variables. This formulation can be solved much faster using the Sinkhorn-Knopp algorithm [Cuturi, 2013]:

(6)

where ⊘ denotes element-wise division. Detailed derivations are provided in Appendix 6.1. With the λ-strong convexity [Peyr´e et al., 2019], the entropic regularized OT could be solved in quadratic time complexity O(nm), where n and m denote the number of data points in the two distributions,

respectively. Since the ID distribution is represented with a fixed number of prototypes, the prototype-based OT has linear time complexity O(Cm).

To evaluate whether an input sample xj ∈ Dte is OOD or not, we derive the cost of moving it to all ID prototypes Tj by decomposing the total transport cost:

As the total transport cost serves as a measure of distributional discrepancy between ID and test data, a higher transport cost T for a test sample indicates a greater deviation from the ID data, suggesting the sample is more likely to be OOD. We emphasize that, while the solution of OT is the transport plan γ, which is often the focus in many applications [Caron et al., 2020; Wang et al., 2022a], our concentration lies on the transport costs between sample pairs. Specifically, it is the product of transport plan and ground cost, serving as a desirable and interpretable measure for OOD detection.

Data Augmentation via Linear Extrapolation As described above, the transport cost can serve as a measure for OOD detection. However, relying solely on the cost to ID prototypes is insufficient for discerning OOD data with smaller distribution shifts from ID data. It is because such OOD data are located closer to ID data in the latent space and tend to incur lower transport costs to ID prototypes. Meanwhile, the mass conservation constraint inherent in OT may exacerbate this issue by enforcing transportation between OOD data and prototypes when ID inputs are sparse. To address this issue, we propose to generate virtual outliers to approximate the OOD region using linear extrapolation. By integrating the transport costs from test inputs to both virtual outliers and ID prototypes, we introduce a contrastive transport cost, which enhances the detection performance, particularly for the OOD inputs with smaller distribution shifts.

Given two representations zi and zj, the linear extrapolation between them is defined as:

Instead of generating outliers aimlessly by enumerating available sample pairs, we construct virtual outliers P ∗ by combining the prototypes P and the average of test inputsM:

where ω > 1 ensures that the generated points P ∗ lies beyond M. The underlying intuition is that M can be expressed as a linear interpolation of the average representations of test ID samples Min and test OOD samples Mout:

where Nin and Nout denote the number of test ID samples and total test samples, respectively. As illustrated in Figure 2, the point M resides between ID and OOD data, guiding the direction for outlier generation in response to distribution shifts. As parameter ω increases, the generated virtual outliers progressively move away from the ID prototypes towards the OOD region in the latent space. By choosing an appropriate parameter, we can control the location of generated virtual outliers to emphasize the detection of OOD inputs with smaller distributions shifts. In contrast to the method [Zhu et al., 2023] that conducts informative extrapolation to synthesize numerous outliers during training with assumed auxiliary outliers, the linear extrapolation is a lightweight operation that does not require training or auxiliary outliers.

Likewise, after generating the virtual outliers P ∗ , we apply the entropic regularized OT between P ∗ and the test inputs to obtain the corresponding transport cost T ∗ . By taking the difference of transport costs from test inputs to the ID prototypes and to the virtual outliers, we derive the contrastive transport cost T -T ∗ as the final OOD score. Given the opposite trends of T and T ∗ in indicating whether a sample is ID or OOD, a higher contrastive transport cost denotes a higher likelihood of the test input being OOD.

4 Experiments

In this section, we preform extensive experiments over OOD detection benchmarks. All the experimental setup adheres to the latest version of OpenOOD, an open repository for benchmarking generalized OOD detection 1 [Zhang et al., 2023b].

4.1 Common Setup

Datasets We assess the performance of our proposed POT on the widely used CIFAR-100 benchmark and ImageNet1k benchmark, which regard CIFAR-100 and ImageNet1k as ID datasets, respectively. For CIFAR-100, we use the standard split, with 50,000 training images and 10,000 test images. For ImageNet-1k, we utilize 50,000 images from the validation set as ID test set. For each benchmark, OpenOOD splits OOD data to Far-OOD and NearOOD based on their degrees of the semantic similarity with ID data. Specifically, the CIFAR-100 benchmark utilizes four Far-OOD datasets: MNIST [Deng, 2012], SVHN [Netzer et al., 2011], Textures [Cimpoi et al., 2014], and Places365 [Zhou et al., 2018], along with two Near-OOD datasets: CIFAR-10 and Tiny Imagenet [Le and Yang, 2015]. For the large-scale ImageNet-1k benchmark, it considers iNaturalist [Horn et al., 2018], Textures [Cimpoi et al., 2014], and OpenImage-O [Wang et al., 2022b] as Far-OOD datasets. In terms of Near-OOD datasets, SSB-hard [Vaze et al., 2022] and NINCO [Bitterwolf et al., 2023] are included.

1 https://github.com/Jingkang50/OpenOOD

Models For the CIFAR-100 benchmark, we utilize ResNet18 [He et al., 2016] as the model backbone, which is trained on the ID training samples for 100 epochs. We evaluate OOD detection methods over three checkpoints. For the ImageNet-1k benchmark, we employ ResNet50 and ViT-b16 [Dosovitskiy and et al, 2021] models pretrained on ImageNet-1k and use the official checkpoints from PyTorch. For more training details, please refer to OpenOOD.

Baselines Since POT performs in a post-hoc manner, we primarily compare against 21 post-hoc OOD detection methods, including OpenMax [Bendale and Boult, 2016], MSP [Hendrycks and Gimpel, 2017], ODIN [Liang et al., 2018], MDS [Lee et al., 2018], MDSEns [Lee et al., 2018], RMDS [Ren et al., 2021], Gram [Sastry and Oore, 2020], EBO [Liu et al., 2020], GradNorm [Huang et al., 2021], ReAct [Sun et al., 2021], MLS [Hendrycks et al., 2022], KLM [Hendrycks et al., 2022], VIM [Wang et al., 2022b], KNN [Sun et al., 2022], DICE [Sun and Li, 2022], RankFeat [Song et al., 2022], ASH [Djurisic et al., 2023], SHE [Zhang et al., 2023c], GEN [Liu et al., 2023], DDE [Garg et al., 2023], NAC-UE [Liu et al., 2024]. The results for DDE are reproduced using the official codebase 2 while the results for the remaining methods are sourced from the implementations in OpenOOD.

Hyperparameter tuning In line with OpenOOD, we use ID and OOD validation sets for hyperparameter selection. Specifically, for the CIFAR-100 benchmark, 1,000 images are held out from the ID test set as the ID validation set, and 1,000 images spanning 20 categories from Tiny ImageNet [Le and Yang, 2015] are reserved as the OOD validation set. For the ImageNet-1k benchmark, 5,000 images from the ID test set and 1,763 images from OpenImage-O are held out for the ID and OOD validation sets, respectively. Please note that all the validation samples are disjoint with the test samples. The validation sets are used to tune hyperparameters, including the entropic regularization coefficient λ and the linear extrapolation parameter ω for POT. Please refer to the Appendix 6.2

2 https://github.com/morganstanley/MSML/tree/main/papers/ OOD Detection via Dual Divergence Estimation

for more details.

Evaluation metrics We report the following widely adopted metrics: (1) area under the receiver operating characteristic curve (AUROC); (2) false positive rate of OOD samples when the true positive rate of ID samples is equal to 95% (FPR95).

Implementation details During the test phase, the test set comprises samples from both the ID test set and OOD dataset. We consider the assumption of having access to the entire test set to be overly restrictive. To this end, we relax the assumption by allowing test inputs arriving in batches, where random batch division is applied to the test set. The default test batch size is set to 512 and we also include ablation studies with varying batch sizes. For each test batch, we employ POT to calculate transport cost scores for the samples. After aggregating the scores across all test samples, we calculate the evaluation metrics for comparison.

4.2 Empirical Results and Analysis

Main results The results for Far-OOD and Near-OOD detection on the CIFAR-100 benchmark are presented in Table 1 and Table 2, respectively. Our proposed POT consistently obtains either the best or second-best results across all datasets and OOD detection metrics. Specifically, on the FarOOD track, POT achieves significant reductions in average FPR95, with decreases of 25.6% and 8.87% compared to the previous leading baselines NAC-UE and DDE, respectively. While Near-OOD samples are considered more intractable to detect due to their similarity in semantic and style with ID samples, POT demonstrates an even greater performance advantage in detecting them, as shown in Table 2. For instance, POT surpasses the next best method, GEN, by 9.96% in average FPR95 and 5.22% in average AUROC. The results of the large-scale ImageNet-1k benchmark are shown in Table 3, where only the AUROC values are reported due to space limitation. As can be seen, POT also consistently outperforms all baseline methods in terms of average performance across different backbones and OOD datasets.

Integration with training methods In the other line of work for OOD detection, training methods employ retraining strategies with training-time regularization to provide a modified model. A important property of post-hoc methods is that they are applicable to different model architectures and training losses. To this end, we examine the performance of post-hoc methods when integrated with established training methods. We evaluate on the Far-OOD track of the ImageNet-1k benchmark using ResNet-50 as the backbone, comparing POT with ASH and NAC-UE, which have achieved top results on this benchmark (see Table 3). Consistent with the experimental setup of NAC-UE, we utilize three training schemes: ConfBranch [DeVries and Taylor, 2018], RotPred [Hendrycks et al., 2019b], and GODIN [Hsu et al., 2020], with softmax cross-entropy (CE) loss as the default training scheme for comparison. We report the average results in Table 4, where Baseline refers to the detection method employed in the original paper and ∆ denotes the subtraction results between the default CE scheme and other training schemes. According to the results, POT has remarkable improvements on the baseline methods, while outperforming ASH and NAC-UE. Importantly, POT maintains stable performance across three training schemes, whereas ASH and NAC-UE both exhibit notable performance degradation compared to the default CE training scheme. Such results demonstrates that POT is generic to be seamlessly integrated with different training methods.

What if training data is unavailable Although existing posthoc methods can be applied to pre-trained models without cumbersome retraining, many require access to at least a portion of training data [Liu et al., 2023]. Considering the scenarios where training data is unavailable, such as commercial data involving privacy, some approaches further explore OOD detection without the requirement for training data such as ASH and GEN. To make POT applicable to such scenarios, we draw inspiration from the work [Tanwisuth et al., 2021] and construct the class prototypes with the neural network weights W ∈ Rd×C of the classification layer f. Each column of the weight matrix W corresponds to a d-dimensional class prototype. The underlying idea is that the process of learning class prototypes with learnable parameters in the latent space, is closely similar to the training process of the linear classification layer. In other words, obtaining prototypes in this way does not require any training data and additional parameters. To verify the effectiveness of POT in such scenarios, we compare it to the baseline methods not requiring access to the training dataset. We report the average AUROC for both Far-OOD and Near-OOD datasets in Table 5. The results show that POT continues to outperform all the competitors across all metrics. Notably, compared to GEN, POT significantly reduces the average FPR95 by 35.78% on the Far-OOD datasets and by 9.21% on the Near-OOD datasets of the CIFAR-100 benchmark.

4.3 Ablation Study and Analysis

Ablation on virtual outliers We compare the OOD detection performance of POT with and without virtual outliers, which are generated to approximate the OOD region using linear representation extrapolation. Figure 3 displays the results, where we visualize the data distribution in the CIFAR100 and ImageNet-1k benchmarks. Using the virtual outliers leads to clearer separability between ID and OOD samples, whereas the transport cost score without virtual outliers exhibits larger overlap. The results show that the introduction of virtual outliers enhances the distinguishability between ID and OOD, resulting in more effective OOD detection.

Parameter analysis We ablate along individual parameters with ImageNet-1k as ID data, shown in Figure 6.4 in Appendix 6.4 , where the ViT-b16 is utilized as the backbone for analysis. In Figure 6.4 (a), we find that increasing the test batch size is beneficial for OOD detection. We hypothesize that this improvement is due to larger batch sizes providing a better empirical approximation of the true test data distribution, thereby enabling more accurate measurement of distributional discrepancies and improving detection performance. Although POT faces a performance degradation with the decrease of test batch size, its performance of lower batch sizes still hold superiority over the baseline methods. For instance, POT already achieves 93.18% average AUROC with the test batch size of 32, which outperforms the competitive rival NAC-UE (see Table 1). In Figure 6.4 (b), we can observe that as the regularization coefficient λ increases, AUROC undergoes an ascending interval followed by a decrease interval, ultimately leading to convergence. This convergence occurs because the entropy term becomes increasingly influential in the optimization of OT as λ increases, eventually leading to a uniform transport plan with maximum entropy, thereby stabilizing OOD detection performance. In Figure 6.4 (c), we find that POT works better with a moderate ω. This is intuitive as a lower ω may lead the generated virtual outliers close to ID samples, while a higher ω can generate virtual outliers far away both ID and OOD samples, leading the transport costs to virtual outliers indistinguishable.

Penultimate layer vs. logit layer In this paper, we follow the convention in feature-based methods by using the representations from the penultimate layer of neural network, as it is believed to preserve more information than the output from the top layer, also referred to as logits. To investigate the impact of layer selection, we provide evaluation on POT using representations from the penultimate layer or logits, with the average results presented in Table 6. From the results, using the representations from the penultimate layer achieves better performance than using the logits.

5 Conclusion

In this paper, we tackle the problem of OOD detection from a new perspective of measuring distribution discrepancy and quantifying the individual contribution of each test input. To this end, we propose a novel method named POT, which utilizes the prototype-based OT to assess the discrepancy between test inputs and ID prototypes and use the obtaining transport costs for OOD detection. By generating virtual outliers to approximate the OOD region, we combine the transport costs to ID prototypes with the costs to virtual outliers, resulting in a more effective contrastive transport cost for identifying OOD inputs. Experimental results demonstrate that POT achieves better performance than 21 current methods. Moreover, we show that POT is pluggable with existing training methods for OOD detection and is applicable to the scenarios where the training data is unavailable, highlighting its generic nature and high practicality.

Acknowledgements

This work was supported in part by the National Natural Science Foundation of China under Grant 62472400, Grant 62072428, Grant 62271465, and in part by the Suzhou Basic Research Program under Grant SYG202338. Xike Xie is the corresponding author.