Unifying Distribution Alignment as a Loss for Imbalanced Semi-supervised Learning

1 INTRODUCTION . Semi-supervised learning ( SSL ) uses a large pool of unlabeled data to learn a classifier despite having access to only a small amount of labeled data . Recently , techniques have been introduced ( Berthelot et al. , 2019 ; 2020 ; Sohn et al. , 2020 ) which simplify the process while at the same time pushing performance to new levels . However , these approaches have focused on the cases where the class distributions are balanced for both the labeled and unlabeled data . At the same time , work within the supervised learning community has shown renewed focus on imbalanced or sometimes long-tailed learning — owing to the fact that most data in the real world is not well-balanced . A variety of methodologies for this setting have been introduced ( Kang et al. , 2020 ; Menon et al. , 2021 ; Ren et al. , 2020 ; Hong et al. , 2021 ) . For instance , many have observed the bias that ordinary supervised learning techniques suffer from — favoring head classes over the less numerous tail classes . Kang et al . ( 2020 ) show that softmax-based classifiers often produce classification weights which correlate with class frequency and thus reduces the class-balanced performance of the model . Resampling strategies ( Chawla et al. , 2002 ; He & Garcia , 2009 ; Buda et al. , 2018 ; Byrd & Lipton , 2019 ) ( which sample from the pool based on the desired distribution ) are often effective as well . By shifting from sampling the data distribution for the majority of training to a more class-balanced regime at the end , one can learn a good representation while mitigating the aforementioned bias at the final classification layer . In addition , a distributional shift ( Hong et al. , 2021 ) can be observed within existing protocols where training occurs on an imbalanced dataset yet evaluation is done with respect to a balanced one . In this work , we study the combined setting of semi-supervised and imbalanced learning . In particular , we consider both FixMatch ( Sohn et al. , 2020 ) and MixMatch ( Berthelot et al. , 2019 ) as base semi-supervised learners . Both employ two losses : a cross entropy loss on the labeled data , and an unsupervised loss that relies on consistency between the classifier outputs among augmented versions of unlabeled examples . 1000 2000 3000 4000 5000 6000 Total Training Epochs 45 50 55 60 65 70 75 A cc ur ac y FixMatch DA , _min = 0.5 CReST+ UDAL ( ours ) Supervised ( 10 % labels ) Supervised ( 100 % labels ) Therefore , vulnerability to bias from the imbalance can happen in three ways : the supervised loss itself , the quality of pseudo-labels derived from the classifier on unlabeled examples , and the pseudolabels themselves can bias the classifier even if they are perfectly predicted . Furthermore , confirmation bias within semi-supervised learning is already a worrisome factor even without any imbalance between the classes ( Arazo et al. , 2020 ) . Within the balanced setting , the ReMixMatch ( Berthelot et al. , 2020 ) approach to semi-supervised learning introduces strong regularization through distribution alignment ( Bridle et al. , 1992 ) ( i.e . modifying the prediction by the ratio of the desired distribution to model distribution ) to help mitigate this . Recently CReST ( Wei et al. , 2021 ) has shown that distribution alignment also confers benefits in the imbalanced setting and aims to progressively rebalance the distribution of pseudo-labels . Similar to resampling approaches ( Kang et al. , 2020 ) , imbalanced learning considers a shift from random sampling to class-balanced sampling , CReST attempts to align the pseudo-labels themselves to a more balanced distribution as training progresses . However , distribution alignment is not the only technique that CReST relies upon to achieve good performance . In addition , CReST requires a generational approach to self-training which accumulates a relatively balanced subset of confident pseudo-labels to augment the labeled set with as a form of self-training . As the generations proceed , this subset is re-sampled to become more and more balanced . Each generation re-initializes the classifier ’ s network , and therefore , the only “ state ” retained is through these accumulated pseudo-labels ( now treated as ordinary supervised labels ) . This process can be extremely costly with respect to training time , and we hypothesize that it is not optimal because it fails to directly address the imbalance in the labeled data . Instead , we seek a simpler solution to imbalanced semi-supervised learning through distribution alignment alone . We ask the question : is this disjoint methodology truly necessary ? Can a single , central approach be devised to address imbalance in semi-supervised learning ? In this work , as our contributions , we identify an affirmative answer to this question by connecting the ideas of progressive distribution alignment from Berthelot et al . ( 2020 ) ; Wei et al . ( 2021 ) and the method of logit adjustment from fully supervised , imbalanced learning ( Menon et al. , 2021 ; Ren et al. , 2020 ; Hong et al. , 2021 ) . Furthermore , this approach can be implemented with only a few lines of code , has significantly reduced training time requirements , and generally outperforms previous work . Finally , it shows significantly better performance characteristics as more labeled data becomes available and readily scales to larger datasets — achieving a 1.6 % increase in accuracy compared to the best existing method on ImageNet-127 . 2 PREREQUISITES . We present background information on both problem settings . First , we formally define the problem setting of Imbalanced Semi-Supervised Learning ( SSL ) . Second , we outline the idea of distribution alignment ( Berthelot et al. , 2020 ; Wei et al. , 2021 ) to improve pseudo-label quality within both the balanced and imbalanced settings of SSL . We revisit a recent method , CReST ( Wei et al. , 2021 ) , which also attempts to address imbalanced semi-supervised learning . 2.1 CLASS-IMBALANCED SEMI-SUPERVISED LEARNING . Semi-supervised learning relies on two sources of data : a labeled set X = { ( xi , yi ) } Ni=1 where each xi is a training example and yi is the corresponding target . Since classification is the focus of this work , we consider yi as a class label within C = { 1 , . . . , C } with a total number of C classes . In imbalanced learning , we expect varying numbers of training examples across classes . Therefore , we denote the number of examples in our labeled set corresponding to class c ∈ C as Nc such that ∑C c=1Nc = N . We assume that the classes are ordered with respect to frequencies and in a descending manner i.e . Nc ≥ Nc+1 . It is often useful to characterize the degree of imbalance by the ratio N1/NC , and we refer to this as the imbalance ratio of the dataset . We use pdata ( y ) and q ( y ) to denote the marginal distributions of the data and the model . When there is no ambiguity , we drop y as pdata and q to simplify presentations . Additionally , we have an unlabeled set of examples U = { ui } Mi=1 for which we have no corresponding target . While we expect that this set is also imbalanced , we additionally make the common assumption ( Wei et al. , 2021 ) that it follows the same class distribution and thus shares the same imbalance ( ratio ) as the labeled set . Finally , an important measure β = NN+M considers the percentage of overall examples that are labeled . 2.2 DISTRIBUTION ALIGNMENT . Distribution alignment ( DA ) was re-introduced in the setting of semi-supervised learning within ReMixMatch ( Berthelot et al. , 2020 ) . To mitigate the tendencies of semi-supervised learning to suffer from confirmation bias , regularization can be added to the pseudo-label inference step . In particular , if we assume the labeled and unlabeled data both come from the same distribution pdata ( although , we do not know particular labels from the unlabeled set ) , we would expect that our model should produce pseudo-labels that follow the same distribution . This marginal distribution of the model q ( y ) can be estimated by moving average , which we denote as q̂ ( y ) or q̂ , as in Berthelot et al . ( 2020 ) . If we denote our current model ’ s predictions on unlabeled examples as q ( y|xu ) , these predictions can be re-scaled through dividing by q̂ and multiplying by pdata . After normalization , we have : q̃ ( y|x ) = Normalize ( q ( y|x ) pdata q̂ ) . ( 1 ) In equation 1 , we assume element-wise operations beween q ( y|x ) , pdata and q̂ . Normalize ( p ) ensures p as a probability distribution which sums to 1 . As noted in Wei et al . ( 2021 ) , it is not always optimal to align the predictions directly to pdata when pdata is imbalanced . Rather , a smoothed form which ( elementwise ) exponentiates the distribution by a factor of α before normalization : p̃α = Normalize ( pαdata ) , 0 ≤ α ≤ 1 . ( 2 ) is used instead of pdata and found to both regularize the predictions as well as combat bias . As α→ 0 , this approaches an alignment against a more uniform distribution . 2.3 LOGIT ADJUSTMENT . While DA is clearly applicable to pseudo-label inference during training , it has no direct effect on the labeled portion . Since a semi-supervised approach relies on both labeled and unlabeled losses , it is critical to address the problem of imbalance at the supervised level as well . For this , we examine a popular technique within supervised learning , often known as logit adjustment ( Menon et al. , 2021 ) , balanced softmax ( Ren et al. , 2020 ) , or LADE ( Hong et al. , 2021 ) . These methods modify the loss computation to compensate for the class imbalance found in the data distribution . Notably , when a data distribution is class-imbalanced , we attempt to minimize the classification loss with respect to this data distribution . However , at evaluation time , we either evaluate on a class-balanced dataset or produce a class-balanced error by averaging the per-class accuracies . This is a shift in distribution which can cause poor performance . Therefore , this shift is integrated into the cross entropy loss : LLA ( y , f ( x ) ) = LCE ( y , f ( x ) + log pdata − log ( Unif ( C ) ) ) ≡ LCE ( y , f ( x ) + log pdata ) ( 3 ) where Unif ( C ) is the discrete uniform distribution over C classes , y is the true label of x , f ( x ) is the vector-valued output of the classifier , and pdata is the marginal class distribution of the data as a vector . As elaborated within Menon et al . ( 2021 ) , this has the effect that instead of optimizing f ( x ) directly , we optimize h ( x ) = f ( x ) + log pdata which aligns the source distribution of f ( x ) correctly to the uniform class distribution seen during evaluation . Menon et al . ( 2021 ) also discusses an inference time procedure which attempts to account for this shift without modifications to the training procedure . Since this is “ for free ” , we include results combined with it in Table 1 as “ LA ( Inf ) ” for the inference time procedure .

Summary: This paper addresses the topic of semi-supervised learning in cases where the underlying data distribution is severely imbalanced. This approach combines distribution alignment with logit adjustment, resulting in an efficient method for solving the aforementioned problem while improving performance in the test setting. Unlike existing state-of-the-art approaches such as CReST, the approach involves no sampling state, and imbalance mitigation is achieved just by modifying the loss functions of the model. Experiments are conducted over three benchmark vision datasets of varying complexity: long-tailed versions of CIFAR10, CIFAR100, and ImageNet-127. Experimental results are competitive with or exceed other methods on the tested datasets with 5x training speedup compared to CReST.