Analyzing Overfitting under Class Imbalance in Neural Networks for Image Segmentation

TL;DR

This paper investigates how class imbalance causes overfitting in image segmentation neural networks, leading to poor generalization of small structures, and proposes loss function modifications to improve accuracy.

Contribution

It provides new insights into overfitting under class imbalance and introduces asymmetric loss functions and regularization techniques to mitigate logit shift.

Findings

Logit activation distributions shift across decision boundaries for under-represented classes.

Proposed loss functions improve segmentation accuracy on challenging datasets.

Methods consistently outperform baseline models across various tasks.

Abstract

Class imbalance poses a challenge for developing unbiased, accurate predictive models. In particular, in image segmentation neural networks may overfit to the foreground samples from small structures, which are often heavily under-represented in the training set, leading to poor generalization. In this study, we provide new insights on the problem of overfitting under class imbalance by inspecting the network behavior. We find empirically that when training with limited data and strong class imbalance, at test time the distribution of logit activations may shift across the decision boundary, while samples of the well-represented class seem unaffected. This bias leads to a systematic under-segmentation of small structures. This phenomenon is consistently observed for different databases, tasks and network architectures. To tackle this problem, we introduce new asymmetric variants of popular loss functions and regularization techniques including a large margin loss, focal loss, adversarial training, mixup and data augmentation, which are explicitly designed to counter logit shift of the under-represented classes. Extensive experiments are conducted on several challenging segmentation tasks. Our results demonstrate that the proposed modifications to the objective function can lead to significantly improved segmentation accuracy compared to baselines and alternative approaches.