eGAN: Unsupervised approach to class imbalance using transfer learning

A dataset is considered imbalanced when there is a significant, or in some cases, extreme disproportion between the number of samples of the different classes in the dataset. The class or classes with large number of samples are called the majority, while the class with few examples are denoted as the minority. In many cases, the machine learning model is required to correctly classify the minority class while minimizing the misclassification of the majority class. However, the skewness in the data often leads machine learning classification methods to favour the majority class.

Class imbalance problem in computer vision is normally approached either at the data level or algorithm level. Using data augmentation, a class with a small number of samples can be expanded into a class with a much larger number of samples. Earlier data augmentation was achieved simply by transforming images via scaling, cropping, flipping, padding, rotation, brightness, contrast, saturation level etc [15]. Now-a-days, synthetic images can also be generated using generative models such as VAE, GAN [8][4]. As a result, a humongous image dataset can be created from the images of the minority class.

At the algorithm level, the objective function is tweaked to heavily penalize the network for mis-classifying the minority class [13] [10]. The most popular is cost-sensitive approach. Here, the classifier is modified to incorporate varying penalty for each of considered groups of examples. By assigning a higher cost to the less represented set of samples its importance is boosted during training.

Transfer learning has been known to help improve the performance of machine learning models [14]. By fine-tuning varying number of layers in the pre-trained image classification model, the pre-trained model can serve as a feature extractor, while adding a classifier head for more specific feature learning for the current task.

In this work, we compared the performance of various pre-trained image classification models for the task of unsupervised image classification with varying imbalance ratios. Our architecture, named eGAN, is developed to serve as a basis for this comparison. Using GAN [4], we reparameterise the task of the discriminator as a classifier which outputs a positive score for majority samples and a negative score for the minority ones. We integrate an encoder module to the GAN network that encodes the minority samples into a latent code from which the generator learns. While most GAN-based architectures focus on the output of the generator, in our proposed approach, as we intuitively adapted the vanilla GAN network and the corresponding loss function to directly classify the majority and the minority class, we are more concerned about the performance of the discriminator.

There have been a lot of work in the last few decades to address class imbalance. Earlier approaches include deliberate undersampling of the majority class or oversampling of the minority class by mere copying [3]. However, for image data the earlier approach leads to loss of useful data information while the latter approach causes overfitting [11]. Data augmentation via rotation, scaling, cropping etc can be considered as a variant of oversampling which copy the same data, however with little modification [15]. VAE and GAN enables the generation of completely new data [8][4]. In recent years, GAN-based approaches have gained much popularity than others and a good number of variants of vanilla GAN have been proposed to address class imbalance [12] [16] [1].

In [6], an ensemble method was proposed based on advanced generative adversarial network to generate new samples for the minority class to restore balance. Our opinion is that the computational demand of such approach is enormous, and many low-income countries of the world do not have access to such computation power. Deep Cascading (DC) with a long sequence of decision trees could help to handle unbalanced data [2]. A DC is a sequence of n classifiers where each sample x passes to the next classifier only if the current one classifies it as positive according to a high-sensitivity decision threshold. However, this works well with foreground-background imbalance unlike the classification task. Transfer learning with GAN was used to generate images from limited data in [14]. Their result showed that knowledge from pre-trained networks can ensure faster convergence and significantly improve the quality of generated images.

In this section, we discuss our proposed approach and the various testbeds that were used in our experiments.

All discriminator scores that are less than zero are classified as minority, otherwise they are classified as majority class. Table 1 shows the result obtained by combining various pre-trained models. Adam optimizer with a learning rate of 1e-4 was used to train all models for 100 epochs.

As can be seen in Figure 2, the model achieved a Nash equilibrium on test data at around 10 epochs. Here, we perform inference on test data at every epoch and plot the number of samples correctly classified. Five layers of DenseNet121 were fine-tuned at each generator and discriminator module, while 422 layers’ weight were kept fixed. At Nash, the discriminator correctly classified roughly 700/1000 of each of the minority and majority test data. This result convinces us that transfer learning with GAN can be used to overcome the challenge of highly imbalanced dataset, owing to the fact that we train only with 10 samples of the minority class and 5000 samples of the majority class. Similar performance is observed in CIFAR-100 with imbalance ratio 1:50.

Without pre-training the discriminator, the effect of the high imbalance in the training set is revealed, as the discriminator is skewed towards the majority class in the training set, thereby missing all the minority samples in the test data. This can be seen in Figure 3 on CIFAR-100. This behaviour pattern is observed on CIFAR-10 as well. We experiment with no pre-training at all, neither in the discriminator nor generator, and observed exact same pattern. Therefore, we can safely conclude that the use of transfer learning helps unsupervised image classification in a highly imbalanced domain.

Training can be stopped as soon as Nash equilibrium is reached, as this point gives the model best performance on the minority and majority class. An acceptable threshold can also be set for the absolute difference of the number of correctly classified samples of both classes. For instance, if the $\mid$ correctly_classified_minority - correctly_classified_majority $\mid$ $\leq$ 20. The precision, recall and F1-score curves on CIFAR-10 averaged over five folds at different imbalance ratios are shown in Figure 2.

We observed that at the early training epochs, typically between 1 and 40 epochs, the generator tries to achieve its objective of fooling the discriminator by generating samples from the majority class fed into the discriminator. That results in more of the minority samples being mis-classified as the discriminator ”knows” the distribution of the majority too well. A drastic change occurs when the generator start generating samples from latent vector, which can fool the discriminator as seen in generator and discriminator loss shown in Figure 4.

In this work, we demonstrates the capability of a GAN-based unsupervised technique to address class imbalance using pre-trained models. We conducts experiment with varying levels of imbalance ratios in the training dataset. Instead of synthesizing artificial images with the generator for data augmentation, we employ the discriminator as a classifier and formulate the loss function accordingly. Experimental results reveal that transfer learning plays a significant role in the model performance. The performance measure of interest plays a significant role in deciding the trained model from which epoch to deploy in production as the model at different epochs favour different evaluation metrics for example sensitivity. Future work will focus on the usage of this approach for anomaly detection task where the distinguishing features between normal (majority) and abnormal (minority) are less profound. Our work can be further explored in object detection tasks in case of imbalance between foreground and background.