Label Dropout: Improved Deep Learning Echocardiography Segmentation Using Multiple Datasets With Domain Shift and Partial Labelling

1. Problem Statement

The paper addresses the challenge of training robust deep learning models for echocardiography (echo) segmentation using multiple diverse datasets. While deep learning can automate the extraction of cardiac biomarkers, clinical translation requires models that generalize across different scanners, operators, and imaging protocols.

Two primary obstacles exist when combining multiple public echo datasets (e.g., CAMUS, Unity Imaging, EchoNet Dynamic):

1. Domain Shift: Variations in image acquisition (scanner types, cone shapes, intensity inhomogeneities) create distribution shifts between datasets.
2. Partial Labelling: Different datasets contain different sets of ground truth labels. For example, one dataset may label the Left Ventricle (LV) and Left Atrium (LA), while another only labels the LV. When combined, the training data is "partially labelled," meaning some structures are missing ground truth in specific samples.

The Core Issue: The authors identify that standard approaches to handling partial labels (specifically Adaptive Cross-Entropy Loss) fail in the presence of domain shift. Instead of learning robust features, the model engages in shortcut learning: it associates specific domain characteristics (e.g., "this image looks like Unity Imaging") with the presence or absence of specific labels. Consequently, the model learns to predict labels only when they are historically present in that specific domain, leading to catastrophic failure when those labels are missing in other domains.

2. Methodology

Datasets and Pre-processing

The study utilized three public 2D echo datasets:

• CAMUS: 500 subjects, A2C/A4C views, labels for LV, LVM, LA.
• Unity Imaging: 400 images, A2C view, labels for LV, LVM, LA.
• EchoNet Dynamic: ~20k images, A2C/A4C views, labels for LV only.
• Pre-processing: Images were resized to $256 \times 256$. Extraneous information outside the ultrasound cone was removed using an nnU-Net trained to segment the cone.

Baseline Models

• Architecture: U-Net.
• Loss Functions:
  1. Standard CCE: Cross-Entropy Loss calculated over all classes (including background). This fails on partial labels due to supervision conflicts (labeling a structure as "background" in one sample and "foreground" in another).
  2. Adaptive Loss: A modified Cross-Entropy Loss where terms for missing ground truth labels are removed from the loss calculation. This is the state-of-the-art approach for partial labels but was hypothesized to be susceptible to shortcut learning under domain shift.

Proposed Solution: Label Dropout

To break the correlation between domain characteristics and label presence, the authors propose Label Dropout:

• Mechanism: During training, for samples containing partially labelled classes, there is a probability $p$ that a specific missing label is artificially "dropped out" (set to background) in the ground truth mask, even if it was originally present.
• Logic: By randomly removing labels that should be present, the model is forced to learn the visual features of the structure regardless of whether the label is present in the specific dataset or the specific sample. This prevents the model from relying on domain cues to decide whether to predict a label.
• Implementation: Applied only to partially labelled classes. If multiple classes are missing, one is chosen randomly to drop.

3. Key Contributions

1. Identification of Shortcut Learning in Partial Labeling: The paper is the first to demonstrate that adaptive loss functions, while effective in controlled environments, fail in multi-domain settings because models learn to associate domain identity with label existence rather than learning anatomical features.
2. Label Dropout Scheme: A novel, simple, and effective technique to disrupt the domain-label correlation, forcing the model to learn invariant representations of cardiac structures.
3. Comprehensive Validation: The study provides a rigorous series of experiments isolating the effects of domain shift, partial labeling, and the proposed solution.

4. Experimental Results

The authors conducted four experiments to validate their hypothesis and method:

• Experiment 1 (Need for Diversity): Training on a single dataset resulted in poor cross-domain performance, confirming the necessity of multi-dataset training.
• Experiment 2 (The Failure of Adaptive Loss): When training on the combined, partially labelled datasets:
  • The Standard Loss model failed completely on missing labels.
  • The Adaptive Loss model (without dropout) also failed to predict missing structures (e.g., LVM and LA in EchoNet Dynamic) because it learned that "EchoNet Dynamic images do not have LVM/LA labels."
  • Data augmentation alone could not overcome this shortcut learning.
• Experiment 3 (Controlled Environment): Using only the CAMUS dataset with artificially removed labels (no domain shift):
  • The Adaptive Loss model performed comparably to a fully labelled benchmark (Dice ~0.86 vs 0.87).
  • This confirmed that Adaptive Loss works if there is no domain shift, isolating domain shift as the cause of failure in Experiment 2.
• Experiment 4 (Label Dropout Efficacy):
  • Setup: Combined CAMUS and Unity Imaging (with LA artificially removed from Unity).
  • Result: Introducing label dropout (even at low probabilities like 10%) significantly improved performance.
  • Quantitative Gains: On the EchoNet Dynamic test set (trained on all 3 datasets with Label Dropout):
    • LVM Dice: Increased from 0.319 (Adaptive only) to 0.518 (Adaptive + Dropout). This is a 62% relative improvement.
    • LA Dice: Increased from 0.553 to 0.689. This is a 25% relative improvement.
    • Statistical significance was confirmed via Wilcoxon Signed Rank Test ($p < 0.05$).

5. Significance and Conclusion

This work is critical for the clinical deployment of AI in echocardiography. It highlights that simply combining diverse datasets with existing loss functions is insufficient and can lead to brittle models that fail in real-world scenarios.

• Robustness: Label Dropout enables the training of models that are robust to both domain shifts and missing annotations, a common reality in medical imaging where datasets are rarely perfectly aligned.
• Generalizability: The technique allows for the creation of more generalizable models without requiring expensive re-annotation of existing datasets.
• Future Work: The authors plan to test this on other architectures (e.g., Transformers) and other imaging modalities (e.g., MRI).

In summary, the paper proves that Label Dropout is a necessary intervention to prevent shortcut learning when training deep segmentation models on heterogeneous, partially labelled medical data, resulting in substantial performance gains for clinically relevant cardiac structures.