Partial Optimal Transport Based Out-of-Distribution Detection for Open-Set Semi-Supervised Learning

1 Introduction

Semi-supervised learning (SSL) is a branch of machine learning that leverages both labeled and unlabeled data to improve the performance of learning tasks [Lee and others, 2013; Laine and Aila, 2016; Tarvainen and Valpola, 2017; Berthelot et al., 2019; Xie et al., 2020; Sohn et al., 2020]. A common assumption of SSL is that the label spaces of labeled and unlabeled data are same. However, in open-world settings, a large volume of unlabeled data need to be collected, which often contain novel categories that are unseen in the labeled data. This gives rise to a more practical setting called Openset Semi-supervised Learning (OSSL) [Yu et al., 2020], in which known categories samples, i.e., inliers, should be classified into correct classes, while unknown categories samples, i.e., outliers, should be identified. However, existing SSL approaches fall short in OSSL settings. For example, a common step in SSL is pseudo-labeling [Yang et al., 2023a], which often assigns inappropriate pseudo-labels to outliers, thus posing challenges to the generalization over inliers and also the detection of outliers.

Identifying outliers, referred to as out-of-distribution (OOD) samples, stands as a pivotal challenge in the realm of OSSL. To tackle the challenge, one intuitive strategy is to integrate existing OOD detection methods into OSSL training. For example, UASD [Chen et al., 2020] and SAFE STUDENT [He et al., 2022] proposed to use maximum softmax prediction probability [Hendrycks and Gimpel, 2016] and energy-discrepancy score [Liu et al., 2020], respectively, to detect outliers. TOSSL [Ma et al., 2023] transformed the problem into a closed-set problem, by classifying OOD samples into a new category via distance-based methods [Lee et al., 2018; Sun et al., 2022]. However, these OOD detection methods are designed for the testing phase, relying on models trained over in-distribution (ID) data. In contrast, OSSL asks for OOD detection during the training phase. The latter is more challenging, because: 1) the training data comprises a mixture of ID and OOD data; 2) the training-phase detection implies working with an under-developed model. An alternative strategy relies on model predictions for OOD detection. For instance, MTCF [Yu et al., 2020] proposed to filter ID samples using noisy label learning, and DS3L [Guo et al., 2020] designed a weighting function for assigning weights to input samples, which is jointly trained with network parameters through bi-level optimization. It is evident that the solely dependence on neural network predictions potentially degrades the performance of OOD detection, due to its overfitting on OOD samples.

In response to these deficiencies, a current trend involves employing multiple one-vs-all classifiers to distinguish if unlabeled samples belong to a known class, designating those not belonging to any specific class as OOD samples. In this context, OpenMatch [Saito et al., 2021] introduced an open-set consistency regularization, while T2T [Huang et al., 2021] used self-supervision and cross-modal matching, SSB [Fan et al., 2023] utilized non-linear transformations and pseudo-negative mining, OSP [Wang et al., 2023] proposed aliasing OOD matching and soft orthogonality regularization, to enhance one-vs-all classifiers. More recently, IOMatch [Li et al., 2023] jointly considered inliers and outliers for scenarios with extremely scarce labels. Nevertheless, these methods fall short in OSSL scenarios with numerous classification categories due to the usage of K classifiers. It results in the overall error rate being computed as 1 - (1 - a)K , where a represents the error rate of each individual classifier, ultimately degrading the performance.

To this end, we introduce a novel binary OOD detector rooted in optimal transport (OT) theory, operating directly in the feature space, so that: 1) it gets rid of the reliance on the network output during the training phase, and thus avoids the overfitting deficiency in MTCF and DS3L; 2) it obviates the need for training numerous binary classifiers in one-vsall strategies, such as T2T, OpenMatch, and IOMatch; 3) our method reallocates weights for unlabeled samples to align the unlabeled distribution with the labeled distribution, thus offering better interpretability for discerning OOD samples; 4) it improves the performance of both closed-set classification and open-set detection within OSSL.

The idea of our proposal originates in tackling the problem of distribution mismatch [Bickel et al., 2009], the fundamental technical challenge of OSSL, where the labeled and unlabeled data may originate from different distributions. It is known that the OT theory [Villani, 2008; Courty et al., 2014] is a mathematical tool for quantifying distribution discrepancies, providing an effective way of measuring the distribution mismatch. However, conventional OT operates under the assumption that both distributions process identical total mass, requiring complete transportation of all mass from one distribution to another. This assumption becomes problematic in the context of OSSL, where only a subset of the unlabeled data aligns with the labeled data, and the rest consists of outliers. In such cases, the assumption is untenable. Therefore, we turn to partial optimal transport (POT) [Caffarelli and McCann, 2010], a variant of OT, which allows the selective transport of a fraction of mass from one distribution to another. And yet, to the best of our knowledge, POT in OSSL remains unexplored, despite its potential utility in OSSL.

In this paper, we first formulate the outlier detection problem in OSSL as a POT problem. We seek the partial optimal transport plan between labeled and unlabeled distributions. With POT theory, we devise a mass score function (MSF) to measure the likelihood of unlabeled samples being outliers based on the mass transported between the two distributions. The MSF is calculated through the optimal transport plan of POT, where the MSF of an unlabeled sample is determined by the sum of its mass transport to the labeled distribution. To train the model in an end-to-end manner, we propose a novel OOD loss. This loss function involves predicting the input using a specially designed output layer referred to as the OODhead, while the target value is the MSF generated by the POT

process. Moreover, this approach enables us to adapt the offthe-shelf SSL methods (e.g. FixMatch) into OSSL settings via multi-task learning.

Extensive experiments across various datasets and settings consistently show the superior performance of our method in the realm of OSSL. Notably, our approach effectively addresses the performance decline witnessed in current methods when confronted with intricate scenarios featuring a substantial number of classification categories. Specifically, when applied to CIFAR-100 with 80 known classes, our method has demonstrated a remarkable improvement, achieving a noteworthy increase of over 7% in closed-set classification accuracy. The main contributions are summarized as follows.

• We directly address the distribution mismatch challenge using OT theory, modeling OOD detection in OSSL as a POT problem. This approach enhances the interpretative mechanism for identifying OOD samples.
• We introduce a novel mass score function (MSF) as a binary OOD detector. Simultaneously, our OOD detector can be integrated with off-the-shelf SSL methods through multi-task learning.
• Our method consistently surpasses current approaches, effectively addressing performance decline in complex scenarios with numerous classification categories.

2 Problem Setup

2.1 ProblemDefinition of OSSL

In OSSL, the labeled data consists exclusively of ID samples, while the unlabeled data and the test data include a mixture of ID and OOD samples. The goal of OSSL is to obtain a model capable of accurately classifying ID samples into their corresponding classes and distinguish OOD samples at test time.

In SSL, a widely used approach is to sample both labeled and unlabeled batches together, combining explicit supervision from labeled samples with the inherent information in unlabeled data. At each training iteration, SSL methods typically sample n labeled samples and m unlabeled samples, where each sample can be represented by a d-dimensional feature vector after undergoing the network's feature extractor. Then, the samples of labeled and unlabeled space constitute two d-dimensional point clouds, which can be represented by L = {li}ni=1 and U = {ui}mi=1 , respectively. We can represent the two point clouds of samples in the form of probability distributions, L and U, indicating the labeled and unlabeled distributions, where δ represents the Dirac function.

As shown in the right half of Figure 1, in OSSL, there exists a mismatch between L and U. Consequently, the unlabeled distribution can be represented as the union of UID and UOOD, where UID aligns with L, while UOOD contrasts with it. Effectively tackling the challenges posed by OSSL necessitates addressing this inherent mismatch between L and U during the training stage.

2.2 OT Preliminaries

Optimal transport (OT), as discussed in [Villani, 2021], is to seek an optimal transport plan between two measures at a minimal cost, resulting in the Wasserstein distance metric, as explained in [Figalli and Glaudo, 2021], providing a geometric way for aligning probability measures.

OT proves to be versatile in addressing distribution shift challenges. [Courty et al., 2014] introduced a domain adaptation method using regularized optimal transport to align source and target domains while preserving their structures. [Feng et al., 2023] introduced an OT-Filter to filter label noise under distribution shift. For open-set domain adaptation, both [Xu et al., 2020] and [Yang et al., 2023b] modeled with partial optimal transport. [Xu et al., 2020] regulated partial optimal transport using the mean cost of transport, while [Yang et al., 2023b] estimated the transport ratio. [Chapel et al., 2020] proposed a PU-learning method based on partial optimal transport, and [Lu et al., 2023] applied OT to the semantically coherent out-of-distribution detection task. However, as far as we know, the potential utility of OT in OSSL has not been investigated yet, despite its possible benefits in OSSL.

In the context of OSSL, a notable challenge arises from the distribution mismatch between the labeled distribution L and the unlabeled distribution U. An insightful approach involves exploring the application of optimal transport theory to detect samples in U that are not aligned with L, essentially UOOD.

When addressing the distribution mismatch challenge in OSSL through OT, the initial step involves defining the optimal transport problem between the labeled distribution L and the unlabeled distribution U. Then, in the context of the optimal transport problem, an admissible coupling represents a joint probability distribution over two measures and must adhere to constraints on marginal distributions, ensuring complete mass transport. The set of all admissible couplings Π(L, U) between the two distributions is given as follows.

where T is a coupling matrix, and Tij describes the amount of mass at li transported to uj. Given a transport cost matrix C, where the element Cij represents the cost of transporting one unit of mass from li to uj, optimal transport addresses the problem of transporting L toward U with minimal cost. A coupling matrix that achieves the minimal cost is called the optimal transport plan and the minimal cost is called the optimal transport distance. The OT problem can be described as follows.

However, when dealing with the obtained population-wise metric, i.e., the transport distance between L and U, it is impractical to identify an individual OOD sample using OT distance. Therefore, it becomes imperative to identify OOD samples by considering the perspective of the optimal transport plan represented by T.

2.3 POT for OSSL

Identifying OOD samples through the optimal transport plan presents a notable challenge. The conventional OT assumes that the two distributions have the same total probability mass, i.e., ∥L∥1 = ∥U∥1 , and that all the mass of an unlabeled sample must be transported. Although this assumption can be well applied to the task involving identical distributions or domain adaptation, it encounters challenges in OSSL settings, where only a portion of samples in U aligns with L. However, each unlabeled sample is transported to specific labeled samples, consequently disrupting the efficient transport of aligned samples from U to L. At the same time, every unlabeled sample carries an identical transport mass, posing a challenge in identifying OOD samples through the transport plan.

A straightforward way is to transport the mass of ID samples between U and L, while maintaining the mass of OOD samples within U. This gives rise to the concept of POT, which allows the transport of a portion of the mass. The left half of Figure 1 offers a simulated example of OOD detection in POT, where the linewidth indicates the amount of transport mass between two samples. It conveys that the transport mass between samples of L and UID significantly outweighs that of L and UOOD. So, with POT, it holds good potentials for identifying individual OOD samples.

Formally, the focus of the POT problem lies in transporting a fraction 0 ≤ s ≤ min(∥L∥1 , ∥U∥1 ) of the mass, while simultaneously minimizing the associated cost. The set of admissible couplings, Πs (L, U) in this case, is defined as:

Therefore, the POT problem can be described as:

Moving forward, our goal is to tackle the OSSL problem using the principles of partial optimal transport.

3 Method

3.1 Formulating POT in OSSL

In this section, our primary focus is on formalizing the OOD detection task in OSSL as a POT problem. We aim to use the labeled distribution as a baseline distribution to detect outliers. Therefore, we redefine the labeled distribution L and k-fold unlabeled distribution U as follows.

where li and ui represent the feature of the i-th sample in labeled and k-fold unlabeled distributions, respectively. The introduced parameter k denotes the redundant mass, which enables transporting k-fold mass from each sample in U. In this way, we represent the resulting problem as a POT problem with the following feasible set.

Incorporating equation constraints into the labeled distribution enhances the interpretability of our approach. As interpreted by Equation 7, the problem is to redistribute weights among unlabeled samples, aligning the unlabeled distribution with the labeled distribution. Moreover, the concise constraints specified in Equation 7 expedites the convergence of POT calculations.

In tackling the POT problem, the discrete optimal transport formulation, in essence, is a convex optimization problem, more precisely, a linear programming problem. Unfortunately, this linear programming problem suffers from a cubic computing complexity. To mitigate this, one way is to leverage entropic regularization, as outlined in [Cuturi, 2013], which is formulated as follows.

where ε > 0 is the regularization coefficient, and H(T) is the entropic regularization term. According to [Benamou et al., 2015], we can recast the regularized OT problem in the language of Kullback-Leibler projections:

The set of feasible solutions, denoted as C in the regularized POT context, is the intersection of two convex subspaces, i.e., C1 ∩ C2, where

Obviously, the above formulation of POT for OSSL can be viewed as a convex optimization problem. Following with previous works ([Dykstra, 1983; Bauschke and Lewis, 2000]) for solving POT, we conclude our computing framework as Algorithm 1.

3.2 Mass Score Function: A POT Based Binary OOD Detector

In Section 3.1, we have formulated POT between probability distributions defined over labeled and unlabeled datasets.

Algorithm 1 POT

Require: labeled distribution L, k-fold unlabeled distribution U, cost matrix C, regularization ε, and error δ.

This formulation prompts the need to create a robust scoring function for assessing the probability of a sample being classified as an OOD sample. To tackle this challenge, we introduce a novel mass score function (MSF) derived from the optimal transport plan in POT. The MSF aggregates the mass transferred from unlabeled samples to labeled ones and functions as a binary OOD detector.

Next, we investigate the mechanism of MSF and explain why POT proves to be effective for OOD detection. The key idea is to leverage the discrepancy between labeled distribution L and the k-fold unlabeled distribution U to identify the portion within U that aligns with L. As shown in Equation 6, we treat L as the reference distribution, where all samples carry identical weight (i.e., undergo the same mass transportation), assured by equality constraints. In contrast, samples in U may have different mass transportation, falling within the range of [0, km ], as guaranteed by inequality constraints.

The cost of transporting unit mass between ID samples, which includes both labeled and unlabeled samples, is relatively low. In line with the principle of POT, only a fraction of the mass within the unlabeled distribution can be feasibly transported. To minimize the overall transportation cost, ID samples in unlabeled dataset obtain a higher allocation of transportation mass. This characteristic contributes to the identification of OOD samples based on the transported mass of each individual sample. Alternatively, considering this from a different perspective, the MSF adjusts the weights for unlabeled samples to align the unlabeled distribution with the labeled distribution. Samples that do not conform to the labeled distribution, namely OOD samples, naturally receive smaller weight assignments.

Thus, in order to derive the mass score function, we initially employ the cosine distance as the transport cost between labeled and unlabeled samples. Subsequently, we use Algorithm 1 to obtain the optimal transport plan. At last, we calculate MSF and scale it to the range [0, 1] for training. The procedure is detailed in Algorithm 2. Following the procedure, we can derive a binary OOD detector based on POT.

3.3 OSSL Training with Proposed OODDetector

In Section 3.2, we present a novel OOD detection method using POT between L and U. While this approach is effective during the training stage for calculating the OOD score for

Algorithm 2 Mass Score Function

Require: Feature of labeled and k-fold unlabeled distributionL and U, entropy regularization coefficient ε, and mass redundancy parameter k.

Distribution: L = ∑ni=1 1n δli and U = ∑mi=1 km δui

Cost: calculate cosine distance matrix M = Cosine(L, U)

OT Plan: calculate optimal transport plan using algorithm1: T = P OT (L, U , M, ε)

OOD Score: calculate optimal transport mass as OOD score for U: ScoreU = T⊤1n

return Scaled ScoreU

each unlabeled sample, challenges arise when applying POT on the test data, since there are no corresponding labeled samples available for reference. To address this problem, we investigate the OOD-head, an additional output layer, to retain the capability of detecting OOD samples, supported by POT.

In particular, we denote the score produced by the OODhead as PredU and the value of MSF predicted by POT as ScoreU. At the training stage, the OOD-head is expected to produce scores close to the values of MSF for unlabeled samples. To this end, we resort to the following mean square error (MSE) loss.

Essentially, the OSSL task can be viewed as a combination of a semi-supervised classification task and an OOD detection task. We therefore formulate the OSSL problem as a multitask learning problem, which is implemented by only using a simple multi-layer feedforward neural network [Caruana, 1997]. More, the two tasks share the same feature extractor and use different prediction heads. By doing this, our method can be integrated with various semi-supervised methods 1 .

In addition, we outline two strategies that we have implemented to further enhance the performance of our OOD-

1 We use supervised and unsupervised losses from FixMatch to validate our approach. Additionally, in the ablation experiment, we apply our method to various other SSL techniques.

Algorithm 3 OSSL Computing Framework

Require: Feature extractor f, prediction head g, and OOD head h. Weight of unsupervised loss λu, weight of OOD loss λOOD. Weak data augmentation α and strong data augmentation A. Entropy regularization parameter ε, and mass redundancy parameter k. Sample a batch of labeled data DL and a batch of unlabeled data DU

Feature Extractor:

Feature of the labeled batch L = f(DL) ∈ Rn×d

Feature of weakly augmented unlabeled batch U = f (α(DU )) ∈ Rm×d

OOD Score:

OOD score of unlabeled batch using Algorithm2:

ScoreU = POD(L, U, ε, k)

Set the OOD score of labeled samples: ScoreL = 1n

Score = CONCAT(ScoreU , ScoreL)

Predict OOD Score:

The score of labeled batch: PredL = h(f(DL))

The score of strongly augmented unlabeled batch:

PredU = h(f (A(DU )))

Pred = CONCAT(PredU , PredL)

Calculate Loss:

calculate supervised loss of FixMatch Lx, unsupervised loss of FixMatch Lu, and LOOD using (11)

loss = Lx + λuLu + λOODLOOD

head. Firstly, we utilize both labeled and unlabeled samples for OOD-head training, as opposed to solely relying on unlabeled samples. Leveraging supervision information from ID labeled samples, we therefore set the OOD score of each labeled sample as 1. Second, we use the same data augmentation strategy with FixMatch for the OOD detection loss LOOD to improve the generalization capability of the OODhead. Using POT, we obtain the MSF for unlabeled samples based on the features of weakly-augmented versions of unlabeled samples. Further, we use the OOD loss to optimize the predicted OOD score for strongly-augmented versions of unlabeled samples. The overall computing framework is described in Algorithm 3.

4 Experimental Results

Metrics. To evaluate the performance of OSSL methods, we employ closed-set classification accuracy to test the performance concerning the known classes. We use AUROC to assess the model's open-set classification ability in distinguishing between inliers and outliers.

Baselines. In terms of OSSL baselines, we evaluate our approach in comparison to various existing methods, including FixMatch, MTCF, T2T, OpenMatch, and IOMatch. Additionally, since our method is an extension of SSL, we also include FixMatch, a powerful and widely adopted method in SSL, as a baseline for reference. To assess the potential of FixMatch for OOD detection, we employ the maximum softmax prediction probability as the score function at test time.

4.1 Performance Evaluation

Results on CIFAR10 and CIFAR100. We evaluate the performance of POT in comparison to baselines on the widely used benchmark datasets for SSL, namely CIFAR10 and CIFAR100. In these experiments, we adapt a randomly initialized Wide ResNet-28-2 [Zagoruyko and Komodakis, 2016] with 1.5Mparameters, in consistency with existing works. In particular, for CIFAR10, we divide it into 6 known classes and 4 unknown classes. For CIFAR100, we consider two settings: one with 80 known classes and 20 unknown classes, and another with 55 known classes and 45 unknown classes, organized according to superclasses. Note that we use the same set of hyper-parameters for all experiments 2 .

Tables 1 and 2 report the results on CIFAR10 and CIFAR100, respectively. On the relatively simple CIFAR-10

2 {µ = 2, B = 64, Ne = 512, Ni = 1024, ε = 0.05, k = 2, λOOD = 0.01}. B indicates the batch size and µ is the relative size of batch size for unlabeled data. Ne indicates the total number of training epochs and Ni is the number of iterations per epoch.

dataset, our approach achieves state-of-the-art (SOTA) performance in terms of both closed-set accuracy ACC and openset quality AUROC. Upon transitioning to the more challenging CIFAR-100 dataset, our method excels, dominating other OSSL competitors in both closed-set accuracy and AUROC. Notably, our approach showcases a noteworthy increase of more than 5% in closed-set accuracy. This substantiates the efficacy of our OOD binary detector, particularly in challenging and complex scenarios.

Results on ImageNet-30. We further evaluate our proposal in more complex and challenging settings. In particular, we use ImageNet-30 dataset [Hendrycks and Gimpel, 2016], a subset of ImageNet, comprising 30 classes. We select the first 20 classes in alphabetical order as known classes, while the remaining 10 classes are designated as unknown classes. We employ the ResNet-18 [He et al., 2016] as the backbone network, the remaining experimental settings are the same as those conducted on the CIFAR dataset. Notably, as shown in Table 3, POT achieves the SOTA performance on the ImageNet-30 dataset, both in terms of inlier accuracy and AUROC.

Comparison with IOMatch. IOMatch treats all OOD samples as a distinct category, framing the OSSL task as an open-set classification challenge. To enable a direct comparison with IOMatch, we configure our method to output the model-predicted category for samples with an OOD score less than 0.5, and samples with a score exceeding 0.5 are directly classified as OOD. This alignment allows for a meaningful comparison using identical evaluation indicators.

Note that IOMatch addresses the OSSL problem with the constraints of limited labels, particularly in few-shot learning scenarios. In contrast, our method is designed for general semi-supervised learning scenarios. In the comparative study, we experimented with a setting featuring 25 labeled samples per class, in contrast to IOMatch's few-shot learning scenario with only 4 labeled samples per class. This intentional deviation in experimental design ensures a more comprehensive and representative assessment of our method's performance in semi-supervised learning contexts.

In Tables 4 and 5, it shows that our approach clearly demonstrates a substantial advantage in open-set classification. In closed-set classification, specifically with 80 known classes, our method outperforms IOMatch by over 9% in closed-set classification tasks. These findings underscore the significant strengths of our approach in effectively handling a wide spectrum of classification categories.

Effect of k and ε. Table 6 shows the performance under 9 different configurations of POT parameters. Notably, when the unlabeled dataset in the CIFAR10 dataset contains 40 percent of OOD samples, the model achieves excellent results with k ranging from 1.5 to 2.5. This observation underscores the validity of our method, particularly the strategy of generating redundant mass for unlabeled distributions.

Further, we examine the effect of the entropy regularization parameter on the model. We find that the model exhibits increased stability with larger values of ε. In addition, based on the results in Table 6, the model consistently outperforms most parameter combinations, demonstrating the robust generalization capabilities inherent in our proposed method.

Effect of λOOD. The impact of varying the weight assigned to the OOD loss is demonstrated in Table 7. The result shows that our approach yields superior performance when tuning the OOD loss weights to align with the order of magnitude of the supervised loss (e.g., weights of 0.1 or 0.01). This suggests that the OOD loss should have a balanced weight with the supervised loss, neither dominating it nor being dominated by it. In practice, since the supervised loss typically carries a smaller magnitude, it suggests to use small weights for the OOD loss when implementing our method.

Training Efficiency. When computing the MSF, POT demonstrates a per-iteration complexity of O(mn), which is notably lower than the computational cost of conventional SSL methods. In our experiments on CIFAR-10 dataset, which includes 50 labeled samples per class and is executed on a single GPU, our method completes one training epoch in a total of 123 seconds, in which the computation of MSF takes only 6 seconds. It shows that our method achieves competitive performance in terms of efficiency.

Adaption to other SSL Methods. The OOD detector we introduced possesses the ability to seamlessly integrate with various off-the-shelf SSL methods, adapting them to OSSL scenarios. Table 8 illustrates the effectiveness of integrating our proposed OOD detector with FlexMatch [Zhang et al., 2021], an advanced algorithm derived from FixMatch. These findings underscore the generality of our method to SSLs.

5 Conclusion

In this work, we introduce a novel approach for OOD detection within the context of OSSL. In particular, we formulate the OOD detection in OSSL as a POT problem. With POT theory, we propose a novel mass score function based on transport mass to measure the likelihood of a sample being an outlier. More, to train the model in an end-to-end manner, we devise an OOD loss, enabling the adaption of off-the-shelf SSL methods to OSSL settings via multi-task learning. Extensive experiments on multiple datasets and OSSL settings are conducted to evaluate our proposals.

Acknowledgements

This work was supported in part by the National Natural Science Foundation of China under Grant 61772492, Grant 62072428, Grant 62271465, in part by the Open Fund Project of Guangdong Academy of Medical Sciences, China under Grant YKY-KF202206, and in part by the Suzhou Basic Research Program under Grant SYG202338. Xike Xie is the corresponding author.