Distribution-Aware Replay for Continual MRI Segmentation

TL;DR

This work tackles distribution shifts in longitudinal MRI segmentation under privacy constraints that preclude data rehearsal and under silent model failures on unseen data. It introduces ccVAE, a two-stage architecture that attaches a conditioned VAE to a UNet to model the feature distribution p(u) through q(z|u) with prior p(z) = $N(0,I)$, enabling feature-space pseudo-replay via samples from p(z|t) and reconstruction-based OoD scoring. The VAE is conditioned on the task identity t and a slice index s, forming ccVAE, which permits robust, slice-aware distribution modeling and memory-efficient rehearsal without sharing patient data. Empirical results on hippocampus and prostate MRI datasets demonstrate improved Dice scores and calibration (low ECE) with strong OoD rejection, validating privacy-preserving continual learning for dynamic clinical environments.

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.

Figures (6)

• Figure 1: (1) The UNet is trained on the first task $\mathcal{T}_1$. (2) The VAE is trained on features $u_1$ with slice and task conditioning. (3) A set of features $\bar{u}_{i<t}$ are synthesized, pseudo-labeled and placed in memory $\mathcal{M}$. The UNet decoder is then trained on $\mathcal{M}$ and the new data of task $\mathcal{T}_t$. (4) During inference, the reconstruction loss between $u$ and $\hat{u}$ is used to classify whether the MRI is ID or OoD.
• Figure 2: Representative slices $s$ of MRI scans from each (a) hippocampus and (b) prostate dataset. The red areas depict the ground truth segmentation masks.
• Figure 3: Augmentations applied to the hippocampus (top row) and prostate (bottom row) datasets to create challenging OoD scenarios.
• Figure 4: Test Dice ($\uparrow$) during the learning trajectory for (a) hippocampus and (b) prostate. New tasks are introduced at 250 epoch intervals. ccVAE (yellow) maintains the most stable segmentation performance throughout the trajectory.
• Figure 5: Four segmentations produced by the model trained on the first prostate dataset. Images (a) and (b) are correctly considered ID and segmented correctly. (c) is correctly considered OoD, but (d) is misclassified.