C^2GR: Coupled Comprehensive Generative Replay for a Continually Learnable Universal Segmentation Model

2026-06-22

AI Summary

This work addresses catastrophic forgetting in task-incremental learning for universal segmentation, where models suffer from dual distribution shifts in both image appearance and segmentation targets. To mitigate this, the authors propose a Coupled Compositional Generative Replay (C²GR) framework that leverages a Bayesian joint diffusion mechanism to simultaneously synthesize historically consistent image-mask pairs. Furthermore, they introduce a Relation-aware Unified Prompt Synchronization (RUPS) strategy to enable co-optimization of the generator and segmenter. The approach effectively alleviates forgetting while preserving data privacy. Evaluated across 20 cross-modal, multi-target medical segmentation tasks, the method achieves performance only 2.44% below that of joint training on all data, significantly outperforming existing continual learning approaches.

Abstract

Universal segmentation models exhibit significant potential for diverse tasks involving different imaging modalities and segmentation objectives. Task-Incremental Learning provides a privacy-preserving approach to continually evolve a universal model on tasks from sequentially-arriving medical departments. However, training the model solely on the incoming task induces forgetting on past tasks, since consecutive tasks exhibit concurrent shifts in image appearance and segmentation objective. To address this problem, we propose a novel Coupled Comprehensive Generative Replay (C^2GR) framework that simultaneously synthesizes image-mask pairs of previous tasks to mitigate forgetting under concurrent appearance and objective shifts. This requires preserving image-mask correspondence for structure-realistic generation and bridging asynchronous optimization of the generator and segmentor for segmentation-oriented generation. Specifically, we propose a Bayesian Joint Diffusion (BJD) method that formulates the correspondence as conditional distributions optimized via conditional denoising. Furthermore, we develop a Relation-aware Unified Prompt Synchronization (RUPS) scheme to simultaneously modulate the generator and segmentor via a shared task-relation-aware prompt for synchronizing their optimization. Experiments on 20 tasks spanning diverse modalities and objectives demonstrate that C^2GR exhibits only a 2.44% drop in overall performance compared to joint training with all task data, effectively alleviating forgetting from the concurrent shifts. Our code will be made publicly available at https://github.com/mar-cry/C2GR.

Problem

• continual learning
• catastrophic forgetting
• universal segmentation
• task-incremental learning
• medical image segmentation

Innovation

• Generative Replay
• Task-Incremental Learning
• Bayesian Joint Diffusion
• Prompt Synchronization
• Universal Segmentation