Distribution-Aware Replay for Continual MRI Segmentation

Nick Lemke,Camila González,Anirban Mukhopadhyay,Martin Mundt

2024-07-31

Abstract:Medical image distributions shift constantly due to changes in patient population and discrepancies in image acquisition. These distribution changes result in performance deterioration; deterioration that continual learning aims to alleviate. However, only adaptation with data rehearsal strategies yields practically desirable performance for medical image segmentation. Such rehearsal violates patient privacy and, as most continual learning approaches, overlooks unexpected changes from out-of-distribution instances. To transcend both of these challenges, we introduce a distribution-aware replay strategy that mitigates forgetting through auto-encoding of features, while simultaneously leveraging the learned distribution of features to detect model failure. We provide empirical corroboration on hippocampus and prostate MRI segmentation.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper primarily addresses the challenges of continual learning in the field of medical image segmentation by proposing a method called "Distribution-Aware Replay." Specifically, this method aims to solve the following two core issues: 1. **Avoiding Catastrophic Forgetting**: Medical image datasets are usually small and change over time, which makes models prone to forgetting previously learned knowledge (catastrophic forgetting). Traditional data replay strategies can effectively mitigate this issue but directly violate medical data privacy regulations. 2. **Identifying and Handling Outlier Samples**: The model needs to cope not only with expected changes in data distribution but also with the ability to identify and handle samples that unexpectedly come from unknown distributions (i.e., out-of-distribution, OoD), which may lead to erroneous and overconfident predictions. To address these issues, the authors propose a two-stage modeling architecture that includes a Conditional Variational Autoencoder (cVAE) to capture the feature distribution of the base network (e.g., UNet) and achieve pseudo-rehearsal to alleviate forgetting. At the same time, this method can utilize the learned feature distribution to detect and reject OoD samples, avoiding the model's "silent failures." To validate the effectiveness of the proposed method, the authors conducted experimental evaluations on two medical image segmentation tasks: hippocampus and prostate MRI image segmentation, and also assessed its OoD detection capability in handling augmented datasets. The experimental results show that the method can maintain stable segmentation performance across different tasks while effectively identifying OoD samples. Additionally, a series of ablation experiments demonstrated the importance of each component of the proposed method.

Distribution-Aware Replay for Continual MRI Segmentation

Continual Learning of Medical Image Classification Based on Feature Replay

Segmentation of Multiple Sclerosis Lesions across Hospitals: Learn Continually or Train from Scratch?

Dynamic memory to alleviate catastrophic forgetting in continual learning with medical imaging

Continual atlas-based segmentation of prostate MRI

Comprehensive Generative Replay for Task-Incremental Segmentation with Concurrent Appearance and Semantic Forgetting

What is Wrong with Continual Learning in Medical Image Segmentation?

Incremental Learning for Heterogeneous Structure Segmentation in Brain Tumor MRI

Continual Learning via Manifold Expansion Replay

Towards Synchronous Memorizability and Generalizability with Site-Modulated Diffusion Replay for Cross-Site Continual Segmentation

Rethinking Exemplars for Continual Semantic Segmentation in Endoscopy Scenes: Entropy-based Mini-Batch Pseudo-Replay

Handling Data Heterogeneity with Generative Replay in Collaborative Learning for Medical Imaging

RECALL+: Adversarial Web-based Replay for Continual Learning in Semantic Segmentation

Distributionally Robust Segmentation of Abnormal Fetal Brain 3D MRI

Learning Towards Synchronous Network Memorizability and Generalizability for Continual Segmentation Across Multiple Sites

Continual Domain Incremental Learning for Privacy-aware Digital Pathology

Decoupling Predictions in Distributed Learning for Multi-Center Left Atrial MRI Segmentation

Distribution-Level Memory Recall for Continual Learning: Preserving Knowledge and Avoiding Confusion

M EMORY R EPLAY WITH D ATA C OMPRESSION FOR C ONTINUAL L EARNING

Selective experience replay compression using coresets for lifelong deep reinforcement learning in medical imaging