Cross-Dataset Generalization in Breast MRI Tumor Classification via Class-Wise Dataset Mixing

Technical Summary: Cross-Dataset Generalization in Breast MRI Tumor Classification via Class-Wise Dataset Mixing

Problem Statement
Breast MRI is a highly sensitive tool for tumor detection, yet its clinical utility is hindered by the time-intensive nature of reviewing hundreds or thousands of slices per exam. While deep learning models often achieve high internal accuracy, they frequently fail when deployed across different institutions due to domain shift and dataset-origin bias. Models may inadvertently learn "shortcuts" based on scanner vendors, acquisition protocols, or reconstruction pipelines rather than genuine tumor features. This issue is exacerbated in breast MRI because public datasets (e.g., Duke, fastMRI, MAMA-MIA) differ significantly in patient populations, annotation styles, and file formats. The core problem addressed is whether standard deep learning pipelines can distinguish between transferable tumor cues and dataset-specific artifacts when trained on multi-source data.

Methodology
The study evaluates two architectures, EfficientNet-B3 (convolutional) and WaveViT-Small (transformer-based with wavelet transforms), using a strict external validation protocol.

• Datasets:
  • Training Sources: Duke Breast Cancer MRI (DICOM, lesion bounding boxes) and fastMRI (HDF5, case-level labels).
  • External Validation: MAMA-MIA (multi-center DCE-MRI, expert-corrected segmentations), used exclusively for testing.
• Data Preparation: All datasets were converted to a common 256×256 three-channel slice-based representation. Patient-level splitting (80/20) was enforced to prevent data leakage, ensuring slices from the same patient did not appear in both training and test sets.
• Experimental Design:
  1. Confounded Baseline: A training set was constructed where the class label was perfectly correlated with the dataset origin (all malignant slices from Duke, all negative slices from fastMRI). This setup was designed to expose if models learn dataset identity as a proxy for pathology.
  2. Proposed Approach (Class-Wise Mixing): A "Mixed Dataset" was constructed where both classes (positive and negative) contained samples from both Duke and fastMRI. Specifically, the positive class included all Duke positives plus 40% of available fastMRI positives, and the negative class included all fastMRI negatives plus 40% of available Duke negatives.
• Training Protocol: Both experiments used identical hyperparameters (AdamW optimizer, learning rate $10^{-4}$, weighted cross-entropy loss, batch size 64) and augmentation strategies (random rotation, vertical flip) to isolate the effect of data composition.

Key Results
The experiments were evaluated on the independent MAMA-MIA cohort.

• Confounded Setup: When trained on the dataset where label and source were perfectly correlated, both models failed to generalize.
  • WaveViT-Small: Accuracy dropped to 0.5048 (near chance), with a recall of 0.9986, indicating the model predicted nearly all external slices as positive.
  • EfficientNet-B3: Accuracy dropped to 0.5265, with a recall of 0.9946.
  • Interpretation: The models learned dataset-origin signatures rather than tumor features, which did not transfer to MAMA-MIA.
• Class-Wise Mixing: Introducing cross-dataset samples into both classes significantly improved external performance without changing the model architecture.
  • WaveViT-Small: Accuracy improved to 0.8463 and F1-score to 0.8625.
  • EfficientNet-B3: Accuracy improved to 0.8884 and F1-score to 0.8994.
  • Both models maintained high recall while significantly improving precision, demonstrating they no longer classified the majority of external slices as positive.

Significance and Claims
The paper claims that the composition of training data is as critical as the choice of neural network backbone for reliable cross-dataset generalization.

• Bias Control: The study demonstrates that "shortcut learning" driven by dataset origin can be effectively mitigated by ensuring that each class contains samples from multiple data sources.
• Methodological Requirement: The authors argue that strict external validation on independent cohorts (like MAMA-MIA) is necessary to reveal dataset-origin bias that internal splits miss.
• Scope: The work focuses on identifying whether classifiers rely on transferable signals rather than optimizing for a deployment-ready patient-level score. The authors note that current results are based on per-slice predictions and that future work could explore patient-level aggregation (e.g., averaging probabilities or majority voting) to further refine clinical utility.

The study concludes that controlling dataset-origin bias through class-wise mixing is a practical and essential methodological step for developing more reliable clinical breast MRI AI systems.