Analyzing overfitting under class imbalance in neural networks for image segmentation
Author(s)
Li, Z
Kamnitsas, K
Glocker, B
Type
Journal Article
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 underrepresented 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 underrepresented 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.
unbiased, accurate predictive models. In particular, in image
segmentation neural networks may overfit to the foreground
samples from small structures, which are often heavily underrepresented 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 underrepresented 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.
Date Issued
2021-03-01
Date Acceptance
2020-12-19
Citation
IEEE Transactions on Medical Imaging, 2021, 40 (3), pp.1065-1077
ISSN
0278-0062
Publisher
Institute of Electrical and Electronics Engineers
Start Page
1065
End Page
1077
Journal / Book Title
IEEE Transactions on Medical Imaging
Volume
40
Issue
3
Copyright Statement
© 2020 IEEE. Personal use of this material is permitted. Permission from IEEE must be obtained for all other uses, in any current or future media, including reprinting/republishing this material for advertising or promotional purposes, creating new collective works, for resale or redistribution to servers or lists, or reuse of any copyrighted component of this work in other works.
Sponsor
Commission of the European Communities
Innovate UK
Identifier
https://ieeexplore.ieee.org/document/9302891
Grant Number
H2020 - 757173
104691
Subjects
cs.CV
cs.CV
Nuclear Medicine & Medical Imaging
08 Information and Computing Sciences
09 Engineering
Publication Status
Published
Date Publish Online
2020-12-22