ABC: Auxiliary Balanced Classifier for Class-imbalanced Semi-supervised Learning

Suppose that we have a labeled dataset $\mathcal{X}=\left\{\left(x_{n},y_{n}\right):n\in\left(1,...,N\right)\right\}$, where $x_{n}\in\mathbb{R}^{d}$ is the $n$th labeled data point and $y_{n}\in\left\{1,\ldots,L\right\}$ is the corresponding label. We also have an unlabeled dataset $\mathcal{U}=\left\{\left(u_{m}\right):m\in\left(1,\ldots,M\right)\right\}$, where $u_{m}\in\mathbb{R}^{d}$ is the $m$th unlabeled data point. We express the ratio of the amount of labeled data as $\beta=\frac{N}{M+N}$. Generally, $\beta<0.5$, because label acquisition is costly and laborious. We denote the number of labeled data points of class $l$ as $N_{l}$, i.e., $\sum_{l=1}^{L}N_{l}=N$, and assume that the $L$ classes are sorted according to cardinality in descending order, i.e., $N_{1}\geq N_{2}\geq\cdots\geq N_{L}$. We denote the ratio of the class imbalance as $\gamma=\frac{N_{1}}{N_{L}}$. Under class-imbalanced scenarios, $\gamma\gg 1$. Following previous CIL studies, we define the half of the classes containing a large amount of data as the majority classes, and the other half of the classes, containing a small amount of data, as the minority classes. Following [34], we assume that $\mathcal{X}$ and $\mathcal{U}$ share the same class distribution, i.e., the labeled and unlabeled datasets are class-imbalanced to the same extent. From $\mathcal{X}$ and $\mathcal{U}$, we generate minibatches $\mathcal{MB}_{\mathcal{X}}=\left\{\left(x_{b},y_{b}\right):b\in\left(1,...,B\right)\right\}\subset\mathcal{X}$ and $\mathcal{MB}_{\mathcal{U}}=\left\{\left(u_{b}\right):b\in\left(1,\ldots,B\right)\right\}\subset\mathcal{U}$ for each iteration of training, where $B$ is the minibatch size. Using these minibatches for training, we aim to learn a model $f:\mathbb{R}^{d}\rightarrow\left\{1,...L\right\}$ that performs effectively on a class-balanced test set.

We attach the ABC to the backbone’s representation layer, so that it can utilize the high-quality representations learned by the backbone. We use FixMatch [29] or ReMixMatch [3] as the backbone, as these two have achieved state-of-the-art SSL performance. FixMatch uses the classification loss calculated from the weakly augmented labeled data point $\alpha\left(x_{b}\right)$ generated by flipping and cropping the image, and the consistency regularization loss calculated from the weakly augmented unlabeled data point $\alpha\left(u_{b}\right)$ and strongly augmented unlabeled data point $\mathcal{A}\left(u_{b}\right)$ generated by Cutout [10] and RandomAugment [8]. ReMixMatch predicts the class label of the weakly augmented unlabeled data point $\alpha\left(u_{b}\right)$ using distribution alignment and sharpening, and assigns the predicted label to the strongly augmented unlabeled data point $\mathcal{A}\left(u_{b}\right)$. These strongly augmented unlabeled data point $\mathcal{A}\left(u_{b}\right)$ and strongly augmented labeled data point $\mathcal{A}\left(x_{b}\right)$ are used to conduct mixup regularization. ReMixMatch also conducts consistency regularization in a manner similar to FixMatch and self-supervised learning using the rotation of the image [11, 39]. FixMatch and ReMixMatch have greatly improved the SSL performance by learning high-quality representations using strong data augmentation. However, these algorithms can be significantly biased toward the majority classes in class-imbalanced settings.

Using FixMatch and ReMixMatch as the backbone of the proposed algorithm, we ensure that the ABC enjoys high-quality representations learned by the backbone, while replacing the backbone’s biased classifier. To train the ABC, we reuse the weakly augmented data and strongly augmented data used by the backbone to decrease the computational cost. Although we use FixMatch and ReMixMatch as the backbone in this study, the ABC can also be combined with other DNN-based SSL algorithms, as long as they use weakly augmented data and strongly augmented data.

To train the ABC to be balanced, we first generate $0/1$ mask $M\left(x_{b}\right)$ for each labeled data point $x_{b}$ using a Bernoulli distribution $\mathcal{B}\left(\cdot\right)$ with the parameter set to be inversely proportional to the number of data points of each class. This setting makes $\mathcal{B}\left(\cdot\right)$ generate mask 1 with high probability for the data points in the minority classes, but with low probability for those in the majority classes. Then, the classification loss is multiplied by the generated mask, so that the ABC can be trained with a balanced classification loss. Multiplying the classification loss by the $0/1$ mask can be interpreted as oversampling of the data points in the minority classes, whereas it can be interpreted as undersampling of those in the majority classes. In representation learning, oversampling and undersampling techniques have shown overfitting and information loss problems, respectively. In contrast, the ABC can overcome these problems because it uses the representations learned by the backbone, which is trained on all data points in the minibatch. The use of the $0/1$ mask to construct the balanced loss, instead of directly creating a balanced subset, allows the backbone and the ABC to be trained from the same minibatches. Therefore, the representations of minibatches calculated for training the backbone can be used again for training the ABC. Consequently, the proposed algorithm only requires a slightly increased time cost compared to training the backbone alone. This is confirmed in Section 4.3. The overall procedure of balanced training with $0/1$ mask for the ABC attached to a representation layer of the backbone is presented in Figure 2. The classification loss for the ABC, $L_{cls}$, with $0/1$ mask $M\left(\cdot\right)$ is expressed as

where $\mathbf{H}$ is the standard cross-entropy loss, $\alpha\left(x_{b}\right)$ is an augmented labeled data point, $p_{s}\left(y|\alpha\left(x_{b}\right)\right)$ is the predicted class distribution using the ABC for $\alpha\left(x_{b}\right)$, and $p_{b}$ is the one-hot label for $x_{b}$.

To increase the margin between the decision boundary and the data points using unlabeled data, we conduct consistency regularization for the ABC, similar to the way in FixMatch. Specifically, we first obtain the predicted class distribution $p_{s}\left(y|\alpha\left(u_{b}\right)\right)$ for a weakly augmented unlabeled data point $\alpha\left(u_{b}\right)$ using the ABC and use it as a soft pseudo-label $q_{b}$. Then, for two strongly augmented unlabeled data points $\mathcal{A}_{1}\left(u_{b}\right)$ and $\mathcal{A}_{2}\left(u_{b}\right)$, we train the ABC to produce their predicted class distributions, $p_{s}\left(y|\mathcal{A}_{1}\left(u_{b}\right)\right)$ and $p_{s}\left(y|\mathcal{A}_{2}\left(u_{b}\right)\right)$, to be close to $q_{b}$.

In class-imbalanced settings, because most unlabeled data points belong to majority classes, most weakly augmented unlabeled data points can be predicted as the majority classes. Then, consistency regularization would be conducted with a higher frequency for the majority classes, which can cause a classifier to be biased toward the majority classes. To prevent this issue, we conduct consistency regularization in a modified manner that is suitable for class-imbalance problems. Specifically, whereas FixMatch minimizes entropy by converting the predicted class distribution for a weakly augmented data point into a one-hot pseudo-label, we directly use the predicted class distribution as a soft pseudo-label. We do not pursue entropy minimization for the ABC because it can accelerate biased classification toward certain classes. Moreover, we once again generate $0/1$ mask $M\left(\cdot\right)$ for each unlabeled data point $u_{b}$ based on a soft pseudo label $q_{b}$, and multiply the consistency regularization loss for $u_{b}$ by the generated mask, so that the ABC can be trained with a class-balanced consistency regularization loss. Note that existing resampling techniques are not applicable to unlabeled data, because they require a label for each data point. In contrast, we make it possible to resample unlabeled data by using the soft pseudo-label and the $0/1$ mask. The consistency regularization loss, $L_{con}$, with $0/1$ mask $M\left(\cdot\right)$ is expressed as

where $\mathbf{I}$ is the indicator function, $\max\left(q_{b}\right)$ is the highest predicted assignment probability for any class, representing the confidence of prediction, and $\tau$ is the confidence threshold. To avoid the unwanted effects of inaccurate soft pseudo-labels during consistency regularization, we only use the weakly augmented unlabeled data points whose confidence is higher than the threshold $\tau$, similar to that in FixMatch. To take full advantage of few unlabeled data points with prediction confidence values that are higher than the confidence threshold $\tau$ in the early stage of training, we gradually decrease the parameter of the Bernoulli distribution $\mathcal{B}\left(\cdot\right)$ for $u_{b}$ from 1 to ${N_{L}}/{N_{\widehat{q}_{b}}}$, where $\widehat{q}_{b}$ is the one-hot pseudo-label obtained from $q_{b}$. Following previous studies [3, 24, 29, 4], we do not backpropagate gradients for pseudo-label prediction. The overall procedure for consistency regularization for the ABC is shown in Appendix A.

Unlike a recent CIL trend to finetune a classifier in a balanced manner after representation learning is completed (i.e., decoupled learning of representations and a classifier) [17, 27], we obtain a balanced classifier by training the proposed algorithm end-to-end. We train the proposed algorithm with the sum of losses from Sections 3.3 and 3.4, and the loss for the backbone, $L_{back}$. The total loss function $L_{total}$ is expressed as

Whereas we use the sum of the losses for the backbone and ABC for training the proposed algorithm, we predict the class labels of new data points using only the ABC. In our experiments in Sections 4.4 and 4.5, we show that the proposed algorithm trained end-to-end produces better performance than competing algorithms with decoupled learning of representations and a classifier, and we analyze possible reasons. We present the pseudo code of the proposed algorithm in Appendix B.

We created class-imbalanced versions of CIFAR-10, CIFAR-100 [21], and SVHN [25] datasets to conduct experiments under various ratios of class imbalance $\gamma$ and various ratios of the amount of labeled data $\beta$. For class-imbalance types, we first consider long-tailed (LT) imbalance in which the number of data points exponentially decreases from the largest to the smallest class, i.e., $N_{k}=N_{1}\times\gamma^{-\frac{k-1}{L-1}}$, where $\gamma=\frac{N_{1}}{N_{L}}$. We also consider step imbalance [5] in which the whole majority classes have the same amount of data and the whole minority classes also have the same amount of data. Two types of class imbalance for the considered datasets are illustrated in Appendix C. For the main setting, we set $\gamma=100$, $N_{1}=1000$, and $\beta=20\%$ for CIFAR-$10$ and SVHN, and $\gamma=20$, $N_{1}=200$ and $\beta=40\%$ for CIFAR-$100$. Similar to [18], we set $\gamma$ of CIFAR-100 to be relatively small because CIFAR-$100$ has only 500 training data points for each class. To evaluate the performance of the proposed algorithm on large-scale datasets, we also conducted experiments on 7.5M data points of 256 by 256 images from the LSUN dataset [37].

We compared the performance of the proposed algorithm with that of various baseline algorithms. Specifically, we considered the following baseline algorithms:

• •

  Deep CNN (vanilla algorithm): This is trained on only labeled data with the cross-entropy loss.

• •

  BALMS [27] (CIL algorithm): This state-of-the-art CIL algorithm does not use unlabeled data.

• •

  VAT [24], ReMixMatch [3], and FixMatch [29] (SSL algorithms): These are state-of-the-art SSL algorithms, but do not consider class imbalance.

• •

  FixMatch+CReST+PDA and ReMixMatch+CReST+PDA (CISSL algorithms): CReST+PDA [34] mitigates class imbalance by using unlabeled data points classified as the minority classes with a higher probability than those classified as the majority classes.

• •

  ReMixMatch+DARP and FixMatch+DARP (CISSL algorithms): These algorithms use DARP [18] to refine the pseudo labels obtained from ReMixMatch or FixMatch.

• •

  ReMixMatch+DARP+cRT and FixMatch+DARP+cRT (CISSL algorithms): Compared to ReMixMatch+DARP and FixMatch+DARP, these algorithms finetune the classifier using cRT [17].

For the structure of the deep CNN used in the proposed and baseline algorithms, we used Wide ResNet-$28$-$2$ [38]. We trained the proposed algorithm for $250,000$ iterations with a batch size of 64. The confidence threshold $\tau$ was set to $0.95$ based on experiments with various values of $\tau$ in Appendix D. We used the Adam optimizer [20] with a learning rate of $0.002$, and used Cutout [10] and RandomAugment [8] for strong data augmentation, following [18]. Similar to [3, 4], we evaluated the performance of the proposed algorithm using an exponential moving average of the parameters over iterations with a decay rate of $0.999$, instead of scheduling the learning rate. In Tables 1-5, we used the overall accuracy and the accuracy only for minority classes as performance measures. We repeated the experiments five times under the main setting, and three times under the step imbalance and other settings of $\beta$ and $\gamma$. We report the average and standard deviation of the performance measures over repeated experiments. For the vanilla algorithm, FixMatch+DARP+cRT, and ReMixMatch+DARP+cRT, which suffered from overfitting, we measured performance every 500 iterations and recorded the best performance. Further details of the experimental setup are described in Appendix E.

The performance of the competing algorithms under the main setting are summarized in Table 1. We can observe that the proposed algorithm achieved the highest overall performance, with improved performance for minority classes. Interestingly, VAT, an SSL algorithm, showed similar performance to the vanilla algorithm, and worse performance than BALMS, a CIL algorithm. Similarly, FixMatch and ReMixMatch, which do not consider class imbalance, showed poor performance for minority classes. Although BALMS mitigated class imbalance, it produced poor overall performance, as it did not use unlabeled data for training. This demonstrates the importance of using unlabeled data for training, even in the class-imbalanced setting. FixMatch+CReST+PDA and ReMixMatch+CReST+PDA mitigated class imbalance by using unlabeled data points classified as the minority classes with a higher probability, but produced lower performance than the proposed algorithm. This may be because even if all unlabeled data points classified as minority classes are additionally used for training, their amount is still less than that of the data in majority classes, while the proposed algorithm uses class-balanced minibatches by generating the $0/1$ mask. Fixmatch+DARP and ReMixMatch+DARP slightly mitigated class imbalance by refining biased pseudo-labels, but resulted in lower performance than the proposed algorithm. This may be because even perfect pseudo labels cannot change the underlying class-imbalanced distribution of the training data. By additionally using a rebalancing technique cRT, FixMatch(ReMixMatch)+DARP+cRT performed better than FixMatch(ReMixMatch)+DARP. However, FixMatch(ReMixMatch)+DARP+cRT still performed worse than FixMatch(ReMixMatch)+ABC, although it also uses high-quality representations learned by FixMatch(ReMixMatch) and techniques for mitigating class imbalance. The superior performance of FixMatch(ReMixMatch)+ABC over FixMatch(ReMixMatch)+DARP+cRT is probably because FixMatch(ReMixMatch)+ABC was trained end-to-end, and the ABC was also trained using unlabeled data. We discuss this in more detail in Sections 4.4 and 4.5. Overall, the algorithms combined with ReMixMatch performed better than the algorithms combined with FixMatch. In addition to the overall accuracy and minority-class-accuracy, we also compared the performance of the competing algorithms in terms of the geometric mean (G-mean) of class-wise accuracy under the main setting in Appendix F.

To evaluate the performance of the proposed algorithm in various settings, we conducted experiments using ReMixMatch, FixMatch, and the CISSL algorithms considered in Table 1, while changing the ratio of class imbalance $\gamma$ and the ratio of the amount of labeled data $\beta$. The results for CIFAR-$10$ are presented in Table 2, and the results for SVHN and CIFAR-$100$ are presented in Appendix G. In Table 2, we can observe that the proposed algorithm achieved the highest overall accuracy with greatly improved performance for minority classes for all settings. Because FixMatch+DARP+cRT and ReMixMatch+DARP+cRT do not use unlabeled data for classifier tuning, the difference in performance between FixMatch(ReMixMatch)+DARP+cRT and the proposed algorithm increased as the ratio of the amount of labeled data $\beta$ decreased and as the ratio of class imbalance $\gamma$ increased. In addition, the difference in performance between FixMatch(ReMixMatch)+CReST+PDA and the proposed algorithm tended to increase as the ratio of class imbalance $\gamma$ increased, because the difference between the number of labeled data points belonging to the majority classes and the number of unlabeled data points classified as the minority classes increases with $\gamma$.

We also conducted experiments under a step-imbalance setting, where the class imbalance was more noticeable. This setting assumes a more severely imbalanced class distribution than the LT imbalance settings, because half of the classes have very scarce data. The experimental results for CIFAR-$10$ are presented in Table 3, and the results for SVHN and CIFAR-$100$ are presented in Appendix H. In Table 3, we can see that the proposed algorithm achieved the best performance, and the performance margin is greater than that of the LT imbalance settings. ReMixMatch+CReST+PDA showed relatively low performance compared to the other algorithms.

To evaluate the performance of the proposed algorithm on a large-scale dataset, we also conducted experiments on the LSUN dataset [37], which is naturally a long-tailed dataset. Among the algorithms considered in Tables 2 and 3, those combined with CReST were excluded for comparison, because CReST requires loading of the whole unlabeled data in the repeated process of updating pseudo-labels, which is not possible for the large-scale LSUN dataset. Instead, we additionally considered FixMatch+cRT and ReMixMatch+cRT for comparison. The experimental results are presented in Table 4. The proposed algorithm showed better performance than the other baseline algorithms. DARP resulted in degradation of the performance, possibly because the scale of the LSUN dataset is very large. Specifically, DARP solves a convex optimization with all unlabeled data points to refine the pseudo labels. As the scale of the unlabeled dataset increases, this optimization problem becomes more difficult to solve and, consequently, the pseudo-labels could be refined inaccurately. Unlike the results for other datasets, the algorithms combined with FixMatch performed better than the algorithms combined with ReMixMatch.

The proposed algorithm requires additional parameters for the ABC, but the number of the additional parameters is negligible compared to the number of parameters of the backbone. For example, the ABC additionally required only $0.09\%$ and $0.87\%$ of the number of backbone parameters for CIFAR-$10$ with $10$ classes and CIFAR-$100$ with $100$ classes, respectively. Moreover, because the ABC shares the representation layer of the backbone, it does not significantly increase the memory usage and training time. Furthermore, we could train the proposed algorithm on the large-scale LSUN dataset without a significant increase in computation cost, because the entire training procedure could be carried out using minibatches of data. In contrast, the algorithms combined with DARP required convex optimization for all pseudo-labels, which significantly increased the computation cost as the number of classes or the amount of data increased. Similarly, it required significant time to train the algorithms combined with CReST, because CReST requires iterative re-training with a labeled set expanded by adding unlabeled data points with pseudo-labels. We present the floating point operations per second (FLOPS) for each algorithm using Nvidia Tesla-V100 in Appendix I.

The ABC can use high-quality representations learned by the backbone when performing balanced classification. To verify this, in Figure 3, we present t-distributed stochastic neighbor embedding (t-SNE) [31] of the representations of the CIFAR-$10$ test set learned by the ABC (without SSL backbone), FixMatch+ABC, and ReMixMatch+ABC on CIFAR-$10$-LT under the main setting. Different colors indicate different classes. As expected, “ABC (without SSL backbone)" failed to learn class-separable representations because sufficient data were not used for training while using the $0/1$ mask. In contrast, by training the backbone (FixMatch or ReMixMatch) together with the ABC, the proposed algorithm could use the entire data and learn high-quality representations. In this example, ReMixMatch produced more separable representations than FixMatch, which shows that the choice of the backbone affects the performance of the proposed algorithm, as expected.

The proposed algorithm can also mitigate class imbalance by using the ABC. To verify this, we compare the confusion matrices of the predictions on the test set of CIFAR-$10$ using ReMixMatch, ReMixMatch+DARP+cRT, and ReMiMatch+ABC trained on CIFAR-$10$ under the main setting in Figure 4. In the confusion matrices, the value in the $i$th row and the $j$th column represents the ratio of the amount of data belonging to the $i$th class to the amount of data predicted as the $j$th class. Each cell has a darker red color when the ratio is larger. We can see that ReMixMatch often misclassified data points in the minority classes (e.g., classes $8$ and $9$ into classes $0$ and $1$). This may be because ReMixMatch does not consider class imbalance, and thus biased pseudo-labels were used for training. ReMixMatch+DARP+cRT produced a more balanced class-distribution compared to ReMixMatch by additionally using DARP+cRT. However, a significant number of data points in the minority classes were still misclassified as majority classes. In contrast, ReMixMatch+ABC classified the test data points in the minority classes with higher accuracy, and produced a significantly more balanced class distribution than ReMixMatch+DARP+cRT, as shown in Figure 4 (c). As both ReMixMatch+DARP+cRT and ReMixMatch+ABC use ReMixMatch to learn representations, the performance gap between these two algorithms results from the different characteristics of the ABC versus DARP+cRT as follows. First, DARP+cRT does not use unlabeled data for training its classifier after representations learning is completed, whereas the ABC uses unlabeled data with unbiased pseudo-labels for its training. Second, whereas DARP+cRT decouples the learning of representations and training of a classifier, the ABC is trained end-to-end interactively with representations learned by the backbone. We also present the confusion matrices of the predictions on the test set of CIFAR-$10$ using FixMatch, FixMatch+DARP+cRT, and FixMatch+ABC as well as the confusion matrices of the pseudo-labels on the same dataset using ReMixMatch, ReMixMatch+DARP+cRT, ReMixMatch+ABC, FixMatch, FixMatch+DARP+cRT, and FixMatch+ABC in Appendix J. Moreover, we compare the ABC and the classifier of DARP+cRT in more detail using the validation loss plots in Appendix K.

We conducted an ablation study on CIFAR-$10$-LT in the main setting to investigate the effect of each element of the proposed algorithm. The results for ReMixMatch+ABC are presented in Table 5, where each row indicates the proposed algorithm with the described conditions in that row. The results are summarized as follows. 1) If we did not gradually decrease the parameter of the Bernoulli distribution $\mathcal{B}\left(\cdot\right)$ when conducting consistency regularization, then an overbalance problem occurred because of unlabeled data misclassified as minority classes. 2) Without consistency regularization for the ABC, the decision boundary did not clearly separate each class. 3) Without using the $0/1$ mask for $L_{cls}$ and $L_{con}$, the ABC was trained to be biased toward the majority classes. 4) Without confidence threshold $\tau$ for consistency regularization, training became unstable and, consequently, the ABC was trained to be biased toward certain classes. 5) Similarly, if hard pseudo-labels, instead of soft pseudo-labels, were used for consistency regularization, then the ABC was biased toward certain classes. 6) If the ABC was solely used without the backbone, the performance decreased because the ABC could not use high-quality representations learned by the backbone. 7) When we used a re-weighting technique [13] instead of a mask for the ABC, training became unstable because of abnormally large gradients calculated for training on the data of the minority classes. 8) The decoupled training of the backbone and ABC resulted in decreased classification performance, as was also analyzed in Section 4.4. Similarly, we present the results of the ablation study for FixMatch+ABC in Appendix L.