MetaBalance: High-Performance Neural Networks for Class-Imbalanced Data

1 INTRODUCTION . Deep learning researchers typically compare their classification models on standard benchmark tasks , but these well-curated datasets are overwhelmingly class-balanced ( Krizhevsky et al. , 2009 ; Russakovsky et al. , 2015 ) . In contrast , real-world datasets are often highly imbalanced . For example , in applications such as cancer detection or environmental disaster prediction , a disproportionate amount of data available comes from negative cases ( Kail et al. , 2020 ; Fotouhi et al. , 2019 ) To make matters worse , methods designed for balanced datasets often fail on imbalanced ones . When the training data is severely skewed , neural networks tend to memorize a small number of samples from minority classes and do not generalize at inference . This shortcoming is particularly harmful since the minority classes are often the ones we care most about . The most common methods to handle class imbalance work by re-sampling data to achieve better balance in the training set or by re-designing the loss to attend more strongly to minority class data . However , these methods do not prevent rapid overfitting to minority samples and are sub-optimal when applied to highly over-parameterized models that memorize data quickly . We propose MetaBalance , a modification of the MAML meta-learning algorithm for training deep neural networks on class-imbalanced data . Intuitively , the proposed training scheme finds a set of network parameters such that , when fine-tuned for one iteration , the model achieves good performance on balanced data . We compare our algorithm to common techniques for class-imbalance and find that MetaBalance significantly outperforms existing methods on tasks with few samples in minority classes . In particular , we evaluate our method on image classification , credit-card fraud detection , loan default prediction , and facial recognition tasks . Finally , we explain why deep neural networks trained with MetaBalance achieve better generalization on class-imbalanced data . 2 BACKGROUND . Class imbalance is problematic because classes with poor representation may be ignored by a model at inference time . Consider , for example , a corpus of credit-card transactions in which only 1 % of transactions are fraudulent . In this case , a fraud classifier can trivially achieve 99 % accuracy simply by predicting that every transaction is non-fraudulent . When identifying fraudulent transactions is difficult , this is often the optimal strategy for a classifier to take . For this reason , data scientists often care about optimizing the performance on a more balanced dataset that does not promote this strong bias . In this section , we review existing methods , both classical and recently proposed , for training in situations of extreme class imbalance . We also give a brief overview of the MAML algorithm as it forms the foundation for our method . 2.1 LEARNING ON CLASS IMBALANCED DATA . Existing methods for handling class imbalance can be broadly split into three groups : re-sampling methods that increase the numbers of minority class samples , methods for reducing majority class samples , and classifier-level methods that modify the training routine to shift the model ’ s focus towards minority samples during training . Oversampling methods focus on generating new minority class samples from the available unbalanced data . A naive approach is to simply replicate points from the minority class , however this does not produce new information about the minority class and is known to lead to severe overfitting to over-sampled examples . To address this problem , Chawla et al . ( 2002 ) propose SMOTE , which generates unique minority samples by linearly interpolating between the existing observations from minority classes . Several improvements to SMOTE have been made with the goal of generating additional training data in a way that results in better decision boundaries after training ( Han et al. , 2005 ; Nguyen et al. , 2011 ; He et al. , 2008 ) . For example , SVMSmote generates the new minority examples along the boundary found by a support vector machine ( Han et al. , 2005 ) . SMOTE and its modifications are intended for tabular data and not for high-dimensional data such as images . However , some strong data augmentation techniques designed for preventing overfitting on images operate in similar ways . For example , mixup generates new images by taking a convex combination of images in the dataset ( Zhang et al. , 2017 ) and CutMix blends two images by cutting a patch from one image and inserting it into another ( Yun et al. , 2019 ) . Both methods produce labels for the new samples by taking a weighted average of labels of the blended images . SMOTE is closely related to mixup , with the main difference being that SMOTE only performs mixing within the minority class , while mixup intermingles samples between all classes . Kim et al . ( 2020 ) propose translating majority samples to minority samples by applying adversarial perturbations and in this way leveraging the diversity of the majority information . Finally , there is a growing body of work proposing GANs for generating realistic samples from minority classes , but training GANs is difficult , and these models are notorious for performing poorly on diverse datasets or memorizing their training data ( Shamsolmoali et al. , 2020 ; Deepshikha & Naman , 2020 ; Ali-Gombe & Elyan , 2019 ; Mullick et al. , 2019 ) . Undersampling is another common technique for handling class-imbalance . In contrast to oversampling , which adds minority class data , undersampling removes from the majority samples to form a balanced dataset . Removing data at random leads to losing critical data points in the majority class and several works propose methods for cleverly choosing samples that can be removed without losing important information about majority classes ( Lin et al. , 2017b ; Wilson , 1972 ; Tomek et al. , 1976 ) . Wilson ( 1972 ) proposes an Edited Nearest Neighbours algorithm ( ENN ) , where majority class data points that do not agree with the predictions of the KNN algorithm are removed . Another method , Cluster Centroids , undersamples by replacing clusters of the majority class , found with a k-means algorithm , by their respective centroids ( Lin et al. , 2017b ) . These methods are problematic when data is extremely high-dimensional , as nearest-neighbor classifiers tend to become uninformative in this regime ( e.g . ` 2 distance is generally not a good measure of similarity between images ) . Furthermore , undersampling prevents the user from leveraging the abundance of majority class data to learn better feature representations . Classifier-level methods modify the training routine to emphasize the minority class samples . Several disparate techniques exist in this category . For example , cost-sensitive learning works by altering the loss on minority class points either through re-weighting the loss or by altering the learning rate ( Elkan , 2001 ; Kukar et al. , 1998 ; Cui et al. , 2019 ; Lin et al. , 2017a ) . Intuitively , applying different weights to training samples is similar to over-sampling those data points with the appropriate frequencies . Other classifier-level methods include regularizers which promote large margins on minority class data or impose constraints on the “ balanced performance ” as measured on a small balanced dataset ( Sangalli et al. , 2021 ; Huang et al. , 2016 ; Li et al. , 2019 ) . Finally , there are postprocessing methods designed to re-scale the scores output by a classifier to achieve better performance ( Richard & Lippmann , 1991 ; Chan et al. , 2019 ) . 2.2 MAML . Meta-learning algorithms , originally designed for few-shot learning , train on a battery of “ tasks ” ( e.g. , classification problems with different label spaces ) with the goal of learning a common feature extractor that can be quickly fine-tuned for high performance on new tasks using little data . Model-Agnostic Meta-Learning ( MAML ) , a popular algorithm for few-shot learning , is compatible with any model and can be adapted to a variety of training algorithms ( Finn et al. , 2017 ) . The MAML pipeline contains an inner and an outer loop . At the beginning of an episode , tasks accompanied by small support and query datasets are sampled . In the inner loop , the model ’ s initial meta-parameters are fine-tuned on support data in order to perform well on the sampled tasks . Then , in the outer loop , the model is evaluated on query data , and its meta-parameters are updated so that the model performs better on the query data after fine-tuning on support data . The outer loop thus requires unrolling the inner fine-tuning procedure and computing the gradient with respect to the meta-parameters . The resulting meta-parameters have the property that they learn quickly from few-shot data when samples are scarce in a new domain . Our algorithm , MetaBalance , is based on the MAML training routine and exploits the few-shot performance of meta-learning to handle the case where minority-class data is scarce . Two works , L2RW and Meta-Weight-Net ( MW-Net ) , have applied the MAML algorithm to learn loss weights for training neural networks in the class-imbalance setting ( Ren et al. , 2018 ; Shu et al. , 2019 ) . Both methods find optimal sample weights by minimizing the meta loss on a balanced holdout set . The main difference between the works is that L2RW learns the sample weights implicitly , while MW-Net trains an explicit weighting function . Our work is different from these algorithms in two aspects : ( 1 ) instead of learning the optimal sample weights , we learn the optimal network parameters directly , ( 2 ) MetaBalance does not require a balanced holdout dataset , which might be problematic in tasks with few minority samples . Indeed , we show that MetaBalance outperforms MW-Net on tasks with extremely small minority classes and that our algorithm prevents overfitting in these scenarios . Finally , several works target class-imbalance in few-shot learning ( FSL ) . Lee et al . ( 2019 ) proposes Bayesian framework for balancing the effect of meta-learning and task-specific learning across tasks and Ochal et al . ( 2021 ) demonstrates that re-balancing strategies are still effective against imbalance in FSL . 3 METABALANCE . Our algorithm , MetaBalance is based on Model-Agnostic Meta-Learning ( MAML ) , an approach designed for few-shot learning ( Finn et al. , 2017 ) . Standard meta-learning schemes consist of “ inner ” and “ outer ” loops , each with their own , potentially distinct , loss function . In the inner loop , the model is fine-tuned by minimizing a loss function defined by a set of support data . The outer loop evaluates the fine-tuned model on a batch of query data from the same task . Our proposed MetaBalance method exploits the fact that meta-learning decouples the inner- and outer-loop loss functions , allowing us to use different class balancing strategies for each . We often find that using the natural class ( im ) balance on the inner loop support set is optimal , as it prevents the optimizer from aggressively overfitting to minority class data . At the same time , we can use rebalanced data in the query set , which guides the algorithm to achieve high accuracy on class-balanced data as training proceeds . This allows us to simultaneously combine the benefits of training on imbalanced data ( prevention of over-fitting ) with the benefits of training on balanced data ( minority classes don ’ t get ignored ) . To formalize the MetaBalance algorithm , consider neural network f with parameters θ that maps samples to predictions . In the inner loop of the MetaBalance training routine , we minimize the loss over sampled batch X to obtain new parameters θ′ = θ − α∇θLX ( fθ ) . Then , in the outer loop , we evaluate the model on a new balanced sample Z based on new parameters θ′ and compute a loss LZ on those predictions as well . We accumulate the outer loss over multiple support and query batches , fine-tuning individually on each support batch , and we finally update θ to minimize the query loss . See Algorithm 1 for details . Sampling strategies in the inner and outer loops : Algorithm 1 MetaBalance Require : Initial parameter vector , θ , outer learning rate , η , inner learning rate , γ. for epochs = 1 , . . . , epoch do LossZ = 0 for step = 1 , . . . , metaStep do Sample unbalanced support batch , X , and balanced query batch , Z Evaluate∇θLXf ( θ ) Compute adapted parameters with gradient descent : θ′ = θ − γ∇θLXf ( θ ) Accumulate the loss over Z : LossZ = LossZ + LZf ( θ′ ) + βLXf ( θ ) end for Update θ ← θ − η∇θLossZ end for The benefits of MetaBalance come from the fact that we can choose different sampling strategies for sampling support and query data in the inner and outer loops . We can improve the performance beyond baselines , often dramatically , by using complementary sampling strategies for the two losses . For each loss , we consider using the natural class balance , over sampling , undersampling , SMOTE , SVMSmote , ENN , and CC ( Chawla et al. , 2002 ; Nguyen et al. , 2011 ; Wilson , 1972 ; Li et al. , 2019 ) . Results are presented and discussed in the following sections .

Summary: The paper proposes a meta-learning-based method to learn under class imbalance. Class imbalance is ubiquitous in many real applications and neural networks tend to learn biased models toward majority classes. Analogous to a popular meta-learning method called MAML, the proposed method consists of 2 nested training loops: outer on the query dataset and inner on the support dataset. In the inner loop, support datasets are randomly sampled from the original dataset by retaining the class imbalance and candidate parameters are obtained by updating the current parameters with the gradient of the loss on the support dataset. During the inner loop, a global loss on both the query and the current support datasets is accumulated.  The model parameters are updated after running the inner loop for a pre-defined number of steps using the gradient of the accumulated loss. By doing so, it is expected to learn a model from different subsamples (support set) of the imbalanced datasets while performing well on a balanced dataset.  The proposed method is compared with several undersampling/oversampling/re-weighting based methods to handle class imbalance. The presented results show that MetaBalance improve accuracy on minority classes in many cases and avoid overfitting to minority/majority class samples.