Mistake-driven Image Classification with FastGAN and SpinalNet

1 INTRODUCTION . Supervised training of deep learning models , like the image classification models we consider in this paper , is the most efficient training approach , but also the most data-intensive . To mitigate this issue , alternative techniques are commonly applied , e.g. , transfer learning from a pretrained model to a new domain ( Ridnik et al. , 2021 ) or data augmentation to expand the available dataset synthetically ( Shorten & Khoshgoftaar , 2019 ) . One particular data augmentation technique is to use Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2020 ) that learn to produce new data that is similar to the distribution of the GAN ’ s training data . GAN-based augmentation has been shown to be successful in aiding the training process ( Bowles et al. , 2018 ; Wang et al. , 2018c ; Tanaka & Aranha , 2019 ) and has found practical adoption in other research domains ( Frid-Adar et al. , 2018 ; Sasmal et al. , 2020 ) . One challenge and limitation of GAN-based augmentation is either the high computational cost to augment each class of the dataset individually or the selection of the classes to augment that best support the model training . In this paper , we focus on a mistake-driven training procedure ( see Figure 1 ) that utilizes GAN-based data augmentation only after the initial model training , albeit transfer learning or training from scratch , and augments only those classes with the lowest accuracy , where the model shows performance weaknesses . This keeps the effort for class-wise GAN training at a lower level than a full-sized data augmentation procedure but still contributes to the model improvement by focusing on the most significant deficits accuracy-wise . The mistake-driven training method is fully compatible with any other training regime as it does not interfere directly with the model training or finetuning step except by querying the model during evaluation and enhancing the dataset for finetuning . Due to this explicit focus on model weaknesses as well as the careful selection of the components in the mistake-driven training pipeline , our method is computationally efficient and shows compet- itive performance on a number of benchmarks . First , we employ the progressive SpinalNet neural network architecture ( Kabir et al. , 2020 ; Chopra , 2021 ) as the image classification model , build on a backbone model , e.g . a pretrained Wide-ResNet ( Zagoruyko & Komodakis , 2016 ) or EfficientNetB7 ( Tan & Le , 2019 ) as used in the experiments . Afterward , the class-wise accuracy on the validation set is computed , and the worst-performing classes are identified . For each of these classes , a sampleefficient GAN is trained , here we choose the light-weight FastGAN method ( Liu et al. , 2021 ) , and use it to generate new samples . Finally , the already trained model is finetuned under consideration of the newly generated samples to improve and balance the model performance among the previously worst-performing classes . We perform an extensive experimental evaluation to test the influence of each component in the mistake-driven training pipeline and our component selection of backbone models , activation functions , and model optimizers . The method is evaluated on five widely used image classification datasets and exceeds the previous state-of-the-art accuracy on four of them . The contribution of our paper is threefold : 1 . During training , our method focuses on those classes with the lowest accuracy after the initial training phase . Only these classes are augmented to boost the accuracy , which leads to better performance . 2 . Due to the light-weight GAN method , the GAN-based augmentation is viable and effective for mistake-driven training even for datasets with only a few example images per class , while simultaneously requiring less computation than other , more complex GAN methods . 3 . Our extensive experiments , including ablation studies , show competitive or better accuracy than the previous state-of-the-art on five datasets . 2 RELATED WORK . This section relates the presented work to the recent literature on data augmentation using GANs . There are two types of image augmentation methods that help in increasing the accuracy of the model and making the model more robust . The first one is a geometrical transformation-based augmentation , and the other one is generative adversarial-based augmentation . Geometrical Transformation & Data Augmentation In the past decade , various data augmentation techniques have been used to improve classification prediction accuracy . The most common ones are those based on geometrical transformations ( Shorten & Khoshgoftaar , 2019 ) e.g. , cropping , flipping , rotation , color space , noise injection , and translation . These geometrical transformations increase the training dataset and improve the test accuracy ( Perez & Wang , 2017 ) . These data augmentation techniques are a best practice for training deep neural networks , especially in computer vision tasks where a vast number of augmentations are available and can be combined . Cubuk et al . ( 2019 ) further introduces a method to perform this combination , called AutoAugment automatically , that learns dataset-specific augmentation strategies to increase the final model accuracy . Synthetic Data Generation Using Generative Models Generative models , GANs in this context , have shown growing capability in generating very realistic data . These generated synthetic images can be utilized for data augmentation in data classification . Researchers have been attempting to enhance high-resolution images with GANs but with little success due to the enormous amount of data required to train GANs . Data augmentation with GANs has recently received increasing attention from the research community . Some research works have attempted to supplement images using GANs for data classification ( Tanaka & Aranha , 2019 ) . However , with limited data , utilizing GANs to augment data becomes a strenuous process . Rashid et al . ( 2019 ) , proposed to augment the ISIC skin lesion classification dataset using GANs . The results produced by data augmentation using GANs outperform ResNet and densenet models . In another line of work , Sasmal et al . ( 2020 ) used DCGAN for generating synthetic colonoscopic images . The generated images help in data augmentation to perform better polyp classification . Furthermore , Bejiga & Melgani ( 2018 ) suggests a GAN-based domain adaptation technique for aerial image classification . They use GANs for unsupervised domain adaptation of aerial remote sensing images . In the same field of work , Saha et al . ( 2021 ) developed TilGAN for the classification of images showing til , i.e. , tumor-infiltrating lymphocytes and non-til images . In addition to the previous research works , Zeng et al . ( 2020 ) propose to use GANs for augmenting data to detect disease severity . In another line of work , GANs were used for data augmentation in the field of multi-domain learning ( Yamaguchi et al. , 2020 ) . These multi-domain GANs learn both the outer and target datasets simultaneously and generate new samples for the target tasks . A detailed survey on the use of GANs for data augmentation is presented by Shorten & Khoshgoftaar ( 2019 ) . In this paper , our focus is to augment high-resolution images with a limited dataset which can further be used for finetuning the model . However , if the amount of the training data is not a constraining factor , IDA GAN ( Yang & Zhou , 2021 ) and Polarity GAN ( Deepshikha & Naman , 2020 ) can be utilized for data augmentation , too . Data Augmentation with Limited Data GANs generally require massive datasets for training ; limited data makes learning the underlying model arduous . Moreover , training GANs with limited data can easily lead to overfitting ( Bowles et al. , 2018 ) , which makes the training more difficult . Transfer learning can be used to avoid the problem of overfitting of GANs , which also results in better performance ( Bengio , 2012 ) . Apart from the traditional way of augmenting data using geometric transformations , some GANs use observations that have been forged ( Wang et al. , 2018b ) for generating new samples . Even though generating new images with limited data using GANs remains a challenging task , but at least the risk of overfitting can be reduced by using a pretrained model on a huge dataset and using that pretrained model for training GANs can produce some good quality images Wang et al . ( 2018b ) . However , using this technique sometimes leads to mode collapse , which is further fixed by Liu et al . ( 2021 ) using skip connections while training . 3 BACKGROUND . 3.1 SHARPNESS-AWARE MINIMIZATION . In deep learning models , we need to optimize the loss function such that the DL model can converge to global minima . To reach the minima , many optimization algorithms such as stochastic gradient descent , Adam ( Kingma & Ba , 2014 ) , and RMSProp are being used as a design choice for optimizing the loss function . Zhang et al . ( 2016 ) show that DL models could memorize the training data and easily overfit to it , due to which the trained model lack generalization ability . However , when a DL model converges to an abrupt global minimum , the value of the loss function remains high in the neighborhood of the global minimum , causing the DL model to lose its generalization ability . For the better generalization of DL models , Sharpness-Aware Minimization ( SAM ) was proposed ( Foret et al. , 2020 ) , where the optimization goal is reformulated such that it considers the output of loss functions from not only the minimum but also at its neighboring points . In that way , SAM minimizes both loss value and loss sharpness . Finally , the minimization of loss value and loss sharpness is done using gradient descent by learning those parameters that converge to a global minima , which has nearly the same loss value at the neighboring points . 3.2 PROGRESSIVE SPINAL NETWORKS . In the human biological network , the spine performs the preprocessing of the input , which then goes to the brain for the final output ( D ’ Mello & Dickenson , 2008 ) . By taking inspiration from the human spinal cord , Kabir et al . ( 2020 ) proposed spinal networks . As the human brain receives information from various sensory neurons , and these signals are processed by the human spine before reaching the brain . Progressive SpinalNet consists of a processing unit and concatenating unit . The input to the Progressive SpinalNet Fully Connected ( FC ) layer is the CNN features . The processing unit can be a single hidden layer or multiple hidden layers . Each layer receives some portion from the input . The output of the processing unit is concatenated , and the following processing unit receives the concatenated output as an input . The gradients are easily propagated back to the first FC layer from the last FC layer due to connections between the layers . This helps in dealing with the vanishing gradient problem . The size of each layer is progressively increased , and the output size is the same for all the FC layers . The Progressive SpinalNet can be deeper as it does not suffer from a vanishing gradient . Contradicting the traditional way of feedforward neural networks , they used forward and introduced sideways and zig-zag interactions , which leads to better network predictions . 3.3 GENERATIVE ADVERSARIAL NETWORKS . The task of Generative Adversarial Networks ( GANs ) is to learn the distribution of the training data and generate new samples . GAN learns the distribution of the input data ρdata and generates synthetic images which have nearly the same distribution as the input image . The GAN architecture consists of basically two deep neural network architecture : ( a ) : the generative network G and ( b ) : the discriminator network D. The generator takes a noise vector z as an input from a known distribution ρz , generally a uniform distribution and outputs G ( z ) which maps to the space of distribution ρg . The generator G tries to produces more realistic images as the training progresses , and the discriminator D improves its ability to discriminate between synthetic and authentic images . Finally , the generator aims to generate images that have the same distribution as the input images , i.e. , ρg = ρdata . This basically leads to min-max optimization problem where the discriminator network D is trained to maximize log ( D ( x ) ) , where D ( x ) is the output of the discriminator and the generator network is trained as to minimize log ( 1-D ( G ( z ) ) . So , the objective function for GANs is a min-max objective function : min G max D J ( D , G ) & = Ex∼px [ log ( D ( x ) ) ] + Ez , ∼pz [ log ( 1−D ( G ( Z ) ) ] Furthermore , the computation cost to train Vanilla GANs is exorbitant , and it takes a large amount of training data to learn the distribution of training data . This limits GANs and its variants to generate high fidelity images when the training data is less , and we have fewer computational resources . To train GAN with limited data and low computational resources , Liu et al . ( 2021 ) proposed lightweight GAN , which can generate high fidelity images even when the training data is significantly less ( 50-100 images ) .

Summary: This paper studies mistake-driven image classification, with the idea of applying data augumentation to classes with lower classification accuracy. Emphasizing the training on classes with low accuracy boosts the overall performance of trained CNN classifier.   The proposed method is quite straightforward, and the paper is clearly written. It first evaluates the classification accuracy of each class on the validation set, then selects the categories with lowest accracy for data augumentation with GAN, and add the augumented class into the training set for fine-tuning. Experiments are tested on 5 dataset, which demonstrate improved performance.