Addressing catastrophic forgetting for medical domain expansion

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 re-training 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 forgetting 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.

Rethinking Architecture Design for Tackling Data Heterogeneity in Federated Learning

“…A schematic description of FedAVG [43] and CWT [7] is illustrated in Figure 1. At its core, FL presents a challenge of data heterogeneity in the distributions of training data across clients, which causes non-guaranteed convergence and model weight divergence for parallel FL methods [20,33,62], and severe catastrophic forgetting problem for serial FL methods [7,16,53]. Among recent developments to the classic parallel FedAVG algorithm [43] have included using server momentum (FedAVGM) to mitigate per-client distribution shift and imbalance [21], globally sharing small subsets of data among all users (FedAVG-Share) [62], using a proximal term to the local objective (FedProx) to reduce potential weight divergence under severely heterogeneous devices [33], and constructing a shared global model in a layer-wise manner by matching and averaging hidden elements (FedMA) [59].…”

“…This problem is particularly prevalent in FL over healthcare data since input images captured by different institutions may vary significantly in local patterns (difference in intensity, contrast, etc.) due to different medical imaging protocols [16,50], as well as in natural data splits across users due to user idiosyncrasies in speaking [30], typing [17], and writing [27]. On the other hand, ViTs have been shown to be significantly less biased towards local patterns as compared to CNNs and instead use self-attention to learn global interactions [48], which may contribute to their surprising robustness to distribution shifts and adversarial perturbations [3,44].…”

“…The batch normalization layer has been shown to be one of the major factors that deteriorate the performance of federated learning methods on non-IID data partitions [20,16]. Hsieh et al [20] demonstrate that group normalization (or layer normalization) can avoid the skew-induced accuracy loss of batch normalization on non-IID data.…”

“…Despite the rich opportunities afforded by FL, there remain fundamental research problems to be tackled before FL can be readily applicable to real-world data distributions. Most current methods that aim to learn a single global model across non-IID devices encounter challenges such as non-guaranteed convergence and model weight divergence for parallel FL methods [20,33,62], and severe catastrophic forgetting problem for serial FL methods [7,16,53].…”

Stabilizing Elastic Weight Consolidation method in practical ML tasks and using weight importances for neural network pruning

“…Subsequently, several papers have appeared that further refine [6] and develop [7], [8], [9], [10] the methodology proposed in [5]. There are also examples of applications of the Elastic Weight Consolidation method in applied tasks [11], [12], [13] and its comparative evaluations for different neural network architectures [12].…”