AutoBalance: Optimized Loss Functions for Imbalanced Data

TL;DR

AutoBalance is a bi-level optimization framework that automatically designs loss functions to improve accuracy and fairness in imbalanced and group-sensitive classification tasks, addressing overfitting and fairness concerns in deep learning.

Contribution

It introduces a novel bi-level optimization approach for automatic loss function design tailored for imbalanced and sensitive data scenarios.

Findings

AutoBalance outperforms state-of-the-art methods in imbalanced classification.

The approach improves fairness and accuracy simultaneously.

Theoretical insights support the benefits of the designed loss functions.

Abstract

Imbalanced datasets are commonplace in modern machine learning problems. The presence of under-represented classes or groups with sensitive attributes results in concerns about generalization and fairness. Such concerns are further exacerbated by the fact that large capacity deep nets can perfectly fit the training data and appear to achieve perfect accuracy and fairness during training, but perform poorly during test. To address these challenges, we propose AutoBalance, a bi-level optimization framework that automatically designs a training loss function to optimize a blend of accuracy and fairness-seeking objectives. Specifically, a lower-level problem trains the model weights, and an upper-level problem tunes the loss function by monitoring and optimizing the desired objective over the validation data. Our loss design enables personalized treatment for classes/groups by employing a parametric cross-entropy loss and individualized data augmentation schemes. We evaluate the benefits and performance of our approach for the application scenarios of imbalanced and group-sensitive classification. Extensive empirical evaluations demonstrate the benefits of AutoBalance over state-of-the-art approaches. Our experimental findings are complemented with theoretical insights on loss function design and the benefits of train-validation split. All code is available open-source.