PaperlayerView exact revision on arXiv

Source-linked AI summary

Analyzing Overfitting under Class Imbalance in Neural Networks for Image Segmentation

Zeju Li, Konstantinos Kamnitsas, Ben Glocker

arXiv:2102.10365v1cs.CV

TL;DR

Strong class imbalance and limited data can cause segmentation networks to overfit under-represented foreground structures, but the resulting behavior is not fully understood. The paper inspects logit distributions, identifies asymmetric shifts toward the decision boundary, and introduces asymmetric loss and regularization variants. Across challenging segmentation tasks, these modifications improve segmentation performance and sensitivity for under-represented classes.

  • Problem

    Class imbalance can cause neural networks to overfit under-represented foreground samples, producing poor generalization and under-segmentation of small structures.

  • Method

    The study inspects network behavior and introduces asymmetric variants of large-margin and focal losses that counter logit shifts in under-represented classes.

  • Results

    The proposed methods substantially improve segmentation performance, including DSC and sensitivity, and outperform previous solutions under class imbalance.

  • Takeaways & Limitations

    Considering biased logit behavior when designing objectives can help build more unbiased networks for imbalanced segmentation.

  • Takeaways & Limitations

    Combined regularization requires smaller hyper-parameter magnitudes because excessive regularization can prevent convergence and harm training performance.

Abstract

from arXiv · show

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.

I. INTRODUCTION

The study examines how strong class imbalance and limited data affect neural-network behavior in image segmentation. It identifies asymmetric logit shifts for under-represented classes and proposes tailored losses and regularization to address them.

  • Problem: Class imbalance is common in medical image segmentation, where foreground regions of interest can be heavily under-represented and ratios of 100:1 or higher are not uncommon.The paper defines the imbalance ratio using background and foreground pixel counts.
  • Problem: With imbalanced training data, models can overfit under-represented classes and generalize poorly, but the resulting network behavior is not well understood.The paper frames this behavioral gap as the motivation for inspecting logits under different training-data amounts.
  • Findings: The study finds that unseen under-represented samples shift toward or across the decision boundary, while over-represented samples remain largely unaffected, causing under-segmentation.The asymmetric shift is presented as a mechanism linking overfitting under imbalance to biased test predictions.
  • Findings: Across four segmentation tasks and two architectures, the authors report decreased sensitivity specifically for under-represented classes under overfitting.The contribution is positioned as a consistent observation rather than a task-specific effect.
  • Approach: The paper proposes asymmetric variants of five loss functions and regularization techniques that explicitly account for the biased behavior of under-represented classes.The modifications target better feature embeddings and improved segmentation accuracy for small structures of interest.
  • Related approaches: Existing re-weighting, data-synthesis, training-strategy, and general regularization approaches do not explicitly address asymmetric logit shifts and show limited success in this setting.The paper contrasts its behavior-focused approach with methods developed for broader imbalance or general overfitting problems.

III. OVERFITTING UNDER CLASS IMBALANCE AND ITS EFFECT ON SEGMENTATION PERFORMANCE

Experiments across strongly imbalanced segmentation tasks show that reducing training data increases overfitting and selectively harms foreground sensitivity. Logit distributions shift for under-represented samples, motivating asymmetric training modifications designed to keep their activations correctly classified.

  • Experimental setup: Experiments use BRATS, ATLAS, small-organ, and KiTS data to study strongly imbalanced segmentation with varying amounts of training data.The study examines challenging tasks spanning four datasets.
  • Segmentation performance: With less training data, test segmentation accuracy decreases while training accuracy increases, indicating easier overfitting.The paper evaluates this pattern using DSC and related sensitivity and precision measures.
  • Segmentation performance: Overfitting primarily reduces sensitivity while precision remains largely stable, and foreground specificity stays above 0.999 across settings.The stable background predictions contrast with the degradation for under-represented foreground classes.
  • General behavior: Across four datasets and model architectures, overfit models systematically under-segment under-represented foreground classes on unseen test data.The authors describe this as a consistent behavior rather than a task- or architecture-specific outcome.
  • Logit distribution shift: Foreground logits shift toward and sometimes across the decision boundary, whereas background logits remain comparatively stable.The shift is observed while monitoring training and unseen test samples in multiple segmentation tasks.
  • Logit distribution shift: Less training data produces larger logit shifts, and the rarest foreground structures shift the most.The paper reports correlations between shift magnitude, overfitting, and foreground-structure size.
  • Mitigation: The proposed asymmetric modifications add a bias for under-represented classes and aim to keep their logits away from the decision boundary.Correct predictions can remain intact when foreground logits move toward, but do not cross, the boundary.

A. Asymmetric large margin loss

The asymmetric large margin loss applies the margin only to rare foreground classes, moving their logits farther from the decision boundary while preserving the class-specific formulation of segmentation losses.

  • A. Asymmetric large margin loss: Logit activations are the final-layer outputs used by softmax to obtain class probabilities in the cross-entropy formulation.The notation uses yij for the one-hot target and pij for the predicted probability.
  • A. Asymmetric large margin loss: The large margin loss increases distances between logits for different classes by adding a margin to every class.Its symmetric form leaves the decision boundary centered.
  • A. Asymmetric large margin loss: Class imbalance can shift unseen foreground samples toward the background class, so the asymmetric variant moves the decision boundary closer to background.This targets the observed direction of foreground logit shift.
  • A. Asymmetric large margin loss: The asymmetric modification sets the margin only for rare classes rather than applying it symmetrically to every class.The rare-class indicator is defined using a one-hot vector r.
  • A. Asymmetric large margin loss: The same symmetric and asymmetric large margin constructions can be derived for DSC loss by substituting the corresponding formulation.The paper formulates cross-entropy and DSC losses with class-specific terms.

B. Asymmetric focal loss

The asymmetric focal variants retain focal attenuation for background samples but remove it for foreground classes, helping keep under-represented foreground logits away from the decision boundary.

  • B. Asymmetric focal loss: Focal loss reduces the weight of well-classified samples and focuses training on samples near the decision boundary.The focus is controlled by the hyper-parameter γ.
  • B. Asymmetric focal loss: Symmetric focal loss can leave every class near the decision boundary, making unseen foreground samples more likely to cross it.This motivates an asymmetric modification for class imbalance.
  • B. Asymmetric focal loss: The asymmetric focal loss removes attenuation for the foreground class to keep its logits away from the decision boundary.Attenuation remains part of the loss design for the other classes.
  • B. Asymmetric focal loss: The proposed focal DSC loss reduces penalties for well-classified samples while preserving a loss magnitude similar to the original formulation.Its asymmetric version is obtained by removing attenuation for the foreground class.

C. Asymmetric adversarial training

The asymmetric adversarial-training variant adds adversarial examples while shaping training to produce a larger separation between foreground logits and the decision boundary.

  • C. Asymmetric adversarial training: Adversarial training augments optimization with difficult samples generated by perturbing the input.The adversarial loss combines the original loss with the loss on the generated sample.
  • C. Asymmetric adversarial training: The chosen loss can be cross-entropy, DSC, or both, while perturbation magnitude and range control the adversarial examples.The direction of the generated samples is denoted by dadv.
  • C. Asymmetric adversarial training: Symmetric adversarial training may preserve the decision boundary and allow unseen foreground samples to shift toward background.The asymmetric variant instead aims to create more space between foreground logits and the boundary.

D. Asymmetric mixup

The asymmetric mixup strategy generates selected hard-label samples near foreground examples, while asymmetric augmentation reduces transformations of already abundant background classes.

  • D. Asymmetric mixup: Mixup generates additional training examples by linearly combining pairs of images and their labels.The mixing coefficient λ is sampled from a beta distribution.
  • D. Asymmetric mixup: Standard mixup centers the decision boundary between classes, which helps little under the studied imbalance.The proposed modification instead uses hard labels for selected augmented samples near foreground examples.
  • D. Asymmetric mixup: Asymmetric mixup labels an augmented sample as foreground when it remains sufficiently close to a foreground sample according to the margin conditions.The conditions require λ or 1 − λ to exceed the margin m and identify one source as foreground.
  • D. Asymmetric mixup: The asymmetric mixup design keeps the decision boundary away from the foreground class.Training samples whose generated label is zero are not used for updating the model.
  • D. Asymmetric mixup: Asymmetric augmentation reduces transformed samples for background classes because additional synthesized background data is unnecessary under strong imbalance.The foreground class receives the broader augmentation treatment.

F. The combination of asymmetric techniques

The study combines asymmetric losses and augmentation strategies to counter overfitting under class imbalance, evaluating them across multiple segmentation tasks and training settings.

  • F. The combination of asymmetric techniques: Asymmetric adversarial training, mixup, and augmentation generate additional training samples, while asymmetric large margin and focal losses modify optimization.The techniques can be integrated into one model to introduce variance or foreground bias from multiple perspectives.
  • F. The combination of asymmetric techniques: The combined objective can also incorporate a Dice similarity coefficient loss alongside the asymmetric techniques.A combined DSC loss is formulated analogously to the other combined objectives.
  • F. The combination of asymmetric techniques: The experiments cover brain tumor, stroke lesion, organ, and kidney tumor segmentation using different architectures and training scenarios.The study keeps baseline and modified-method hyper-parameters consistent and varies regularization settings across tasks.

3) Small organ segmentation:

Experiments on organ and kidney tumor segmentation evaluate asymmetric techniques under limited-data settings, with sensitivity and DSC emphasized despite possible distance-metric trade-offs.

  • 3) Small organ segmentation:: The organ experiments use 30 CT scans, test on 10 cases, and compare models trained with 20 cases or 5 cases.The task segments thirteen abdominal organs with DeepMedic under full and reduced training-set sizes.
  • 3) Small organ segmentation:: Kidney tumor experiments use KiTS19 with 210 CT images and a 3D U-Net employing extensive intensity and geometric augmentation.DeepMedic with similar augmentation produced lower accuracy, so asymmetric regularization was evaluated primarily with the U-Net.
  • 3) Small organ segmentation:: The evaluation reports DSC, sensitivity, precision, and 95% Hausdorff distance, with higher sensitivity and DSC sometimes accompanied by worse Hausdorff distance.False positives can increase distance-based metrics even when initial segmentation quality improves.
  • 3) Small organ segmentation:: Across tasks, asymmetric methods improve sensitivity for under-represented classes, whereas prior class-imbalance approaches show little sensitivity improvement.The results support mitigating overfitting under class imbalance through asymmetric modifications.
  • 3) Small organ segmentation:: Increasing tumor-sample weighting from 50% to 80% reduces performance, while F-score training produces only slight sensitivity gains under severe imbalance.The baseline experiments attribute the degradation to memorization and overfitting of under-represented samples.
  • 3) Small organ segmentation:: Validation-based threshold adjustment is unreliable when unseen test logits shift relative to training and validation distributions.Without test-set ground truth, thresholds cannot be optimized for the specific test distribution.

4) Asymmetric adversarial training:

Asymmetric regularization improves logit separation and segmentation under class imbalance, while the benefits of individual methods depend on augmentation, modality, and task conditions.

  • 4) Asymmetric adversarial training:: Without data augmentation, modified adversarial training further improves sensitivity and substantially boosts performance when little training data is available.With augmentation, original adversarial training adds little, whereas the proposed modifications consistently improve segmentation.
  • 4) Asymmetric adversarial training:: Original mixup can improve the well-represented kidney class but lowers sensitivity for the under-represented class.Asymmetric mixup helps BRATS substantially but is less effective for single-channel ATLAS and KiTS.
  • 4) Asymmetric adversarial training:: Asymmetric augmentation improves DSC and sensitivity for under-represented classes in most cases but can hurt performance when augmentation is strong and training data is sufficient.The reported explanation is that strong asymmetric augmentation may make the background under-represented.
  • 4) Asymmetric adversarial training:: Combining asymmetric techniques gives the overall best segmentation results with improved sensitivity in all cases, unlike combining symmetric counterparts.The symmetric combination often decreases sensitivity and does not mitigate overfitting.
  • 4) Asymmetric adversarial training:: On BRATS with 5% training data, asymmetric modifications expand foreground logit variance and keep unseen foreground logits on the correct side of the decision boundary.Asymmetric mixup yields the largest improvement by enlarging the foreground logit region and reducing unseen-foreground logit shift.
  • 4) Asymmetric adversarial training:: On KiTS with 10% training data, asymmetric augmentation further reduces unseen tumor logit shift beyond standard augmentation.The asymmetric large margin loss also reduces tumor-logit shift toward the kidney class.
  • 4) Asymmetric adversarial training:: Across tasks and two network architectures, the proposed asymmetric methods substantially improve DSC and sensitivity under class imbalance.The conclusion reports that they outperform previous solutions and motivate inspecting logit distributions as a network-analysis tool.

SUPPLEMENTARY MATERIAL

The supplementary analysis interprets focal and large margin losses as sample-level re-weighting mechanisms whose gradient effects shape how logits are distributed around the decision boundary.

  • SUPPLEMENTARY MATERIAL: The supplementary material derives the gradient formulations for cross-entropy, focal loss, and large margin loss from normalized network probabilities.The notation uses x_i, y_i, p_i, z_i, θ, and sample-specific weights to connect losses to parameter gradients.
  • SUPPLEMENTARY MATERIAL: Sample-level re-weighting multiplies each sample loss by a scalar, changing the network-output gradient and fitting higher-weight samples more strongly.The supplementary derivation separates loss-design effects from architecture-related gradient terms.
  • SUPPLEMENTARY MATERIAL: Focal loss decreases the weights of well-classified samples, reducing how strongly the model fits easy cases.Its scalar weight decreases as the product of predicted probability and label increases.
  • SUPPLEMENTARY MATERIAL: Large margin loss increases gradients, especially for well-classified samples, and keeps network-output distributions farther from the decision boundary.Its scalar weight grows with the margin-adjusted network output.

VIII. THE MAGNITUDE OF FOCAL DSC LOSS

The proposed focal DSC loss behaves similarly to the original focal loss while preserving a comparable loss magnitude, supporting its use with combined loss functions. The experiments emphasize reduced regularization when asymmetric techniques are combined.

  • Focal DSC loss behavior: The proposed focal DSC loss has similar behavior to the original focal loss.This similarity is reported for the formulation compared in Figure 9.
  • Focal DSC loss behavior: Its loss magnitude does not change substantially compared with existing solutions, making it easier to combine with other losses.This property is particularly relevant to experiments using 3D U-Net, which combines cross-entropy and DSC loss.
  • Combined regularization: Combining asymmetric regularization techniques can over-regularize the network and prevent convergence or degrade training performance.The experiments therefore use smaller regularization magnitudes for combined regularization.

X. SENSITIVITY ANALYSIS

The sensitivity analyses examine hyper-parameter choices across segmentation settings and show that asymmetric methods generally improve results, though their effects depend on the dataset, regularization magnitude, post-processing, and rare-class characteristics.

  • Hyper-parameter sensitivity: Asymmetric large margin loss generally improves results, but some hyper-parameters can produce unexpected performance.The experiments use DeepMedic with 5% BRATS training data, and m = 0.5 is identified as an example of an unfavorable choice.
  • Hyper-parameter sensitivity: Asymmetric adversarial training improves results with larger perturbations, while asymmetric mixup improves results with m < 0.8.The reported favorable condition is l > 2.5 for asymmetric adversarial training.
  • Hyper-parameter sensitivity: Asymmetric focal loss is robust across the selected hyper-parameters and is recommended as the first choice for new applications.This recommendation comes from the controlled sensitivity experiments.
  • Dataset-dependent effects: Asymmetric mixup is most effective for BRATS but provides limited improvements for ATLAS and KiTS.The authors associate this difference with multi-channel information in BRATS and substantial foreground-background intensity overlap in ATLAS and KiTS.
  • Multi-class segmentation: In abdominal segmentation, asymmetric focal loss improves overall results and increases sensitivity for most rare classes.The evaluation uses 25% of the training data and identifies classes 4, 5, 8, 9, 10, 11, 12, and 13 as rare.
  • Multi-class segmentation: 4.9%: asymmetric focal loss improves the average DSC of rare abdominal classes.The method nevertheless decreases esophagus performance because post-processing can remove correct regions for this very small structure.
  • Post-processing effects: Without post-processing, asymmetric regularization can improve DSC while worsening Hausdorff distance through additional distant false positives.Connected-component post-processing can remove many such predictions and yield better or similar HD in the main results.
Loading 2102.10365v1…