DoSReMC: Domain Shift Resilient Mammography Classification using Batch Normalization Adaptation

This paper introduces DoSReMC, a domain shift resilient framework that enhances cross-domain generalization in mammography classification by fine-tuning only batch normalization and fully connected layers alongside adversarial training, thereby addressing performance degradation caused by data distribution variations without requiring full model retraining.

The Gist

The Big Problem: The "Accent" Problem

Imagine you hire a brilliant detective (an AI) to find hidden clues in a specific city. This detective is trained for months on City A. They learn exactly how the streets look, how the buildings are painted, and the specific slang people use there. They become an expert at spotting criminals in City A.

Now, you send this detective to City B.

• In City A, the streetlights are yellow. In City B, they are white.
• In City A, people wear red coats. In City B, they wear blue.
• The "accent" of the data is different.

Even though the detective knows how to spot a criminal (the core skill hasn't changed), they get confused by the new lighting and colors. They start missing clues or flagging innocent people as suspects. In the medical world, this is called Domain Shift. An AI trained on mammograms from one hospital (or machine brand) often fails when shown images from a different hospital or machine.

The Culprit: The "Adjustable Glasses"

The researchers discovered that the main reason the AI gets confused isn't because it forgot how to see the cancer (the "eyes" or Convolutional Layers are still working fine). The problem is the Adjustable Glasses the AI wears, known in technical terms as Batch Normalization (BN) layers.

Think of these glasses as a filter that automatically adjusts the brightness and contrast of the image based on what the AI has seen before.

• When the AI was trained in City A, the glasses were calibrated to "City A brightness."
• When the AI looks at City B, those same glasses are still set to "City A brightness." The image looks washed out or too dark, and the AI can't see the cancer clearly.

The Solution: DoSReMC (The "Quick Glasses Swap")

The paper introduces a new method called DoSReMC. Instead of firing the detective and hiring a new one, or retraining the detective for months on City B, they simply swap the glasses.

Here is how they did it:

1. Keep the Brain: They left the detective's brain (the complex convolutional filters) exactly as it was. It already knows what cancer looks like.
2. Swap the Glasses: They only adjusted the "glasses" (the BN and Fully Connected layers) to match the new lighting of City B.
3. The Result: The detective can now see City B clearly without needing to relearn everything from scratch.

The Secret Weapon: The "Imposter" Game

To make the glasses even better, they added a trick called Adversarial Training.

Imagine the detective is playing a game against a "Domain Detective" (an imposter).

• The Domain Detective tries to guess: "Is this image from City A or City B?"
• The Main Detective tries to trick the Domain Detective by making the images look so similar that the Domain Detective can't tell the difference.

By playing this game, the Main Detective learns to ignore the "accent" (the specific machine brand or hospital) and focus purely on the "crime" (the cancer). This makes the AI robust enough to work in City A, City B, or even a brand new City C.

Why This Matters (The Real-World Impact)

1. Saves Time and Money: Usually, to fix an AI for a new hospital, you have to retrain the whole model, which takes huge amounts of computing power and time. DoSReMC is like changing the lenses on a camera instead of buying a whole new camera. It's fast and cheap.
2. Fairness: Currently, an AI might work great for patients in New York but fail for patients in a rural clinic with different equipment. This method helps ensure the AI works well for everyone, regardless of which machine took their X-ray.
3. Safety: By reducing "false positives" (thinking a healthy breast is cancerous) and "false negatives" (missing actual cancer), this method makes AI a safer tool for doctors to use in real life.

Summary

The paper says: "Don't retrain the whole brain; just adjust the glasses."

By realizing that the AI's "glasses" (Batch Normalization) are the main reason it fails when moving between different hospitals, the researchers created a lightweight, efficient system that swaps those glasses to fit the new environment. This allows AI to detect breast cancer accurately across different machines and hospitals without needing a massive, expensive overhaul.

Technical Summary

1. Problem Statement

Deep learning models for breast cancer detection via mammography often suffer from domain shift, where performance degrades significantly when applied to data from different sources (e.g., different hospitals, scanner manufacturers, or acquisition protocols).

• Root Cause: The paper identifies that Batch Normalization (BN) layers are a primary source of this domain dependence. BN layers rely on moving averages of mean and variance computed during training. When the input distribution of a target domain (e.g., a different X-ray machine) differs from the source domain, these fixed statistics cause a mismatch in feature distributions, leading to poor generalization.
• Current Limitations: Conventional solutions often involve retraining the entire network (computationally expensive) or using complex data augmentation. There is a lack of architectural analysis specifically isolating the role of BN layers in mammography domain shift.

2. Methodology: DoSReMC Framework

The authors propose DoSReMC (Domain Shift Resilient Mammography Classification), a framework designed to adapt models to new domains without retraining the entire network.

A. Core Strategy: Selective Fine-Tuning

Instead of fine-tuning all layers, DoSReMC freezes the convolutional layers (which learn robust, domain-invariant feature extractors) and only fine-tunes:

1. Batch Normalization (BN) layers: Specifically the learnable scaling ($\gamma$) and shifting ($\beta$) parameters, and the moving averages of statistics.
2. Fully Connected (FC) layers: The final classification heads.

• Rationale: This preserves the high-level feature representations learned from large pre-trained datasets while adapting the normalization statistics to the specific pixel intensity distributions of the target domain.

B. Integration with Domain-Adversarial Training (DAT)

To further enhance cross-domain generalization, the authors integrate a Partial Domain-Adversarial Training scheme:

• Architecture: A Domain Head is added to classify whether an image belongs to the source or target domain.
• Mechanism: A Gradient Reversal Layer (GRL) is applied only to the BN and FC layers. This forces the model to learn features that confuse the domain classifier (making features domain-invariant) while still accurately classifying benign vs. malignant cases.
• Efficiency: By restricting adversarial training to BN/FC layers, the method avoids disrupting the pre-trained convolutional filters.

C. Datasets

The study utilizes three large-scale Full-Field Digital Mammography (FFDM) datasets:

1. HCTP (Hacettepe-Mammo): A new, in-house dataset from Türkiye containing 157,463 pathologically confirmed images from GE scanners. It includes detailed BI-RADS scores and lesion masks.
2. VinDr-Mammo: A public dataset (5,000 studies) primarily from Siemens scanners.
3. CSAW: A public dataset (24,697 studies) from Hologic scanners.

3. Key Contributions

1. New Dataset (HCTP): Introduction of the largest mammography dataset in Türkiye, featuring pathologically confirmed diagnoses and diverse radiological findings (masses, calcifications, etc.).
2. Architectural Insight: The first comprehensive analysis demonstrating that BN layers are the primary bottleneck for cross-domain generalization in mammography. The study proves that BN statistics learned on one scanner type fail on others, causing feature collapse.
3. DoSReMC Framework: A novel, cost-efficient adaptation strategy that achieves performance comparable to full fine-tuning by updating only BN and FC layers.
4. Partial Adversarial Training: A demonstration that applying adversarial training only to normalization and classification layers yields better cross-domain stability than applying it to the entire network, while reducing computational costs by ~90% in backward passes.

4. Key Results

The experiments were conducted across various transfer scenarios (e.g., training on NYU/HCTP, testing on VinDr/CSAW).

• Impact of BN Statistics:
  • Models trained on NYU (Siemens/Hologic) performed poorly on HCTP (GE) and VinDr (Siemens) when using training-time BN statistics.
  • Test-Time BN: Recomputing BN statistics on the test batch significantly improved performance (e.g., +13% PR-AUC on VinDr), confirming that the convolutional features were valid but the normalization was mismatched.
• DoSReMC Performance:
  • Fine-tuning only BN and FC layers achieved results nearly identical to full fine-tuning (within ~1% PR-AUC) on the target domain.
  • Cross-Domain Generalization: The DoSReMC model trained on HCTP+VinDr achieved a PR-AUC of 0.82 on the unseen CSAW (Hologic) dataset, significantly outperforming models that were fully fine-tuned on specific domains (which suffered from catastrophic forgetting).
• Comparison with Alternatives:
  • Full Fine-Tuning: Often led to catastrophic forgetting when moving to distant domains.
  • Alternative Normalization (LayerNorm/GroupNorm): Replacing BN with LN or GN did not improve robustness; standard BN with adaptation remained superior.
  • Histogram Matching: Simple pixel-level preprocessing failed to improve performance, highlighting the need for feature-level adaptation.
• Computational Efficiency:
  • The partial adaptation strategy (freezing conv layers) reduced memory consumption by ~20% and backward pass time by ~10x compared to full-domain adversarial training.

5. Significance and Conclusion

• Clinical Relevance: DoSReMC provides a practical pathway for deploying AI in diverse clinical environments. It allows hospitals to adapt a pre-trained model to their specific scanner hardware (GE, Siemens, Hologic) with minimal computational resources and without needing massive labeled datasets for retraining.
• Theoretical Contribution: The paper shifts the focus from data-centric solutions to architecture-centric solutions, proving that the "domain gap" in medical imaging is largely a normalization issue rather than a feature extraction issue.
• Future Outlook: The authors suggest that in Federated Learning scenarios, only the fine-tuned BN parameters need to be shared between institutions, drastically reducing communication overhead and preserving patient privacy.

In summary, DoSReMC demonstrates that adapting Batch Normalization statistics is sufficient and necessary to overcome domain shift in mammography, offering a robust, efficient, and scalable solution for real-world AI deployment.