[Paper Review] MetaBalance: High-Performance Neural Networks for Class-Imbalanced Data

Arpit Bansal, Micah Goldblum+4 more|arXiv (Cornell University)|Jun 17, 2021

Imbalanced Data Classification Techniques34 references4 citations

TL;DR

MetaBalance proposes a meta-learning approach to train deep neural networks on class-imbalanced data by decoupling training strategies in inner and outer optimization loops, significantly improving generalization on minority classes without overfitting. It outperforms standard re-sampling and re-weighting methods across image classification, fraud detection, loan default prediction, and facial recognition tasks.

ABSTRACT

Class-imbalanced data, in which some classes contain far more samples than others, is ubiquitous in real-world applications. Standard techniques for handling class-imbalance usually work by training on a re-weighted loss or on re-balanced data. Unfortunately, training overparameterized neural networks on such objectives causes rapid memorization of minority class data. To avoid this trap, we harness meta-learning, which uses both an ''outer-loop'' and an ''inner-loop'' loss, each of which may be balanced using different strategies. We evaluate our method, MetaBalance, on image classification, credit-card fraud detection, loan default prediction, and facial recognition tasks with severely imbalanced data, and we find that MetaBalance outperforms a wide array of popular re-sampling strategies.

Motivation & Objective

• Address the challenge of poor generalization in deep neural networks trained on highly imbalanced datasets, where minority classes are often ignored or overfitted.
• Overcome the limitations of standard re-sampling and re-weighting techniques, which fail to prevent overfitting in over-parameterized models.
• Develop a training framework that leverages meta-learning to learn robust representations by balancing inner-loop and outer-loop optimization strategies.
• Improve model performance on minority classes without sacrificing accuracy on majority classes, achieving better overall generalization.
• Demonstrate the effectiveness of MetaBalance across diverse real-world applications with severe class imbalance, including image classification and financial risk prediction.

Proposed method

• Employ a meta-learning framework with separate inner and outer optimization loops, allowing distinct sampling and loss strategies in each.
• In the inner loop, fine-tune the model on a balanced subset of data (e.g., using oversampling or data augmentation) to adapt to minority class patterns.
• In the outer loop, optimize the initial model parameters using a standard, unweighted loss on the full imbalanced dataset to maintain generalization.
• Use a bi-level optimization objective where the outer loop minimizes validation loss computed after one inner-loop update, enabling the model to learn parameters robust to class imbalance.
• Integrate data augmentation techniques like mixup and CutMix in the inner loop to generate diverse minority-class samples without explicit oversampling.
• Decouple the training strategy across loops: use balanced data for fast adaptation (inner) and imbalanced data for meta-optimization (outer), preventing memorization of minority samples.

Experimental results

Research questions

• RQ1 Can meta-learning be effectively leveraged to improve generalization in deep neural networks trained on highly imbalanced data?
• RQ2 Does decoupling sampling and loss strategies across inner and outer optimization loops lead to better performance on minority classes compared to standard re-sampling or re-weighting?
• RQ3 How does MetaBalance compare to naive training, oversampling, and undersampling in terms of accuracy across both majority and minority classes?
• RQ4 To what extent does MetaBalance prevent overfitting to minority class samples, especially in over-parameterized models?
• RQ5 Can MetaBalance generalize across diverse domains such as image classification, fraud detection, and facial recognition with severe class imbalance?

Key findings

• MetaBalance achieves 29% overall test accuracy on a severely imbalanced CIFAR-10 variant, significantly outperforming the naive model (16.14%) and prior methods.
• On the same imbalanced CIFAR-10 dataset, MetaBalance attains an average of 30% accuracy on minority classes, compared to only 14.3% for the prior state-of-the-art with prior adjustment.
• The model trained with MetaBalance maintains 60% accuracy on the majority class at both train and test time, avoiding the 98.3% overfitting seen in the naive model.
• Even when the threshold of the naive model is adjusted to match MetaBalance’s majority class accuracy, its overall performance (27.03%) remains below MetaBalance’s 29%.
• MetaBalance reduces overfitting to minority class samples, avoiding the gerrymandering of decision boundaries seen in standard training, which memorizes rare instances.
• Undersampling methods perform poorly on the majority class (due to massive data loss), while MetaBalance maintains strong performance on both majority and minority classes, achieving superior overall accuracy.