Analyzing Overfitting Under Class Imbalance in Neural Networks for Image Segmentation.

Area of Science:

• Computer Vision
• Machine Learning
• Artificial Intelligence

Background:

• Class imbalance in training data leads to neural network overfitting, particularly affecting the segmentation of small, under-represented structures.
• This overfitting results in poor model generalization and biased predictions, a common issue in image segmentation tasks.
• The study investigates the underlying network behavior contributing to this bias, specifically focusing on logit activation shifts.

Purpose of the Study:

• To provide new insights into the problem of overfitting caused by class imbalance in neural networks.
• To introduce novel asymmetric loss functions and regularization techniques designed to mitigate logit shift.
• To improve the accuracy and generalization of image segmentation models trained on imbalanced datasets.

Main Methods:

• Empirical analysis of neural network behavior under class imbalance, observing logit activation shifts.
• Development of asymmetric variants of loss functions (e.g., large margin loss, focal loss) and regularization techniques (e.g., adversarial training, mixup, data augmentation).
• Extensive experimentation on challenging image segmentation tasks to evaluate the effectiveness of proposed methods.

Main Results:

• Observed a consistent shift in logit activation distributions across the decision boundary for under-represented classes during testing.
• Demonstrated that the proposed asymmetric loss functions and regularization techniques effectively counter logit shift.
• Achieved significantly improved segmentation accuracy for small structures compared to baseline and alternative approaches.

Conclusions:

• The proposed modifications to objective functions effectively address the logit shift problem in imbalanced datasets.
• These techniques lead to enhanced segmentation accuracy and better generalization for under-represented classes.
• The study offers a robust solution for developing unbiased and accurate predictive models in image segmentation.