CoDG-Net: Structure-Guided Style Diffusion and Collaborative Learning to Mitigate Catastrophic Forgetting in Medical Image Domain Generalization

1 Introduction

Domain Generalization (DG) in medical imaging is a key technique for achieving robust medical image segmentation under multi-center and multi-device conditions [Li et al., 2020; Li et al., 2021a; Li et al., 2021b; Wang et al., 2022]. DG requires the model to rely solely on source-domain samples and maintain high performance on unseen target domains, without accessing any target-domain data during training [Su et al., 2023; Li et al., 2024; Cheng et al., 2023;

Robey et al., 2021; Liu et al., 2021; Robey et al., 2021]. In medical imaging, this capability is not only about crossdomain performance, but is also directly related to the accuracy and safety of clinical diagnosis [Liu et al., 2025]. However, in medical image DG segmentation, the problem of Catastrophic Forgetting (CF) has long been overlooked. Due to privacy and regulatory constraints, source-domain data cannot be freely transferred across institutions, making common continual learning strategies such as mixed-data training or replay of past data infeasible in medical scenarios. This makes CF even more challenging in medical image DG tasks.

To simultaneously improve model generalization and mitigate the CF problem, we conduct a statistical analysis of feature distributions on cross-domain data generated by augmentation methods and derive our motivations as follows: (1) Cross-domain data exhibit consistent anatomical structures, while their differences mainly lie at the style level, such as contrast, texture, and brightness (i.e., style inconsistency). For example, as shown in Fig.1(a), fundus images acquired from different devices share similar anatomical structures, whereas inter-domain discrepancies are primarily reflected in contrast and brightness; (2) Combining traditional augmentations with diffusion-based generation can better cover the cross-domain data distribution and thus enhance domain robustness. As illustrated in Fig.1(b), applying traditional augmentation together with diffusion-based generation on the source domain yields extended domains with richer style information; (3) Unlike training that only fits the source-domain distribution, DG imposes additional generalization constraints (e.g., style diversity) on the parameters within a single training phase to improve cross-domain robustness. This drives the model away from the "sourcedomain optimum" toward a compromise that is more robust but no longer optimal on the source domain, thereby reducing its distribution-fitting capability on the source domain, as shown in Fig.1(c).

Motivated by the above, we aim to simultaneously improve the model's generalization ability on unseen domains and effectively alleviate catastrophic forgetting. First, we propose a novel structure-guided style diffusion augmentation method. It uses the Fourier transform to extract frequencydomain signals from pseudo-domain images generated by traditional style augmentation methods. We then perform diffusion generation on the amplitude spectrum (style), and recombine it with the phase spectrum to obtain new samples with a broader coverage of feature distributions. To effectively mitigate catastrophic forgetting, we design a collaborative learning network with a dual-branch interactive structure (CoDG-Net). The network contains two complementary branches: a source-domain image restoration branch and a generalized segmentation branch. To promote sufficient interactive optimization between the feature spaces of these two branches and to balance source-domain knowledge preservation with target-domain generalization ability, we propose a novel learning bias-guided strategy. It adaptively regulates knowledge transfer at both the layer level and the task level, improving segmentation accuracy while effectively alleviating the CF problem. The main contributions of this paper are summarized as follows:

• We propose a structure-guided style diffusion augmentation method. Under the constraint of preserving anatomical structure consistency, this method introduces a diffusion model to better fit the style distributions of unknown domains, while performing diffusion only on the amplitude spectrum to significantly reduce the difficulty of generation.
• We investigate the CF problem in medical image domain generalization. We propose a collaborative learning network with a dual-branch interaction structure and design a novel learning bias-guided strategy for it, which guides knowledge transfer in the network and effectively alleviates this problem.
• Experimental results show that our method outperforms existing SOTA methods on both single-source and multisource domain generalization datasets, and achieves a lower forgetting rate on source-domain data.

2 RelatedWork

2.1 Domain Generalization inMedical Imaging

Domain generalization (DG) aims to train a model on one or multiple source domains so that it can directly generalize to unseen target domains, which is crucial for robust medical image analysis under varying conditions [Yoon et al., 2024]. To enhance the cross-domain generalization ability of medical image segmentation models, some medical DG methods adopt data augmentation techniques-such as geometric/color perturbations, frequency-domain perturbations, and style transfer-to enrich style variations of sourcedomain samples and simulate potential target-domain distributions [Liu et al., 2024; Huang et al., 2025; Dou et al., 2019]. Other methods exploit domain-invariant feature extraction, cross-domain alignment, adversarial learning, or normalization strategies to reduce inter-domain discrepancies [Choi et al., 2023; Pei et al., 2024; Zhou et al., 2022b; Seo et al., 2020].

2.2 Catastrophic Forgetting in Continual Learning

Catastrophic Forgetting (CF) is a core problem in Continual Learning (CL), referring to the phenomenon that a model significantly forgets previously learned knowledge when learning new tasks or new sites [Wang et al., 2024]. Common mitigation strategies include parameter regularization, knowledge distillation, and data replay [Nori et al., 2025; Zhao et al., 2024; Ma et al., 2022]. In CL-based solutions for DG, Xu et al. [Xu et al., 2025] address cross-site continual segmentation with a joint pursuit of memorization and generalization, and alleviate CF via site-modulated replay and gradient alignment. However, their method assumes that multisite data arrive sequentially and explicitly relies on replayed samples, essentially following a replay-based strategy; moreover, historical site data remain accessible during training.

In contrast, the medical image DG setting considered in this paper has no access to target-domain data during training, thus lacking the cross-domain alignment conditions that most Domain Adaptation (DA) methods rely on. The model is trained only once on a fixed source domain, and after training it can neither access the original source data nor perform any form of task-level data replay. Therefore, existing replaybased CL methods are difficult to directly transfer to the DG setting. To address these limitations, we propose CoDG-Net, which tackles the CF problem in DG under the constraints of no replay and no target-domain data.

3 Method

3.1 ProblemDefinition and Overview

In domain generalization tasks, there exist a source domain DS and a target domain DT . The source domain is defined as DS = {xSi , yS i }NS i=1 , where xSi is the i-th sample in the source domain, yS i is the corresponding segmentation mask, and NS is the total number of samples. The goal of this paper is to train a segmentation model Sseg using only samples from DS , and to ensure its generalization ability to unseen target domains.

As shown in Fig. 2, to address the problem of limited style diversity in source-domain data, we first adopt a traditional B´ezier-curve-based intensity mapping strategy [Zhou et al., 2022b] to explore pseudo-domain styles and obtain a pseudodomain DP . We then introduce a structure-guided style diffusion augmentation module. This module exploits the anatomical consistency of medical images and performs style expansion on source-domain data in the Fourier domain, generating an extended domain DG = {xGi , yG i }NG i=1 with diverse style distributions. Here, each xGi preserves the structural semantics of xSi while simulating the style characteristics of potential target domains. On this basis, in order to adapt to the distribution of DG while suppressing catastrophic forgetting, we propose a collaborative learning network with a dualbranch interactive structure (CoDG-Net). One branch trains a segmentation model Sseg : xG → yG, and the other branch trains a reconstruction model Srec : xG → xS . Through the collaborative training of these two tasks, CoDG-Net establishes a strong interaction mechanism in the feature space, effectively balancing cross-domain generalization and sourcedomain memory preservation.

3.2 Structure-Guided Style Diffusion Augmentation

In previous studies, a large body of work has demonstrated that image style and structure are encoded in the amplitude spectrum and phase spectrum of the Fourier domain, respectively [Qiao et al., 2024; Lv et al., 2022; Piotrowski and Campbell, 1982; Oppenheim and Lim, 2005]. Therefore, to further enhance sample diversity and improve the model's coverage of potential target domains, we propose a novel structure-guided style diffusion augmentation method. The core idea of this method is to exploit the explicit constraint of image structure provided by the Fourier phase spectrum, and to combine it with the strong generative capability of diffusion models to explore unknown amplitude (style) distributions. As shown in Fig. 3, this process not only generates style-diverse DG, but also strictly preserves anatomical structural consistency, thereby avoiding structural distortions in the generated images.

Forward Process: First, we apply the fast Fourier transform (FFT) to convert the input pseudo-domain image into the frequency domain and decouple it into an amplitude spectrum A(u, v) and a phase spectrum Φ(u, v). We take the amplitude spectrum as the diffusion target, while the phase spectrum is preserved for subsequent structural guidance. Specifically, in the forward noising process, we degrade the amplitude spectrum by adding Gaussian noise, as follows:

Where x0 denotes the original amplitude spectrum data, and xt is the noisy data at time step t. The term ¯αt = ∏ts=1 αs is the cumulative attenuation coefficient, where αt = 1-βt and βt represents the noise strength at each time step. N(·; u, σ2 ) denotes a normal distribution with mean u and variance σ2 .

Reverse Process: The goal of the reverse denoising process is to restore the original data x0 from the noisy data xt, by introducing Φ(u, v) as a structural prior condition into the denoising network. The process is defined as:

Where Σθ(xt, t) is the variance term, representing the uncertainty of generation; θ denotes the learnable parameters; and uθ(xt, t,Φ) is the mean prediction function conditioned on the constraint information Φ. It represents the prediction of xt-1 from xt at step t, and is given by:

Where ϵθ(xt, t, Φ) is the noise component predicted by the model. The model learns to approximate the noise ϵ added in the forward process, based on xt, the time step t, and the constraint condition Φ.

Sampling process: We propose a continuous compositional sampling mechanism in the frequency domain, which aims to generate samples that are both style-diverse and structurally reliable. Specifically, we first apply the Fourier transform to extract the phase spectrum Φ(·) and the amplitude spectrumA(·) of the original image. Then, we sample an initial noise xT ∼ N(0, I) from a standard normal distribution as a latent style code. Under the guidance of the original Φ(·), we perform the reverse denoising process to generate a new amplitude spectrum A′ (u, v). Finally, we recombine the generated A′ (·) with the original Φ(·), and reconstruct the final extended-domain image f(x, y) via the inverse Fourier transform (iFFT):

Since this method performs diffusion-based generation only on the amplitude spectrum, it substantially reduces the difficulty of the diffusion process. At the same time, the phase spectrum serves as guidance to preserve critical lesion and organ boundaries with pixel-level accuracy.

3.3 Collaborative Learning Network with Dual-Branch Interaction

To explicitly model the cross-domain style mapping patterns, we design CoDG-Net with two parallel branches, Srec and Sseg (as shown in Fig. 3). The network captures style mapping information indirectly through collaborative learning. Specifically, Srec is dedicated to learning the inverse mapping "from the extended domain back to the source domain": when the network attempts to restore multi-style extended samples to the source-domain appearance, the intermediate layers inevitably encode rich style transfer relationships. These inter-domain mapping cues are then passed to Sseg via feature interaction between the two branches, enabling Sseg to truly learn cross-domain appearance variations during training, rather than merely adding generalization regularization around the source-domain optimum.

Moreover, not all intermediate-layer features are mutually beneficial for the two branches. If cross-branch interactions are indiscriminately enabled at every layer, or if a unified soft-gating/attention fusion is applied across all layers, task-irrelevant noise can be propagated and the reconstruction branch may exert excessive interference on the main segmentation branch, thereby degrading segmentation performance. Therefore, we propose a Learning Bias-Guidance Strategy. The layer-wise strategy selectively opens feature interactions at key layers to suppress ineffective information propagation, while the task-wise strategy explicitly biases the global optimization toward the segmentation task, thus enhancing crossdomain robustness and effectively mitigating, or even reversing, catastrophic forgetting on the source domain.

Layer-Wise Bias-Guided Strategy

In multi-task collaborative learning, feature layers at different network depths usually carry different semantic or style information, and not all layers need cross-task interaction at all times. To adaptively control the interaction behavior of each layer, we propose a layer-wise bias guidance strategy. This strategy aims to decide whether the current layer should focus on "self feature enhancement" or "cross-task feature fusion". Specifically, we define an Interaction Flag to replace traditional hard-coded connections. For the l-th layer of the network, between training steps t and t + k, we compute the evolution trend of the feature map F (t,l) . We use the Euclidean distance d(·, ·) to measure the consistency change of features. If the feature F (t,l) of one branch at step t is closer to the feature F (t+k,l) other of the other branch after k update steps, it indicates that this layer tends to absorb knowledge from the other task, and we mark it as a Collaborative State. Conversely, if the feature is more inclined to maintain its own update direction, we mark it as an Independent State. To prevent the model from falling into a local optimum or staying in a single state for a long time, we introduce a stochastic perturbation mechanism, which forcibly flips the current state with probability δ. Finally, according to the computed state, we generate a layer-wise interaction factor f (l) layer ∈ {0, 1} to dynamically turn on or off the feature interaction channel at that layer (see Algorithm 1 for the detailed procedure).

Algorithm 1: Layer-wise Bias-Guided Strategy.

Input: training step t; intermediate feature maps

F (t,l) seg and F (t,l) rec ; feature maps after k update

• 3 ▷ True → collaborate; False → independent
• 4 With probability δ: Flag ← ¬Flag
• 5 Compute state: State ← Flag seg ∨ Flag rec
• 6 if State then

• 8 else

• 10 Return f (l) layer

Task-Wise Bias-Guided Strategy

Although dual-branch collaborative learning can promote feature extraction, our ultimate goal is to improve the generalization performance of segmentation. If the two tasks are treated as completely equal, the reconstruction task may overly dominate the gradient update direction, which can damage segmentation performance. To address this issue, we propose a task-bias guidance strategy as a global regulation mechanism. Specifically, we first compute the Euclidean distances of features before and after model updates along different directions, and use the relative magnitudes of these distances to evaluate feature effectiveness. We define the smaller Euclidean distance as the effective optimization distance of the current model, because it represents the direction toward which the features tend to converge. To clarify the primary and secondary roles of the tasks, we further assign a sign (direction) to this distance according to the optimization objective: if the current optimization trend of the model moves toward the reconstruction model (i.e., is biased to the reconstruction task), we assign a positive value to the effective distance; conversely, if the model tends to deviate from the reconstruction model (i.e., focuses on itself or the segmentation task), we assign a negative value. Based on this logic, we obtain the task interaction factor f (l) task for the l-th layer (see Algorithm 2 for the detailed computation process).

Algorithm 2: Task-wise Bias-Guided Strategy.

Input: training step t; feature maps at intermediate layer l, F (t,l) seg and F (t,l) rec ; feature maps after k update steps, F (t+k,l) seg and F (t+k,l) rec

• 1 Compute dirseg:

• 2 Compute dirrec:

• 4 Return f (l) task

3.4 Loss Function

To enhance the dual-branch collaborative mechanism of CoDG-Net and explicitly impose the bias-guided strategy, we propose a novel loss function LCO. This loss computes the discrepancy between intermediate-layer feature maps and adjusts it using the layer-wise guidance factor f (l) layer and the task-wise guidance factor f (l)

task , as follows:

Where N is the number of intermediate layers, M is the number of pixels in each feature map, f (l) layer and f (l) task are the adjustment factors for the l-th layer, and F (t+k,l) seg,i andF (t+k,l) rec,i denote the i-th feature value at the l-th layer of Sseg and Srec at stage t +k, respectively.

For the loss of Sseg, we combine the Dice loss with the collaborative loss LCO to enhance gradient stability and pixellevel segmentation capability:

Where α and β are the fusion weights of the two losses. Similarly, the total loss of Srec combines the commonly used MSE loss LMSE for reconstruction with the collaborative loss LCO:

Where α′ and β′ are the fusion weights of the two losses.

4 Experiments

4.1 Experimental Setup

Datasets. We conduct experiments on two medical image domain generalization datasets, covering both single-source and multi-source settings, to verify the effectiveness and generalization ability of our method: the BraTS 2018 dataset [Menze et al., 2014] and a fundus image dataset [Orlando et al., 2020]. The BraTS 2018 dataset contains 210 high-grade glioma cases and 75 low-grade glioma cases. Each case includes four MRI modalities: T2, FLAIR, T1, and T1CE. We use this dataset to perform single-source domain generalization experiments. The fundus image dataset consists of images from four domains collected by different institutions and different scanning devices. The main task is to segment the optic cup and optic disc in fundus images, and this dataset is used for multi-source domain generalization experiments.

Implementation details. All experiments were conducted on Ubuntu 22.04 with Python 3.9 and PyTorch 2.4.1, using 4 NVIDIA RTX 4090 GPUs for training. For a fair comparison, we adopt the same data preprocessing steps as the competing methods. For the BraTS 2018 dataset, we first slice the 3D volumes along the scanning direction, then resize each slice to 240 × 240, and normalize the pixel intensities to the range [-1, 1]. The slices are then split into training, validation, and test sets with a ratio of 7:1:2. For the fundus image dataset, we resize the images from 800×800 to 384×384, and follow the official train-validation split provided by the dataset. For training the diffusion model, we set the number of epochs to 100 and the batch size to 8. We use the Adam optimizer with an initial learning rate of 0.0001, the maximum variance is set to 0.02, and a linear noise schedule is adopted. For training the CoDG-Net collaborative learning network, both the segmentation branch and the reconstruction branch use U-Net as the backbone. The coefficient of LCO is decayed by a factor of 0.5 every 5 epochs. The random perturbation probability δ is set to 0.05, and the hyperparameters α, β, α′ , β′ are all set to 0.5.

Evaluation metrics. To evaluate the segmentation performance of the model on target domains, we use the Dice similarity coefficient, Average Surface Distance (ASD), and Hausdorff Distance (HD) for quantitative assessment. To quantitatively measure the severity of catastrophic forgetting, we adopt the concept of the forgetting measure commonly used in Continual Learning (CL) [Hayes et al., 2018]. Different from the standard CL setting where tasks arrive sequentially, we define the Source-domain Forgetting Measure (SFM) as the performance gap on the source domain between a baseline model trained only on the source domain (Empirical Risk Minimization, ERM) and the proposed domain generalization model. A lower SFM indicates better preservation of source-domain knowledge. An ideal domain generalization model should significantly improve target-domain performance while keeping SFM close to zero (or even negative, which implies positive transfer).

4.2 Comparative Experiments

To comprehensively evaluate the performance of our method, we compare it with existing state-of-the-art (SOTA) approaches on both single-source and multi-source domain generalization tasks, in order to highlight its effectiveness. The experimental results on the BraTS dataset are reported in Table 1, and the results on the Fundus dataset are presented in Table 2.

The results show that our method achieves excellent generalization performance on both single-source and multi-source domain generalization tasks. This is mainly attributed to the proposed structure-guided style augmentation strategy, which can significantly enlarge the coverage of generated samples over unseen domains, enabling the model to better adapt to different domain distributions.

In addition, to further verify the capability of CoDG-Net in alleviating catastrophic forgetting during domain generalization, we compare the SFM values of different methods, as reported in Table 3. The results show that, compared with the baseline model trained only on source-domain data, existing domain generalization methods generally suffer from catastrophic forgetting on the source domain. In contrast, the proposed CoDG-Net can effectively mitigate this problem. We also provide a visual comparison of the predictions on the source and target domains in Fig. 4 (results on the Fundus dataset are included in the supplementary material).

4.3 Ablation Study and Analysis

We conduct ablation experiments to evaluate the contribution of the key components in our method, including pseudodomain style exploration (SE), style diffusion augmentation (SD), dual-branch interaction structure (DBS), layer-wise bias-guided strategy (LW), and task-wise bias-guided strategy (TW). The ablation results are reported in Table 4, where the full model achieves the best performance. Compared with the baseline model (VA1) without domain generalization, introducing the proposed structure-guided style augmentation (VA2, VA3, VA4) significantly improves the performance on the target domain, with Dice increasing from 29.69 to 67.59. However, this also leads to severe forgetting on the source domain, where SFM reaches 5.19. This further validates the rationality of our overall design. It is worth noting that when the SD module is applied only on source-domain data, the degree of source-domain forgetting is relatively small, but the performance gain on the target domain is also limited. This indicates that, without SE, the model can only perform diffusion generation around the source-domain distribution and cannot effectively produce samples from other domains, which demonstrates the necessity of the SE module.

When the dual-branch interaction structure is introduced for collaborative learning, simply enabling interactions at all intermediate layers without any bias guidance (VA5) provides only limited protection of source-domain performance, and even leads to a performance drop on the target domain compared with VA4. This indicates that not all intermediatelayer features possess effective cross-task complementarity; blindly enhancing layer-wise interactions may introduce noisy interference and thus weaken the overall performance of the model. The two bias-guided strategies proposed in CoDG-Net play a crucial role in both preserving sourcedomain performance and improving the generalization ability of the model. Compared with VA5, VA6 better protects the source-domain performance and avoids introducing additional noise. Compared with VA5, VA7 encourages the model to focus more on the segmentation task, thereby improving its segmentation performance. However, using only one of the two strategies cannot achieve globally optimal performance; only when both strategies are applied simultaneously can the model obtain the best trade-off between source-domain protection and target-domain generalization.

In addition, we analyze the activation frequency of collaborative interactions at each intermediate layer during training, as shown in Fig. 5. The vertical axis indicates the proportion of training iterations in which the corresponding layer is in the "open" state. We observe that shallow layers generally have lower activation frequencies, while deeper layers tend to be activated more often, suggesting that not all intermediate-layer interactions are beneficial for segmentation performance. If interactions are blindly and frequently enabled at these layers, it is easy to introduce noisy interference, leading to performance degradation as in VA5. This further demonstrates the necessity of the proposed layer-wise bias-guided strategy.

We also conduct a hyper-parameter sensitivity study on the random perturbation probability δ, as shown in Fig. 6. Using T2 as the source domain on the BraTS dataset, the results show that the T2 source domain is not highly sensitive to δ. However, when δ = 0.05, the standard deviation across different random seeds is smaller, indicating more stable model performance. Moreover, on the three target domains (Flair, T1, and T1CE), the model achieves the best performance when δ is set to 0.05.

5 Conclusion

In this paper, we address the often-overlooked problem of catastrophic forgetting in domain generalization for medical image segmentation and propose a collaborative learning network with a dual-branch interaction structure, termed CoDG-Net. We first introduce a structure-guided style diffusion augmentation method, which provides pseudo-domain samples with more diverse styles and broader coverage for domain generalization. Then, we design a novel bias-guided learning strategy to dynamically regulate knowledge transfer across different layers and between tasks in CoDG-Net, thereby balancing source-domain knowledge retention and segmentation accuracy improvement. Experimental results demonstrate that our method outperforms existing SOTA approaches on both single-source and multi-source domain generalization benchmarks.