# Analyzing Overfitting under Class Imbalance in Neural Networks for Image Segmentation

arXiv: https://arxiv.org/abs/2102.10365 · HTML: https://mlsanity.com/paper/2102.10365/

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 55 references that are in the map.
- First posted on arXiv 20 Feb 2021.

## Builds on (most cited first, 29 of 29)

- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation](https://mlsanity.com/paper/1606.04797/) (2016)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Class-Balanced Loss Based on Effective Number of Samples](https://mlsanity.com/paper/1901.05555/) (2019)
- [A Survey on Deep Learning in Medical Image Analysis](https://mlsanity.com/paper/1702.05747/) (2017)
- [Learning Imbalanced Datasets with Label-Distribution-Aware Margin Loss](https://mlsanity.com/paper/1906.07413/) (2019)
- [Virtual Adversarial Training: A Regularization Method for Supervised and Semi-Supervised Learning](https://mlsanity.com/paper/1704.03976/) (2017)
- [Decoupling Representation and Classifier for Long-Tailed Recognition](https://mlsanity.com/paper/1910.09217/) (2019)
- [Large-Scale Long-Tailed Recognition in an Open World](https://mlsanity.com/paper/1904.05160/) (2019)
- [A systematic study of the class imbalance problem in convolutional neural networks](https://mlsanity.com/paper/1710.05381/) (2017)
- [BBN: Bilateral-Branch Network with Cumulative Learning for Long-Tailed Visual Recognition](https://mlsanity.com/paper/1912.02413/) (2019)
- [Large-Margin Softmax Loss for Convolutional Neural Networks](https://mlsanity.com/paper/1612.02295/) (2016)
- [Adversarial Examples Improve Image Recognition](https://mlsanity.com/paper/1911.09665/) (2019)
- [Additive Margin Softmax for Face Verification](https://mlsanity.com/paper/1801.05599/) (2018)
- [Efficient Multi-Scale 3D CNN with Fully Connected CRF for Accurate Brain Lesion Segmentation](https://mlsanity.com/paper/1603.05959/) (2016)
- [Long-Tailed Classification by Keeping the Good and Removing the Bad Momentum Causal Effect](https://mlsanity.com/paper/2009.12991/) (2020)
- [Boundary loss for highly unbalanced segmentation](https://mlsanity.com/paper/1812.07032/) (2018)
- [A Novel Focal Tversky loss function with improved Attention U-Net for lesion segmentation](https://mlsanity.com/paper/1810.07842/) (2018)
- [The KiTS19 Challenge Data: 300 Kidney Tumor Cases with Clinical Context, CT Semantic Segmentations, and Surgical Outcomes](https://mlsanity.com/paper/1904.00445/) (2019)
- [Brain Tumor Segmentation with Deep Neural Networks](https://mlsanity.com/paper/1505.03540/) (2015)
- [Striking the Right Balance with Uncertainty](https://mlsanity.com/paper/1901.07590/) (2019)
- [Imbalanced Deep Learning by Minority Class Incremental Rectification](https://mlsanity.com/paper/1804.10851/) (2018)
- [3D Segmentation with Exponential Logarithmic Loss for Highly Unbalanced Object Sizes](https://mlsanity.com/paper/1809.00076/) (2018)
- [Overfitting of neural nets under class imbalance: Analysis and improvements for segmentation](https://mlsanity.com/paper/1907.10982/) (2019)
- [Asymmetric Loss Functions and Deep Densely Connected Networks for Highly Imbalanced Medical Image Segmentation: Application to Multiple Sclerosis Lesion Detection](https://mlsanity.com/paper/1803.11078/) (2018)
- [Understanding Generalization and Optimization Performance of Deep CNNs](https://mlsanity.com/paper/1805.10767/) (2018)
- [Small Organ Segmentation in Whole-body MRI using a Two-stage FCN and Weighting Schemes](https://mlsanity.com/paper/1807.11368/) (2018)

## Built on by (most cited first, 30 of 34)

- [CarveMix: A Simple Data Augmentation Method for Brain Lesion Segmentation](https://mlsanity.com/paper/2108.06883/) (2021)
- [Bootstrapping Semi-supervised Medical Image Segmentation with Anatomical-aware Contrastive Distillation](https://mlsanity.com/paper/2206.02307/) (2022)
- [Textual Query-Driven Mask Transformer for Domain Generalized Segmentation](https://mlsanity.com/paper/2407.09033/) (2024)
- [Pseudo-Bag Mixup Augmentation for Multiple Instance Learning-Based Whole Slide Image Classification](https://mlsanity.com/paper/2306.16180/) (2023)
- [Estimating Model Performance under Domain Shifts with Class-Specific Confidence Scores](https://mlsanity.com/paper/2207.09957/) (2022)
- [ACTION++: Improving Semi-supervised Medical Image Segmentation with Adaptive Anatomical Contrast](https://mlsanity.com/paper/2304.02689/) (2023)
- [Joint Learning of Blind Super-Resolution and Crack Segmentation for Realistic Degraded Images](https://mlsanity.com/paper/2302.12491/) (2023)
- [Imbalance-Aware Culvert-Sewer Defect Segmentation Using an Enhanced Feature Pyramid Network](https://mlsanity.com/paper/2408.10181/) (2024)
- [ChessMix: Spatial Context Data Augmentation for Remote Sensing Semantic Segmentation](https://mlsanity.com/paper/2108.11535/) (2021)
- [Tackling the Class Imbalance Problem of Deep Learning Based Head and Neck Organ Segmentation](https://mlsanity.com/paper/2201.01636/) (2022)
- [Deep Clustering via Center-Oriented Margin Free-Triplet Loss for Skin Lesion Detection in Highly Imbalanced Datasets](https://mlsanity.com/paper/2204.02275/) (2022)
- [Weakly Supervised Lesion Detection and Diagnosis for Breast Cancers with Partially Annotated Ultrasound Images](https://mlsanity.com/paper/2306.06982/) (2023)
- [ViT-DeiT: An Ensemble Model for Breast Cancer Histopathological Images Classification](https://mlsanity.com/paper/2211.00749/) (2022)
- [Multi-Scale Cross Contrastive Learning for Semi-Supervised Medical Image Segmentation](https://mlsanity.com/paper/2306.14293/) (2023)
- [Joint Optimization of Class-Specific Training- and Test-Time Data Augmentation in Segmentation](https://mlsanity.com/paper/2305.19084/) (2023)
- [Deep Transfer Learning for Kidney Cancer Diagnosis](https://mlsanity.com/paper/2408.04318/) (2024)
- [Multi-organ segmentation: a progressive exploration of learning paradigms under scarce annotation](https://mlsanity.com/paper/2302.03296/) (2023)
- [MIAFEx: An Attention-based Feature Extraction Method for Medical Image Classification](https://mlsanity.com/paper/2501.08562/) (2025)
- [Exact Reformulation and Optimization for Direct Metric Optimization in Binary Imbalanced Classification](https://mlsanity.com/paper/2507.15240/) (2025)
- [Context Label Learning: Improving Background Class Representations in Semantic Segmentation](https://mlsanity.com/paper/2212.08423/) (2022)
- [Unsupervised Brain Tumor Segmentation with Image-based Prompts](https://mlsanity.com/paper/2304.01472/) (2023)
- [The Effects of Mixed Sample Data Augmentation are Class Dependent](https://mlsanity.com/paper/2307.09136/) (2023)
- [How scanning probe microscopy can be supported by Artificial Intelligence and quantum computing](https://mlsanity.com/paper/2406.19397/) (2024)
- [MSA-UNet3+: Multi-Scale Attention UNet3+ with New Supervised Prototypical Contrastive Loss for Coronary DSA Image Segmentation](https://mlsanity.com/paper/2504.05184/) (2025)
- [Beyond Fairness: Age-Harmless Parkinson's Detection via Voice](https://mlsanity.com/paper/2309.13292/) (2023)
- [Imbalanced Semi-Supervised Learning via Label Refinement and Threshold Adjustment](https://mlsanity.com/paper/2407.05370/) (2024)
- [Credit Card Fraud Detection in the Nigerian Financial Sector: A Comparison of Unsupervised TensorFlow-Based Anomaly Detection Techniques, Autoencoders and PCA Algorithm](https://mlsanity.com/paper/2407.08758/) (2024)
- [Out-of-distribution data supervision towards biomedical semantic segmentation](https://mlsanity.com/paper/2507.12105/) (2025)
- [Identifying Group Anchors in Real-World Group Interactions Under Label Scarcity](https://mlsanity.com/paper/2509.20762/) (2025)
- [A unified framework for detecting point and collective anomalies in operating system logs via collaborative transformers](https://mlsanity.com/paper/2512.23380/) (2025)

## Abstract (arXiv)

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.