[Paper Review] Addressing catastrophic forgetting for medical domain expansion

Sharut Gupta, Praveer Singh+14 more|arXiv (Cornell University)|Mar 24, 2021

Machine Learning in Healthcare4 citations

TL;DR

This paper proposes a novel approach to mitigate catastrophic forgetting in medical domain expansion by combining Elastic Weight Consolidation (EWC) with adaptive modulation of batch normalization statistics. The method preserves performance on original datasets while fine-tuning on new medical imaging data from different institutions or systems, outperforming state-of-the-art methods with statistically significant gains in model robustness and generalization.

ABSTRACT

Model brittleness is a key concern when deploying deep learning models in real-world medical settings. A model that has high performance at one institution may suffer a significant decline in performance when tested at other institutions. While pooling datasets from multiple institutions and retraining may provide a straightforward solution, it is often infeasible and may compromise patient privacy. An alternative approach is to fine-tune the model on subsequent institutions after training on the original institution. Notably, this approach degrades model performance at the original institution, a phenomenon known as catastrophic forgetting. In this paper, we develop an approach to address catastrophic forget-ting based on elastic weight consolidation combined with modulation of batch normalization statistics under two scenarios: first, for expanding the domain from one imaging system's data to another imaging system's, and second, for expanding the domain from a large multi-institutional dataset to another single institution dataset. We show that our approach outperforms several other state-of-the-art approaches and provide theoretical justification for the efficacy of batch normalization modulation. The results of this study are generally applicable to the deployment of any clinical deep learning model which requires domain expansion.

Motivation & Objective

• Address the challenge of catastrophic forgetting when deploying deep learning models in new medical institutions or imaging systems.
• Develop a method that maintains high performance on original datasets while adapting to new domain data without retraining from scratch.
• Overcome privacy and data-sharing limitations by enabling incremental domain expansion via fine-tuning instead of centralized retraining.
• Provide theoretical justification for the effectiveness of batch normalization modulation in mitigating distribution shift across domains.
• Demonstrate the generalizability of the approach to real-world clinical deployment of deep learning models in multi-institutional settings.

Proposed method

• Apply Elastic Weight Consolidation (EWC) to constrain updates to critical parameters learned from the original dataset, preserving knowledge through parameter importance regularization.
• Introduce adaptive modulation of batch normalization (BN) statistics—adjusting running mean and variance during fine-tuning to align internal activations with the new domain.
• Combine EWC with BN modulation to jointly stabilize performance on both source and target domains during incremental fine-tuning.
• Use patient-level averaging of image predictions across views to generate final patient-level breast density assessments.
• Train a ResNet-50 model with ImageNet pretraining, using cross-entropy loss and Adam optimizer with early stopping and model checkpointing.
• Apply data augmentation via random flips and rotations (up to 45°) during training to improve robustness and generalization.

Experimental results

Research questions

• RQ1 Can EWC combined with BN modulation effectively reduce catastrophic forgetting when fine-tuning a medical DL model on a new institution’s data?
• RQ2 How does the proposed method compare to state-of-the-art approaches in maintaining performance on both source and target domains?
• RQ3 To what extent does BN modulation help align internal feature distributions across different imaging systems and institutions?
• RQ4 Does the method preserve model performance on the original dataset while achieving improved accuracy on the new domain?
• RQ5 Is the approach scalable and effective in real-world clinical settings with privacy constraints and heterogeneous imaging protocols?

Key findings

• The proposed method significantly outperformed baseline fine-tuning and other state-of-the-art approaches in maintaining performance on the original DMIST dataset after fine-tuning on MGH data.
• The combination of EWC and BN modulation achieved a Cohen’s Kappa score of 0.72 on the original DMIST dataset and 0.75 on the MGH dataset, demonstrating strong generalization across domains.
• The method showed substantial improvement over EWC alone and BN modulation alone, indicating synergistic benefits from combining both techniques.
• Theoretical analysis confirmed that BN modulation effectively reduces internal covariate shift between domains, supporting its role in mitigating forgetting.
• Statistical analysis using the Wilcoxon signed-rank test confirmed that performance gains were significant at p < 0.05 across multiple evaluation metrics.
• The approach maintained high model robustness even when fine-tuning on data from different imaging systems (e.g., Hologic vs. GE) and institutions, validating its generalizability.