Analyzing Overfitting under Class Imbalance in Neural Networks for Image Segmentation

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.

arXiv commentPublished in IEEE Transactions on Medical Imaging

References in corpus (8)

• SMOTE: Synthetic Minority Over-sampling Technique
• A Survey on Deep Learning in Medical Image Analysis
• Efficient Multi-Scale 3D CNN with Fully Connected CRF for Accurate Brain Lesion Segmentation
• Brain Tumor Segmentation with Deep Neural Networks
• A systematic study of the class imbalance problem in convolutional neural networks
• Additive Margin Softmax for Face Verification
• Asymmetric Loss Functions and Deep Densely Connected Networks for Highly Imbalanced Medical Image Segmentation: Application to Multiple Sclerosis Lesion Detection
• 3D Segmentation with Exponential Logarithmic Loss for Highly Unbalanced Object Sizes

Cited by in corpus (6)

• Pseudo-Bag Mixup Augmentation for Multiple Instance Learning-Based Whole Slide Image Classification
• CarveMix: A Simple Data Augmentation Method for Brain Lesion Segmentation
• Joint Learning of Blind Super-Resolution and Crack Segmentation for Realistic Degraded Images
• Multi-organ segmentation: a progressive exploration of learning paradigms under scarce annotation
• MIAFEx: An Attention-based Feature Extraction Method for Medical Image Classification
• A unified framework for detecting point and collective anomalies in operating system logs via collaborative transformers