Title: Towards Generic Semi-Supervised Framework for Volumetric Medical Image Segmentation
Venue: NeurIPS 2023
Accepted: yes
Keywords: Volumetric Medical Image Segmentation

Semi-supervised Learning

Unsupervised Domain Adaptation

Semi-supervised Domain Generalization
Abstract: Volume-wise labeling in 3D medical images is a time-consuming task that requires expertise. As a result, there is growing interest in using semi-supervised learning (SSL) techniques to train models with limited labeled data. However, the challenges and practical applications extend beyond SSL to settings such as unsupervised domain adaptation (UDA) and semi-supervised domain generalization (SemiDG). This work aims to develop a generic SSL framework that can handle all three settings. We identify two main obstacles to achieving this goal in the existing SSL framework: 1) the weakness of capturing distribution-invariant features; and 2) the tendency for unlabeled data to be overwhelmed by labeled data, leading to over-fitting to the labeled data during training. To address these issues, we propose an Aggregating & Decoupling framework. The aggregating part consists of a Diffusion encoder that constructs a "common knowledge set" by extracting distribution-invariant features from aggregated information from multiple distributions/domains. The decoupling part consists of three decoders that decouple the training process with labeled and unlabeled data, thus avoiding over-fitting to labeled data, specific domains and classes. We evaluate our proposed framework on four benchmark datasets for SSL, Class-imbalanced SSL, UDA and SemiDG. The results showcase notable improvements compared to state-of-the-art methods across all four settings, indicating the potential of our framework to tackle more challenging SSL scenarios. Code and models are available at: https://github.com/xmed-lab/GenericSSL.

Towards Generic Semi-Supervised Framework for Volumetric Medical Image Segmentation

Abstract

Volume-wise labeling in 3D medical images is a time-consuming task that requires expertise. As a result, there is growing interest in using semi-supervised learning (SSL) techniques to train models with limited labeled data. However, the challenges and practical applications extend beyond SSL to settings such as unsupervised domain adaptation (UDA) and semi-supervised domain generalization (SemiDG). This work aims to develop a generic SSL framework that can handle all three settings. We identify two main obstacles to achieving this goal in the existing SSL framework: 1) the weakness of capturing distribution-invariant features; and 2) the tendency for unlabeled data to be overwhelmed by labeled data, leading to over-fitting to the labeled data during training. To address these issues, we propose an Aggregating & Decoupling framework. The aggregating part consists of a Diffusion encoder that constructs a common knowledge set by extracting distribution-invariant features from aggregated information from multiple distributions/domains. The decoupling part consists of three decoders that decouple the training process with labeled and unlabeled data, thus avoiding over-fitting to labeled data, specific domains and classes. We evaluate our proposed framework on four benchmark datasets for SSL, Class-imbalanced SSL, UDA and SemiDG. The results showcase notable improvements compared to state-ofthe-art methods across all four settings, indicating the potential of our framework to tackle more challenging SSL scenarios. Code and models are available at: https://github.com/xmed-lab/GenericSSL.

1 Introduction

Labeling volumetric medical images requires expertise and is a time-consuming process. Therefore, the use of semi-supervised learning (SSL) is highly desirable for training models with limited labeled data. Various SSL techniques [1, 2, 3, 4, 5, 6, 7] have been proposed, particularly in the field of semi-supervised volumetric medical image segmentation (SSVMIS), to leverage both labeled and unlabeled data. However, current SSVMIS methods [8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18] assume that the labeled and unlabeled data are from the same domain, implying they share the same distribution. In practice, medical images are often collected from different clinical centers using various scanners, resulting in significant domain shifts. These shifts arise due to differences in patient populations, scanners, and scan acquisition settings. As a consequence, these SSVMIS methods have limitations in real-world application scenarios and frequently encounter overfitting issues, leading to suboptimal results.

To address this limitation, researchers have increasingly focused on Unsupervised Domain Adaptation (UDA) techniques. These techniques leverage both labeled (source domain) and unlabeled data (target

Real-world

Real-world

Real-world

(a) SSL

(b) UDA

(c) SemiDG

Unlabeled data from different domains

Labeled data from different domains

Domain shifts

Sample

more

data

Sample

more

data

Training set

Testing set

Data from any

unseen domains

domain) for training, but the data originate from different domains. Furthermore, Semi-supervised Domain Generalization (SemiDG), a more stringent scenario, has garnered significant interest. SemiDG utilizes labeled and unlabeled data from multiple domains during training and is evaluated on an unseen domain. Currently, methods for these three scenarios are optimized separately, and there is no existing approach that addresses all three scenarios within a unified framework. However, given that all training stages involve labeled and unlabeled data, it is intuitive to explore a generic SSL-based framework that can handle all settings and eliminate the need for complex task-specific designs. Therefore, this paper aims to develop a generic framework that can handle existing challenges in real-world scenarios, including:

Scenario 1: SSL (Figure 1(a)): The sample data used for both training and testing are from the same domain, representing the standard SSL setting.

Scenario 2: UDA (Figure 1(b)): The sampled data originate from two domains, with the labels of the target domain being inaccessible, representing the UDA setting.

Scenario 3: SemiDG (Figure 1(c)): The sampled data encompasses multiple domains, with only a limited number of them being labeled, representing the SemiDG setting.

Potential similarities can be found and summarized as follows: (1) in the training stage, both labeled data and unlabeled data are used; (2) in the scenario of the real-world application domain, whether the distribution shifts in SSL or the domain shifts in UDA and SemiDG can all be regarded as sampling bias , i.e. , the main difference is how we sample the data in Figure 1.

Now we wonder whether the existing SSVMIS methods are powerful enough to handle this general task. Experimental results show that the existing SSL methods do not work well on UDA and SemiDG settings, as shown in Table 3 & 4, and vice versa (Table 2). One of the main obstacles lies in the severe over-fitting of these models, which is caused by the dominance of labeled data during training . Specifically, the state-of-the-art SSVMIS methods are mainly based on two frameworks: (1) Teacher-student framework [1], where a student model is first trained with the labeled data, and a teacher model obtained from the EMA of the student model generates the pseudo label to re-train the student model with labeled data, see Figure 2(a); (2) CPS (Cross Pseudo Supervision) [3] framework, which leverages the consistency between two perturbed models and the pseudo label generated by one of the networks will be used to train the other network, see Figure 2(b). The predicting modules in these two main frameworks are trained with both labeled and unlabeled data; however, the labeled data, with precise ground truths as supervision, converges more rapidly compared with the unlabeled data. Thus, the training process will easily get overwhelmed by the supervised training task , as shown in Figure 3. Another challenge lies in the fact that existing SSVMIS methods fail to address the issue of distribution shifts, let alone domain shifts, resulting in a limitation in capturing features that are invariant to changes in distribution.

𝑥

!

𝑥

"

𝑦

Teacher

Student

𝑥

!

𝑥

"

EMA

𝑦

𝑦

𝑦

";$

!

𝑓(𝜃

!

)

𝑓(𝜃

𝟐

)

𝑦

";$

"

𝑦

"

𝑥

!

𝑥

"

𝑦

𝑦

";%

𝑦

𝑦

";&

𝑦

";%,&

Used for prediction

(b) CPS

(a) Mean Teacher

𝐸(𝜉)

𝐷(𝜉)

𝐷(𝜃)

𝐷(𝜓)

Labeled data flow

Unlabeled data flow

UA-MT(Teacher-Student-based)

dnsun

0.6

sup

loss

0.4

weight value

0.2

0.0

-0.2

-0.4

0

50

100

150

200

250

300

CLD (CPS-based)

4

dnsun

sup

3

loss

weightvalue

2

1

{337373

0

-1

-2

0

20

40

60

80

100

120

140

(c) Ours

Based on the similarities and the main weaknesses of the mainstream SSVMIS methods, we argue that a generic framework is possible if we can solve the over-fitting issue and design a powerful methods to capture the distribution-invariant features. To tackle the above issues and design a generic SSVMIS methods for the real-world application scenarios, this work proposes a novel Aggregating & Decoupling (A&D) framework. Specifically, A&D consists of an Aggregating stage and a Decoupling stage. In the Aggregating stage, based on the recent success of the diffusion model [19, 20], we propose a Diff-VNet to aggregate the multi-domain features into one shared encoder to construct a common knowledge set to improve the capacity of capturing the distribution-invariant features. To solve the over-fitting issue, in the Decoupling stage, we decouple the decoding process to (1) a labeled data training flow which mainly updates a Diff-VNet decoder and a difficulty-aware V-Net decoder to generate high-quality pseudo labels; (2) an unlabeled data training flow which mainly updates another vanilla V-Net decoder with the supervision of the pseudo labels. The denoising process of the Diff-VNet decoder provides the domain-unbiased pseudo labels while the difficulty-aware V-Net decoder class-unbiased pseudo labels. We also propose a re-parameterizing & smoothing combination strategy to further improve the quality of the pseudo labels.

The key contributions of our work can be summarized as follows: (1) we unify the SSL, Class Imbalanced SSL, UDA, and SemiDG for volumetric medical image segmentation with one generic framework; (2) we state the over-fitting issues of the current SSL methods and propose to solve it by an efficient data augmentation strategy and decoupling the decoders for labeled data and unlabeled, respectively; (3) we introduce the Diffusion V-Net to learn the underlying feature distribution from different domains to generalize the SSL methods to more realistic application scenarios; (4) The proposed Aggregating & Decoupling framework achieves state-of-the-art on representative datasets of SSL, class-imbalance SSL, UDA, and SemiDG tasks. Notably, our method achieves significant improvements on the Synapse dataset (12.3 in Dice) and the MMWHS dataset in the MR to CT setting (8.5 in Dice). Extensive ablation studies are conducted to validate the effectiveness of the proposed methods.

2 Related Work

2.1 Semi-supervised Segmentation & the Class Imbalance Issue

Semi-supervised segmentation aims to explore tremendous unlabeled data with supervision from limited labeled data. Recently, self-training-based methods [3, 4, 21] have become the mainstream

of this domain. Approaches with consistency regularization strategies [22, 3, 21] achieved good performance by encouraging high similarity between the predictions of two perturbed networks for the same input image, which highly improved the generalization ability. In the medical image domain, the data limitation issue is more natural and serious. Existing approaches [23, 24, 10, 14, 13, 17, 16, 25] to combat the limited data have achieved great success but are bottlenecked by the application scenarios and cannot handle more challenging but practical settings such as UDA and SemiDG.

Class Imbalance Issue Class imbalance issue is a significant problem to extend the existing SSLbased methods to more practical setting, since medical datasets have some classes with notably higher instances in training samples than others. In natural image domain, different means are proposed to solve this issue, including leveraging unlabeled data [26, 27, 28, 29, 30], re-balancing data distributions in loss [31, 30, 32], debiased learning [6, 5] etc. In medical image domain, this issue is more severe but only few work [15, 33, 25] noticed this problem. Incorporating the class-imbalance awareness is crucial for the generalization of the SSL methods.

2.2 Unsupervised Domain Adaptation & Semi-supervised Domain Generalization

Domain adaptation (DA) aims to solve the domain shifts by jointly training the model with source domain data and target domain data. Unsupervised Domain Adaptation (UDA) [34, 35, 36, 37] is the most challenging and practical setting among all the DA setting, since no target labels are required. In this context, UDA is becoming increasingly important in the medical image segmentation field, and as a result, a myriad of UDA approaches have been developed for cross-domain medical image segmentation by leveraging: generative adversarial-based methods [38, 39, 40, 41, 42, 43, 44, 45, 46], semi-supervised learning techniques [47, 48, 49], and self-training as well as contrastive learning techniques [50, 51], etc . Though with promising adaptation results, these methods highly rely on the unlabeled target domain information, which hinders the generalizability.

Domain generalization (DG) is a more strict setting, where the difference with DA is that the model does not use any information from the target domain. Unlike unsupervised domain adaptation, semi-supervised domain generalization (SemiDG) does not assume access to labeled data from the target domains. Existing SemiDG methods [52, 53] leverage various unusual strategies to solve the domain shifts, e.g. , meta-learning [52], Fourier Transformation [54], compositionality [55] etc. , which are not general and have unsatisfactory performance on the tasks such as UDA and SSL.

Compared to these prior works, our work is the first to unify SSL, Imbalanced SSL, UDA, and SemiDG settings. This extension not only amplifies the scope and versatility of SSL-based frameworks in medical image segmentation but also stands in stark contrast to earlier approaches that remained confined to singular domains such as SSL or UDA.

2.3 Diffusion Model

Denoising diffusion models [56, 19, 57, 58] have shown significant success in various generative tasks [59, 60, 61, 62, 63], due to the powerful ability of modeling the underlying distribution of the data, conceptually having a greater capacity to handle challenging tasks. Noticing this property, there has been a rise in interest to incorporate them into segmentation tasks, including both natural image segmentation [64, 65, 66], and medical image segmentation [67, 68, 20] Given the notable achievements of diffusion models in these respective domains, leveraging such models to develop generation-based perceptual models would prove to be a highly promising avenue to push the boundaries of perceptual tasks to newer heights.

3 Method

3.1 Overview of the Aggregating & Decoupling Framework

In this section, we will introduce our Aggregating & Decoupling (A&D) framework, as shown in Figure 4, which consists of an Aggregating stage and a Decoupling stage. The training pipeline is illustrated in Algorithm 1.

The Aggregating stage aims to construct a common knowledge set across domains based on the idea that all data share common underlying high-level knowledge, such as texture information.

!

!;#

"

$

"

!

!

!;%

!

$;&

!

$;#

!

$;%

+ t-step noise

+SVDA

+ SVDA

!

!;%

!

!,#

!

$,#

!

$;%

#

$("

!

, "

$

; ()

*("

$

; +)

*("

!

; ()

*("

!

; ,)

& Smooth

#

#,%

#

Argmax

Softmax

Class-wise Diff.

Re-weighting

Pseudo label

Ground truth

Distribution shift

Diffusion

V-Net

Diffusion

V-Net

V-Net

V-Net

Forward & back propagation

Stop gradient

Loss supervision

Aggregating Stage

Knowledge distillation

By aggregating information from multiple domains and jointly trained, the encoder can capture the underlying domain-invariant features. To achieve this, we introduce a powerful yet efficient Sampling-based Volumetric Data Augmentation (SVDA) strategy to enlarge the distribution diversity and leverage the diffusion model to capture the invariant features of the diversified data.

In the decoding of the existing SSL methods, the decoders are simultaneously trained with both labeled and unlabeled data, which leads to coupling and over-fitting issues and further hinder the extending to the general SSL. The Decoupling stage aims to solve these issues by decoupling the labeled and unlabeled data training flows. Concretely, for the labeled data flow, (1) a diffusion decoder is mainly used to guide the diffusion encoder to learn the distribution-invariant representations through the diffusion backward process, and thus produces the domain-unbiased pseudo labels ; (2) a vanilla V-Net decoder with the proposed difficulty-aware re-weighting strategy is mainly to avoid the model over-fit to the easy and majority classes, and thus produces the class-unbiased pseudo labels . Then, for the unlabeled data flow, the domain- and class-unbiased pseudo labels are ensembled through a proposed Reparameterize & Smooth (RS) strategy to generate high quality pseudo labels. Finally, the pseudo labels are used to supervise the training of an additional V-Net decoder for prediction only.

3.2 Aggregating Stage

Assume that the entire dataset comprises of N L labeled samples { ( x l i , y i ) } N L i =1 and N U unlabeled samples { x u i } N U i =1 , where x i ∈ R D × H × W is the input volume and y i ∈ R K × D × H × W is the groundtruth annotation with K classes. The goal of the aggregating stage is to augment the data with SVDA and encode the labeled ( x l , y ) and unlabeled data x u to high-level distribution-invariant features for denoising labeled data flow h l ; ξ difficulty-aware labeled data flow h l ; ψ , and unlabeled data flow h u .

Sampling-based Volumetric Data Augmentation (SVDA) Instead of the time-consuming traditional data augmentation used in [69] which cascaded all the augmentations, SVDA build upon an augmentation set and N aug operations are randomly sampled to apply to both the labeled and unlabeled data. The augmentation set consists of 3D spatial-wise transforms (random crop, random rotation, random scaling) and voxel-wise transforms (Gaussian blur, brightness, contrast, gamma). N aug is empirically set to 3.

Algorithm 1: Training Pipeline of A&D.

Input:

Labeled samples { ( x l , y ) } L , unlabeled samples { x u } U

i i N i =1 i N i =1

Output: l u u

Diffusion encoder E ( x , x ; ξ ) with V-Net decoder D ( x ; θ ) for inference

1 Initialization;

2 for batched data ( x l i , y i ) , x u i do

3 Add SVDA on x l and x u to obtain ˆ x l and ˆ x u ;

4 Add t step noise on y i to obtain the noisy label y t by Eq. 1;

5 Train the diffusion model E ( x l ; ξ ) and D ( x l ; ξ ) with (ˆ x l i , y t ) by Eq. 2;

6 Train the diffusion encoder E ( x l ; ξ ) and difficulty-aware decoder D ( x l ; ψ ) by Eq. 6;

7 Generate domain-unbiased probability map p u ; ξ with E ( x l , x u ; ξ ) + D ( x l ; ξ ) by DDIM [57];

8 Generate class-unbiased probability map p u ; ψ with D ( x l ; ψ ) with forward pass;

9 Ensemble the two maps and obtain the pseudo label y ξ,ψ by Eq. 7;

10 Train E ( x u ; ξ ) and V-Net decoder D ( x u ; θ ) with unlabeled data (ˆ x u , y ξ,ψ ) and by Eq. 8;

11 end

Diffusion for Capturing Invariant Features We follow Diff-UNet [20] to use diffusion model for perception but modify it to a V-Net version and remove the additional image encoder. Given the labeled volume data x l ∈ R D × W × H and its label y ∈ R D × W × H , we first convert the label to the one-hot format y 0 ∈ R K × D × W × H and add successive t step noise ϵ to obtain the noisy label y t ∈ R K × D × W × H , which is the diffusion forward process:

Then, the noisy label is concatenated with the image x l as the input of the Diff-VNet. Concretely, the high-level features are different for different data flows. For the denoising flow, i.e. , D ( x l ; ξ ) as decoder, the diffusion encoder takes concatenation concat ([ y t , x l ]) and time step t as input to generate the time-step-embedded multi-scale features h l ; ξ i ∈ R i × F × D 2 i × W 2 i × H 2 i where i is the stage and F is the basic feature size. h l ; ξ i are further used by D ( x l ; ξ ) to predict the clear label y 0 . For the difficulty-aware training flow and the unlabeled data flow, i.e. , D ( x l ; ψ ) and D ( x u ; θ ) as decoders, the encoder only takes x l and x u as input to obtain the multi-scale features h l ; ψ i , h u i , respectively. Note that h l ; ψ i , h u i are with same shapes with h l ; ξ i .

3.3 Decoupling Stage

The decoupling stage consists of four steps: supervised denoising training to generate domainunbiased pseudo labels, supervised difficulty-aware training to generate class-unbiased pseudo labels, pseudo labeling to ensemble the two pseudo labels and unsupervised training to get the final predictor.

Supervised Denoising Training with the Diffusion Decoder D ( x l ; ξ ) Taking h l ; ξ i as inputs, D ( x l ; ξ ) decodes the features to predict the clear label y 0 as domain-unbiased pseudo label. The objective function is defined as follow:

where L DiceCE ( x, y ) = 1 2 [ L CE ( x, y ) + L Dice ( x, y )] is the combined dice and cross entropy loss.

Supervised Difficulty-aware Training with D ( x l ; ψ ) To alleviate the common class imbalance issue in SSVMIS, we design a Difficulty-aware Re-weighting Strategy (DRS) to force the model to focus on the most difficult classes ( i.e. the classes learned slower and with worse performances). The difficulty is modeled in two ways with the probability map p l ; ξ produced by diffusion decoder D ( x l ; ξ ) : learning speed and performance. We use Population Stability Index [70] to measure the

learning speed of each class after the e th iteration:

where λ k denotes the Dice score of p l ; ξ of k th class in e th iteration and △ = λ k,e -λ k,e -1 . du k,e and dl k,e denote classes not learned and learned after e th iteration. τ is the number accumulation iterations and set to 50 empirically. Then, we define the difficulty of k th class after e th iteration as:

where the classes learned faster have smaller d k,e , the corresponding weights in the loss function will be smaller to slow down the learn speed. We also accumulate 1 -λ k,e for τ iterations to obtain the reversed dice weight w λ k,e and weight d k,e . In this case, classes with lower dice scores will have larger weights in the loss function, which forces the model to pay more attention to these classes. The overall difficulty-aware weight of k th class is defined as:

where α is empirically set to 1 5 in the experiments to alleviate outliers. The objective function of the supervised difficulty-aware training is defined as follow:

Pseudo Labeling with Reparameterize & Smooth (RS) Strategy The domain-unbiased p u ; ξ probability map is generated by iterating the diffusion model ( E ( x l , x u ; ξ ) + D ( x l ; ξ ) ) t times with the Denoising Diffusion Implicit Models (DDIM) method [57]. The class-unbiased p u ; ψ probability map can be obtained by D ( x l ; ψ ) with stopped gradient forward pass. We ensemble p u ; ξ and p u ; ψ to generate high-quality pseudo labels. However, when combining these two maps, we found that the denoised probability map p u ; ψ is too sparse, i.e. , with very high confidence of each class. This property is benefit for the fully-supervised tasks, but in this situation, it will suppress p u ; ψ and is not robust to noise and error. Thus, we re-parameterize p u ; ψ with the Gumbel-Softmax to add some randomness and using Gaussian blur kernel to remove the noise brought by this operation. The final pseudo label is:

Unsupervised Training with D ( x u ; θ ) Finally, we can use the pseudo label y ξ,ψ to train D ( x u ; θ ) in an unsupervised manner. The objective function of the unsupervised training is defined as:

To better leverage the domain- and class-unbiased features, we also transmit with knowledge distillation strategy: θ = w ema × θ +(1 -w ema ) × ( ξ + ψ ) / 2 , w ema = 0 . 99 . The overall training function of the A&G framework is:

where µ is empirically set as 10 and follow [15] to use the epoch-dependent Gaussian ramp-up strategy to gradually enlarge the ratio of unsupervised loss. In the inference stage, only the diffusion encoder E ( x l , x u ; ξ ) and D ( x u ; θ ) are used to generate the predictions.

4 Experiments

4.1 Datasets and Implementation Details

We evaluate our proposed A&D framework on four datasets for four tasks, i.e. , LASeg dataset [72] for SSL, Synapse dataset [73] for class imbalanced SSL, MMWHS dataset [74] for UDA, and M&Ms

dataset [75] for SemiDG. For more details, please refer to the Appendix. We implement the proposed framework with PyTorch, using a single NVIDIA A100 GPU. The network parameters are optimized by SGD with Nesterov and momentum of 0.9. We employ a 'poly' decay strategy follow [69]. For more implementation details, e.g. , data preprocessing, learning rates, batch sizes, etc. , please refer to the Appendix. We evaluate the prediction of the network with two metrics, including Dice and the average surface distance (ASD). For the LASeg dataset, we employ additional metrics, Jaccard and HD95, following [8]. Final segmentation results are obtained using a sliding window strategy.

4.2 Experiment results

Our proposed A&D framework achieves state-of-the-art on all the four settings, i.e. , SSL (Table 2), imbalanced SSL (Table 1), UDA (Table 3), and SemiDG (Table 4). In particular, our method achieves

𝑥

!

𝑥

"

𝑦

";$,&

𝐸(𝜉)

𝐷(𝜉)

𝐷(𝜃)

𝐷(𝜓)

𝑥

!

𝑥

"

𝑦

";$,&

𝐸(𝜉)

𝐷(𝜉)

𝐷(𝜃)

𝐷(𝜓)

𝑥

!

𝑥

"

𝑦

";$,&

𝐸(𝜉)

𝐷(𝜉)

𝐷(𝜃)

𝐷(𝜓)

𝑥

!

𝑥

"

𝑦

";$,&

𝐸(𝜓)

𝐷(𝜉)

𝐷(𝜃)

𝐷(𝜓)

𝐸(𝜉)

𝐸(𝜃)

𝑥

!

𝑥

"

𝑦

";$,&

𝐸(𝜓)

𝐷(𝜉)

𝐷(𝜃)

𝐷(𝜓)

𝐸(𝜉)

𝐸(𝜃)

(a)

(b)

(c)

(d)

(e)

significant improvements over the existing state-of-the-art approach on the Synapse dataset with 20% labeled data and the MMWHS dataset in the MR to CT setting, demonstrating a substantial increase of 12.3 in Dice score for the Synapse dataset and 8.5 for the MMWHS dataset.

4.3 Analyses

Architecture Analysis As shown in Figure 5, we test different architectures for the framework, the corresponding results are in Table 5. When using three separate networks (a)(b), the performance drop significantly, the reason is that this structure can not learn the domain-invariant features. Using pure V-Net (c) or diffusion-based networks (d) also leads to inferior results, especially for the pure

diffusion model, it is hard to train due to the limited labeled data, high-quality pseudo labels are hard to obtain, which further hinder the training of the unsupervised branch.

Ablation on the components We analyze the effectiveness of different components in our method. According to the results in Table 6, on MMWHS MR to CT setting (UDA), when removing the diffusion model, performance decreases the most, which indicated the importance ability of the diffusion model for capturing the distribution-invariant features. The result of removing the SVDA also

Table 6: Ablation study of the components in our framework on 20% labeled Synapse setting (IBSSL), MMWHSMRtoCTsetting (UDA) and 2% labeled 2% labeled M&Ms setting (SemiDG).

indicates that when the data is not diverse enough, the diffusion model cannot capture effective underlying features. On Synapse dataset for IBSSL, the results are slightly different with those in the UDA setting, the DRS plays more important role than the Diffusion. In 5% M&Ms dataset for SemiDG, the results are quite aligned with results in the UDA setting.

Hyper-parameter Selection We evaluate the performance of our method under different time step t of the diffusion model and the number of sampled augmentations, as shown in Table 7.

Table 7: Ablation study of the noise step t and the number of sampled augmentations on MMWHS CT to MR setting.

CT to MR

MR to CT

GT

Pred

GT

Pred

2D View

3D View

Visualization We also present the visualization results to further analyze our method. As shown in Figure 6, for the UDA task, our framework can generate smoother volumetric objects. Our framework also has detect minority classes, as can be seen in Appendix, which indicates the effectiveness of incorporating the class-imbalance awareness with the proposed difficulty-aware re-weighing strategy.

Limitations The diffusion process introduces additional training costs; however, the inference efficiency remains unaffected as only the decoder trained using unlabeled data is utilized during inference. Moreover, the failure cases are mainly in the M&Ms dataset for SemiDG setting. Specifically, our method usually fails on the first and the last slices along the depth axis. Due to the restricted depth dimension (less than 10), the 2D slices and the corresponding masks vary significantly. In such a case, it is hard for our volumetric framework to capture depth-wise information for the first or last slice with only one neighboring slice as a reference, and thus leads to false positive results.

5 Conclusion

In this paper, we propose a generic framework for semi-supervised learning in volumetric medical image segmentation, called Aggregating & Decoupling. This framework addresses four related settings, namely SSL, class imbalanced SSL, UDA, and SemiDG. Specifically, the aggregating part of our framework utilizes a Diffusion encoder to construct a 'common knowledge set' by extracting distribution-invariant features from aggregated information across multiple distributions/domains. On the other hand, the decoupling part involves three decoders that facilitate the training process by decoupling labeled and unlabeled data, thus mitigating overfitting to labeled data, specific domains, and classes. Experimental results validate the effectiveness of our proposed method under four settings.

The significance of this work lies in its ability to encourage semi-supervised medical image segmentation methods to address more complex real-world application scenarios, rather than just developing frameworks in ideal experimental environments. Furthermore, we have consolidated all four settings within a single codebase, enabling the execution of any task using a single bash file by merely adjusting the arguments. We believe that this consolidated codebase will be convenient for further research and beneficial for the community.

References

A. Tarvainen and H. Valpola, 'Mean teachers are better role models: Weight-averaged consistency targets improve semi-supervised deep learning results,' Advances in neural information processing systems , vol. 30, 2017.

K. Sohn, D. Berthelot, N. Carlini, Z. Zhang, H. Zhang, C. A. Raffel, E. D. Cubuk, A. Kurakin, and C.-L. Li, 'Fixmatch: Simplifying semi-supervised learning with consistency and confidence,' Advances in neural information processing systems , vol. 33, pp. 596-608, 2020.

X. Chen, Y. Yuan, G. Zeng, and J. Wang, 'Semi-supervised semantic segmentation with cross pseudo supervision,' in CVPR , pp. 2613-2622, 2021.

Y. Wang, H. Wang, Y. Shen, J. Fei, W. Li, G. Jin, L. Wu, R. Zhao, and X. Le, 'Semi-supervised semantic segmentation using unreliable pseudo-labels,' in CVPR , pp. 4248-4257, 2022.

X. Wang, Z. Wu, L. Lian, and S. X. Yu, 'Debiased learning from naturally imbalanced pseudo-labels,' in CVPR , pp. 14647-14657, 2022.

B. Chen, J. Jiang, X. Wang, P. Wan, J. Wang, and M. Long, 'Debiased self-training for semi-supervised learning,' in Advances in Neural Information Processing Systems , 2022.

H. Wu, X. Li, and K.-T. Cheng, 'Exploring feature representation learning for semi-supervised medical image segmentation,' IEEE Transactions on Neural Networks and Learning Systems , 2023.

L. Yu, S. Wang, X. Li, C.-W. Fu, and P.-A. Heng, 'Uncertainty-aware self-ensembling model for semisupervised 3d left atrium segmentation,' in MICCAI , pp. 605-613, 2019.

S. Li, C. Zhang, and X. He, 'Shape-aware semi-supervised 3d semantic segmentation for medical images,' in MICCAI , pp. 552-561, Springer, 2020.

X. Luo, W. Liao, J. Chen, T. Song, Y. Chen, S. Zhang, N. Chen, G. Wang, and S. Zhang, 'Efficient semi-supervised gross target volume of nasopharyngeal carcinoma segmentation via uncertainty rectified pyramid consistency,' in MICCAI , pp. 318-329, 2021.

X. Luo, J. Chen, T. Song, and G. Wang, 'Semi-supervised medical image segmentation through dual-task consistency,' in Proceedings of the AAAI Conference on Artificial Intelligence , vol. 35, pp. 8801-8809, 2021.

Y. Wu, M. Xu, Z. Ge, J. Cai, and L. Zhang, 'Semi-supervised left atrium segmentation with mutual consistency training,' in MICCAI , pp. 297-306, Springer, 2021.

J. Wang and T. Lukasiewicz, 'Rethinking bayesian deep learning methods for semi-supervised volumetric medical image segmentation,' in CVPR , pp. 182-190, 2022.

H. Wu, Z. Wang, Y. Song, L. Yang, and J. Qin, 'Cross-patch dense contrastive learning for semi-supervised segmentation of cellular nuclei in histopathologic images,' in CVPR , pp. 11666-11675, 2022.

Y. Lin, H. Yao, Z. Li, G. Zheng, and X. Li, 'Calibrating label distribution for class-imbalanced barelysupervised knee segmentation,' in MICCAI , pp. 109-118, 2022.

C. You, R. Zhao, L. H. Staib, and J. S. Duncan, 'Momentum contrastive voxel-wise representation learning for semi-supervised volumetric medical image segmentation,' in MICCAI , pp. 639-652, 2022.

Y. Wu, Z. Wu, Q. Wu, Z. Ge, and J. Cai, 'Exploring smoothness and class-separation for semi-supervised medical image segmentation,' in MICCAI , pp. 34-43, 2022.

H. Wu, X. Li, Y. Lin, and K.-T. Cheng, 'Compete to win: Enhancing pseudo labels for barely-supervised medical image segmentation,' IEEE Transactions on Medical Imaging , 2023.

J. Ho, A. Jain, and P. Abbeel, 'Denoising diffusion probabilistic models,' Advances in Neural Information Processing Systems , vol. 33, pp. 6840-6851, 2020.

Z. Xing, L. Wan, H. Fu, G. Yang, and L. Zhu, 'Diff-unet: A diffusion embedded network for volumetric segmentation,' arXiv preprint arXiv:2303.10326 , 2023.

J. Fan, B. Gao, H. Jin, and L. Jiang, 'Ucc: Uncertainty guided cross-head co-training for semi-supervised semantic segmentation,' in CVPR , pp. 9947-9956, 2022.

Y. Ouali, C. Hudelot, and M. Tami, 'Semi-supervised semantic segmentation with cross-consistency training,' in CVPR , pp. 12674-12684, 2020.

X. Li, L. Yu, H. Chen, C.-W. Fu, and P.-A. Heng, 'Semi-supervised skin lesion segmentation via transformation consistent self-ensembling model,' BMVC , 2018.

X. Li, L. Yu, H. Chen, C.-W. Fu, L. Xing, and P.-A. Heng, 'Transformation-consistent self-ensembling model for semisupervised medical image segmentation,' IEEE Transactions on Neural Networks and Learning Systems , vol. 32, no. 2, pp. 523-534, 2020.

H. Wang and X. Li, 'Dhc: Dual-debiased heterogeneous co-training framework for class-imbalanced semi-supervised medical image segmentation,' in MICCAI , pp. 582-591, Springer, 2023.

C. Wei, K. Sohn, C. Mellina, A. Yuille, and F. Yang, 'Crest: A class-rebalancing self-training framework for imbalanced semi-supervised learning,' in CVPR , pp. 10857-10866, 2021.

Y. Fan, D. Dai, A. Kukleva, and B. Schiele, 'Cossl: Co-learning of representation and classifier for imbalanced semi-supervised learning,' in CVPR , pp. 14574-14584, 2022.

H. Chen, Y. Fan, Y. Wang, J. Wang, B. Schiele, X. Xie, M. Savvides, and B. Raj, 'An embarrassingly simple baseline for imbalanced semi-supervised learning,' arXiv preprint arXiv:2211.11086 , 2022.

Y. Oh, D.-J. Kim, and I. S. Kweon, 'Daso: Distribution-aware semantics-oriented pseudo-label for imbalanced semi-supervised learning,' in CVPR , pp. 9786-9796, 2022.

Z. Lai, C. Wang, S.-c. Cheung, and C.-N. Chuah, 'Sar: Self-adaptive refinement on pseudo labels for multiclass-imbalanced semi-supervised learning,' in CVPR , pp. 4091-4100, 2022.

Y. Hong, S. Han, K. Choi, S. Seo, B. Kim, and B. Chang, 'Disentangling label distribution for long-tailed visual recognition,' in CVPR , pp. 6626-6636, 2021.

S. Yu, J. Guo, R. Zhang, Y. Fan, Z. Wang, and X. Cheng, 'A re-balancing strategy for class-imbalanced classification based on instance difficulty,' in CVPR , pp. 70-79, 2022.

H. Basak, S. Ghosal, and R. Sarkar, 'Addressing class imbalance in semi-supervised image segmentation: A study on cardiac mri,' in MICCAI , pp. 224-233, 2022.

J.-Y. Zhu, T. Park, P. Isola, and A. A. Efros, 'Unpaired image-to-image translation using cycle-consistent adversarial networks,' in Proceedings of the IEEE international conference on computer vision , pp. 22232232, 2017.

Y. Huo, Z. Xu, H. Moon, S. Bao, A. Assad, T. K. Moyo, M. R. Savona, R. G. Abramson, and B. A. Landman, 'Synseg-net: Synthetic segmentation without target modality ground truth,' IEEE transactions on medical imaging , vol. 38, no. 4, pp. 1016-1025, 2018.

J. Hoffman, E. Tzeng, T. Park, J.-Y. Zhu, P. Isola, K. Saenko, A. Efros, and T. Darrell, 'Cycada: Cycleconsistent adversarial domain adaptation,' in International conference on machine learning , pp. 1989-1998, Pmlr, 2018.

Y.-H. Tsai, W.-C. Hung, S. Schulter, K. Sohn, M.-H. Yang, and M. Chandraker, 'Learning to adapt structured output space for semantic segmentation,' in CVPR , pp. 7472-7481, 2018.

Q. Dou, C. Ouyang, C. Chen, H. Chen, B. Glocker, X. Zhuang, and P.-A. Heng, 'Pnp-adanet: Plug-and-play adversarial domain adaptation network at unpaired cross-modality cardiac segmentation,' IEEE Access , vol. 7, pp. 99065-99076, 2019.

C. Chen, Q. Dou, H. Chen, J. Qin, and P.-A. Heng, 'Synergistic image and feature adaptation: Towards cross-modality domain adaptation for medical image segmentation,' in Proceedings of the AAAI conference on artificial intelligence , vol. 33, pp. 865-872, 2019.

M. Bateson, H. Kervadec, J. Dolz, H. Lombaert, and I. Ben Ayed, 'Source-relaxed domain adaptation for image segmentation,' in MICCAI , pp. 490-499, Springer, 2020.

C. Chen, Q. Dou, H. Chen, J. Qin, and P. A. Heng, 'Unsupervised bidirectional cross-modality adaptation via deeply synergistic image and feature alignment for medical image segmentation,' IEEE transactions on medical imaging , vol. 39, no. 7, pp. 2494-2505, 2020.

D. Zou, Q. Zhu, and P. Yan, 'Unsupervised domain adaptation with dual-scheme fusion network for medical image segmentation.,' in IJCAI , pp. 3291-3298, 2020.

X. Han, L. Qi, Q. Yu, Z. Zhou, Y. Zheng, Y. Shi, and Y. Gao, 'Deep symmetric adaptation network for cross-modality medical image segmentation,' IEEE transactions on medical imaging , vol. 41, no. 1, pp. 121-132, 2021.

K. Yao, Z. Su, K. Huang, X. Yang, J. Sun, A. Hussain, and F. Coenen, 'A novel 3d unsupervised domain adaptation framework for cross-modality medical image segmentation,' IEEE Journal of Biomedical and Health Informatics , vol. 26, no. 10, pp. 4976-4986, 2022.

X. Bian, X. Luo, C. Wang, W. Liu, and X. Lin, 'Dda-net: Unsupervised cross-modality medical image segmentation via dual domain adaptation,' Computer Methods and Programs in Biomedicine , vol. 213, p. 106531, 2022.

H. Shin, H. Kim, S. Kim, Y. Jun, T. Eo, and D. Hwang, 'Cosmos: Cross-modality unsupervised domain adaptation for 3d medical image segmentation based on target-aware domain translation and iterative self-training,' arXiv preprint arXiv:2203.16557 , 2022.

C. S. Perone, P. Ballester, R. C. Barros, and J. Cohen-Adad, 'Unsupervised domain adaptation for medical imaging segmentation with self-ensembling,' NeuroImage , vol. 194, pp. 1-11, 2019.

Y. Xia, D. Yang, Z. Yu, F. Liu, J. Cai, L. Yu, Z. Zhu, D. Xu, A. Yuille, and H. Roth, 'Uncertainty-aware multi-view co-training for semi-supervised medical image segmentation and domain adaptation,' Medical image analysis , vol. 65, p. 101766, 2020.

X. Liu, F. Xing, N. Shusharina, R. Lim, C.-C. Jay Kuo, G. El Fakhri, and J. Woo, 'Act: Semi-supervised domain-adaptive medical image segmentation with asymmetric co-training,' in MICCAI , pp. 66-76, Springer, 2022.

L. Liu, Z. Zhang, S. Li, K. Ma, and Y. Zheng, 'S-cuda: Self-cleansing unsupervised domain adaptation for medical image segmentation,' Medical Image Analysis , vol. 74, p. 102214, 2021.

M. Gu, S. Vesal, M. Thies, Z. Pan, F. Wagner, M. Rusu, A. Maier, and R. Kosti, 'Confuda: Contrastive fewshot unsupervised domain adaptation for medical image segmentation,' arXiv preprint arXiv:2206.03888 , 2022.

X. Liu, S. Thermos, A. O'Neil, and S. A. Tsaftaris, 'Semi-supervised meta-learning with disentanglement for domain-generalised medical image segmentation,' in MICCAI , pp. 307-317, Springer, 2021.

X. Liu, S. Thermos, A. Chartsias, A. O'Neil, and S. A. Tsaftaris, 'Disentangled representations for domain-generalized cardiac segmentation,' in Statistical Atlases and Computational Models of the Heart. M&Ms and EMIDEC Challenges: 11th International Workshop, STACOM 2020, Held in Conjunction with MICCAI 2020, Lima, Peru, October 4, 2020, Revised Selected Papers 11 , pp. 187-195, Springer, 2021.

H. Yao, X. Hu, and X. Li, 'Enhancing pseudo label quality for semi-supervised domain-generalized medical image segmentation,' in Proceedings of the AAAI Conference on Artificial Intelligence , vol. 36,

pp. 3099-3107, 2022.

X. Liu, S. Thermos, P. Sanchez, A. Q. O'Neil, and S. A. Tsaftaris, 'vmfnet: Compositionality meets domain-generalised segmentation,' in MICCAI , pp. 704-714, Springer, 2022.

J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli, 'Deep unsupervised learning using nonequilibrium thermodynamics,' in International Conference on Machine Learning , pp. 2256-2265, PMLR, 2015.

J. Song, C. Meng, and S. Ermon, 'Denoising diffusion implicit models,' arXiv preprint arXiv:2010.02502 , 2020.

A. Q. Nichol and P. Dhariwal, 'Improved denoising diffusion probabilistic models,' in International Conference on Machine Learning , pp. 8162-8171, PMLR, 2021.

A. Nichol, P. Dhariwal, A. Ramesh, P. Shyam, P. Mishkin, B. McGrew, I. Sutskever, and M. Chen, 'Glide: Towards photorealistic image generation and editing with text-guided diffusion models,' arXiv preprint arXiv:2112.10741 , 2021.

P. Dhariwal and A. Nichol, 'Diffusion models beat gans on image synthesis,' Advances in Neural Information Processing Systems , vol. 34, pp. 8780-8794, 2021.

C. Saharia, J. Ho, W. Chan, T. Salimans, D. J. Fleet, and M. Norouzi, 'Image super-resolution via iterative refinement,' IEEE Transactions on Pattern Analysis and Machine Intelligence , 2022.

J. Ho, T. Salimans, A. Gritsenko, W. Chan, M. Norouzi, and D. J. Fleet, 'Video diffusion models,' arXiv preprint arXiv:2204.03458 , 2022.

B. L. Trippe, J. Yim, D. Tischer, T. Broderick, D. Baker, R. Barzilay, and T. Jaakkola, 'Diffusion probabilistic modeling of protein backbones in 3d for the motif-scaffolding problem,' arXiv preprint arXiv:2206.04119 , 2022.

T. Amit, T. Shaharbany, E. Nachmani, and L. Wolf, 'Segdiff: Image segmentation with diffusion probabilistic models,' arXiv preprint arXiv:2112.00390 , 2021.

T. Chen, L. Li, S. Saxena, G. Hinton, and D. J. Fleet, 'A generalist framework for panoptic segmentation of images and videos,' arXiv preprint arXiv:2210.06366 , 2022.

Y. Ji, Z. Chen, E. Xie, L. Hong, X. Liu, Z. Liu, T. Lu, Z. Li, and P. Luo, 'Ddp: Diffusion model for dense visual prediction,' arXiv preprint arXiv:2303.17559 , 2023.

J. Wu, H. Fang, Y. Zhang, Y. Yang, and Y. Xu, 'Medsegdiff: Medical image segmentation with diffusion probabilistic model,' arXiv preprint arXiv:2211.00611 , 2022.

J. Wolleb, R. Sandkühler, F. Bieder, P. Valmaggia, and P. C. Cattin, 'Diffusion models for implicit image segmentation ensembles,' in International Conference on Medical Imaging with Deep Learning , pp. 1336-1348, PMLR, 2022.

F. Isensee, P. F. Jaeger, S. A. Kohl, J. Petersen, and K. H. Maier-Hein, 'nnu-net: a self-configuring method for deep learning-based biomedical image segmentation,' Nature methods , vol. 18, no. 2, pp. 203-211, 2021.

H. Jeffreys, 'An invariant form for the prior probability in estimation problems,' Proceedings of the Royal Society of London. Series A. Mathematical and Physical Sciences , vol. 186, no. 1007, pp. 453-461, 1946.

L.-Z. Guo and Y.-F. Li, 'Class-imbalanced semi-supervised learning with adaptive thresholding,' in ICML , pp. 8082-8094, PMLR, 2022.

Z. Xiong, Q. Xia, Z. Hu, N. Huang, C. Bian, Y. Zheng, S. Vesal, N. Ravikumar, A. Maier, X. Yang, et al. , 'A global benchmark of algorithms for segmenting the left atrium from late gadolinium-enhanced cardiac magnetic resonance imaging,' Medical image analysis , vol. 67, p. 101832, 2021.

B. Landman, Z. Xu, J. Igelsias, M. Styner, T. Langerak, and A. Klein, '2015 miccai multi-atlas labeling beyond the cranial vault-workshop and challenge,' 2015.

X. Zhuang and J. Shen, 'Multi-scale patch and multi-modality atlases for whole heart segmentation of mri,' Medical image analysis , vol. 31, pp. 77-87, 2016.

V. M. Campello, P. Gkontra, C. Izquierdo, C. Martin-Isla, A. Sojoudi, P. M. Full, K. Maier-Hein, Y. Zhang, Z. He, J. Ma, et al. , 'Multi-centre, multi-vendor and multi-disease cardiac segmentation: the m&ms challenge,' IEEE Transactions on Medical Imaging , vol. 40, no. 12, pp. 3543-3554, 2021.

Y. Bai, D. Chen, Q. Li, W. Shen, and Y. Wang, 'Bidirectional copy-paste for semi-supervised medical image segmentation,' in CVPR , pp. 11514-11524, 2023.

M. Jafari, S. Francis, J. M. Garibaldi, and X. Chen, 'Lmisa: A lightweight multi-modality image segmentation network via domain adaptation using gradient magnitude and shape constraint,' Medical Image Analysis , vol. 81, p. 102536, 2022.

H. Li, Y. Wang, R. Wan, S. Wang, T.-Q. Li, and A. Kot, 'Domain generalization for medical imaging classification with linear-dependency regularization,' Advances in Neural Information Processing Systems , vol. 33, pp. 3118-3129, 2020.

Q. Liu, Q. Dou, and P.-A. Heng, 'Shape-aware meta-learning for generalizing prostate mri segmentation to unseen domains,' in MICCAI , pp. 475-485, Springer, 2020.

A Appendix

GT

A&D (ours)

CLD

2D View

3D View

mage0033

image0033

R

A.1 More Details of Datasets and Implementation

The hyper-parameters for different datasets are shown in Table 8.

The details of the datasets and the pre-processing operations are as follows.

LASeg Dataset for SSL The Atrial Segmentation Challenge (LASeg) dataset [72] provides 100 3D gadolinium-enhanced MR imaging scans (GE-MRIs) and LA segmentation masks for training and validation. Following previous work [8, 17], we split the 100 scans into 80 for training and 20 for evaluation. We use the processed datasets from [8] where all the scans were cropped centering at the heart region for better comparison of the segmentation performance of different methods and normalized as zero mean and unit variance. In the training stage, SS-Net [17] and BCP [76] use test set for validation to select the best model, which is unreasonable. We use labeled data instead and achieve better performances.

Synapse Dataset for Class Imbalanced SSL The Synapse [73] dataset has 13 foreground classes, including spleen (Sp), right kidney (RK), left kidney (LK), gallbladder (Ga), esophagus (Es), liver(Li), stomach(St), aorta (Ao), inferior vena cava (IVC), portal & splenic veins (PSV), pancreas (Pa), right adrenal gland (RAG), left adrenal gland (LAG) with one background and 30 axial contrast-enhanced abdominal CT scans. We randomly split them as 20,4 and 6 scans for training, validation, and testing, respectively. Following DHC [25], we run the experiments 3 times with different random seeds.

MMWHSDataset for UDA Multi-Modality Whole Heart Segmentation Challenge 2017 dataset (MMWHS) [74] is a cardiac segmentation dataset including two modality images (MR and CT). Each modality contains 20 volumes collected from different sites, and no pair relationship exists between modalities. Following the previous work [41], we choose four classes of cardiac structures. They are the ascending aorta (AA), the left atrium blood cavity (LAC), the left ventricle blood cavity (LVC), and the myocardium of the left ventricle (MYO). For the pre-processing, follow [41], (1) all the scans

were cropped centering at the heart region, with four cardiac substructures selected for segmentation; (2) for each 3D cropped image top 2% of its intensity histogram was cut off for alleviating artifacts; (3) each 3D image was then normalized to [0, 1]. To make a fair comparison, we keep the test set the same with prior arts [38, 41, 44].

M&Ms Dataset for SemiDG The multi-center, multi-vendor & multi-disease cardiac image segmentation (M&Ms) dataset [75] contains 320 subjects, which are scanned at six clinical centers in three different countries by using four different magnetic resonance scanner vendors, i.e., Siemens, Philips, GE, and Canon. We consider the subjects scanned from different vendors are from different domains (95 in domain A, 125 in domain B, 50 in domain C, and another 50 in domain D). We use each three of them as the source domain for training and the rest as the unseen domain for testing. For the pre-processing, (1) all the scans were cropped with four cardiac substructures selected for segmentation; (2) for each 3D cropped image top and bottom 0.5% of its intensity histogram was cut off for alleviating artifacts; (3) each 3D image was then normalized to [0, 1]. Since the data has very few slices on the z-axis (less than 16), the previous work used 2D-based solutions. In this work, since we aim to design a generic framework for volumetric medical image segmentation, we stacked the z-axis to 32 to meet the minor requirement for our encoder with four down-sampling layers. This dataset can also be considered as the extreme case of 3D segmentation tasks.

𝑝

!;#

𝑝

!;$

RS

Softmax

Sum

Argmax

Pseudo label (red)

on original image

A.2 More Analyses

Comparison of Computational Costs We compared the parameters and the inference time of our method with the most SOTA methods in different settings: SS-Net [17] on LASeg dataset for SSL, CLD [15] on Synapse dataset for IBSSL, and vMFNet [55] as well as EPL [54] on M&Ms dataset for SemiDG, as shown in Table 9. From the table we can observe that our method has the fewest parameters except for vMFNet. Although vMFNet has only 20 M parameters, their performances (73.88% and 72.3%) on SSL and UDA tasks are significantly lower than ours (89.40% and 90.0%).

𝑝

!;#

𝑝

!;$

RS

Softmax

Sum

Argmax

Pseudo label

of MYO

on original image

ddm_tumpkp_1

As for the inference time, we can observe that our A&D is the fastest. This can be attributed to our effective aggregating and decoupling strategies, enabling our method to exclusively utilize the unlabeled branch for inference.

Visualization of the Re-parameterize & Smooth (RS) As shown in Figure 8, the output probability map p u ; ξ of diffusion with DDIM D ( ξ ) is with very high confidence with its prediction, however, the results are not stable since the unlabeled data is unseen during the training process of the diffusion decoder, especially for some problematic classes (MYO of MMWHS, Figure 9) with ambiguous boundaries and noise. Thus, if we sum it with the map p u ; ψ generated by the V-Net decoder D ( ψ ) , the error regions (e.g., upper right corner) with high confidence will surpass some correct regions of p u ; ψ with lower confidence and further harm the quality of the final pseudo label. Moreover, in some cases, the two output probability maps have complementary properties (Figure 8), indicating the effectiveness of ensembling them for the high-quality pseudo labels.

Ablation on the Effectiveness of Decoupling the Labeled and Unlabeled Data Training Flows Based on our final framework, we add an additional training process with labeled data on the decoder D ( x u ; θ ) trained with unlabeled data to verify the effectiveness of the decoupling idea. Compared with the final A&D framework, when adding an additional labeled data training branch, the performance in terms of Dice drops from 90.03% to 86.94% on the MR to CT setting of the MMWHSdataset. The result indicates that when the predictor is trained with labeled and unlabeled data, it may get over-fitted to the easier labeled data flow, which verifies the effectiveness of the key idea of our decoupling stage.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles semi-supervised medical image segmentation over multiple domains, generalizing unsupervised domain adaptation and semi-supervised learning. A multiple branch encoder-decoder architecture with a shared encoder is proposed. The first branch is a diffusion model with V-Net architecture trained over labeled images and their noisy labels, where the encoder aims to extract domain-invariant features. The second decoder is trained over labeled images to combat class-imbalance by reweighting the loss with respect to Population Stability Index. The third branch is trained over unlabeled images with pseudolabels, where pseudolabels are ensembles of probability maps predicted by the other two decoders. To account for the sparse predictions of diffusion model decoder, its probability predictions are reparamatrized via Gumbel-Softmax prior to ensembling into pseudolabels. Each branch is trained over a combination of cross-entropy and dice loss, with input images and labels varying with respect to the branch.\n\nSTRENGTHS:\nAdopting diffusion models for generalized semi-supervised segmentation and domain adaptation is novel. Abundant experiments over multiple tasks and datasets demonstrate the benefit of the proposed method, particularly for unsupervised domain adaptation.\n\nWEAKNESSES:\nUnsupervised/semi-supervised segmentation and domain adaptation in medical imaging are well-explored domains with several missing related works:\n\nGu et al. “ConFUDA: Contrastive Fewshot Unsupervised Domain Adaptation for Medical Image Segmentation”, 2022\n\nXia et al. “Uncertainty-aware multi-view co-training for semi-supervised medical image segmentation and domain adaptation”, 2020\n\nBian et al. “DDA-Net: Unsupervised cross-modality medical image segmentation via dual domain adaptation”, 2021\n\nPerone et al. “Unsupervised domain adaptation for medical imaging segmentation with self-ensembling”, 2019\n\nShin et al. “COSMOS: Cross-Modality Unsupervised Domain Adaptation for 3D Medical Image Segmentation based on Target-aware Domain Translation and Iterative Self-Training”, 2022\n\nBateson et al. “Source-Relaxed Domain Adaptation for Image Segmentation”, 2021\n\nLiu et al. “S-CUDA: Self-cleansing unsupervised domain adaptation for medical image segmentation”, 2021\n\nLiu et al. “ACT: Semi-supervised Domain-adaptive Medical Image Segmentation with Asymmetric Co-Training”, 2022\n\nParticularly, Liu et al. 2022 propose decoupled semi-supervised training and domain adaptation branches, a similar motivation to the proposed method. The authors should highlight their novelty against these methods that are not discussed.\n\nQUESTIONS:\nNone\n\nLIMITATIONS:\nLimitations are discussed", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-27"}

comments: SUMMARY:
This paper addressed both the semi-supervised issue and the domain adaptation issue for medical image segmentation. The proposed method incorporated the diffusion model and used a consistency-based regularization for training. Performance on several public datasets is good for the segmentation tasks.

STRENGTHS:
1. Applying a diffusion model to reduce the domain gaps is reasonable and achieves performance gains according to the results.
2. Extensive experiments and sufficient comparisons.
3. The motivation is sound. Considering both tasks together is now the research attention.

WEAKNESSES:
1. There are some recent papers considering the two tasks as well. See:
[1] Bai Y, Chen D, Li Q, et al. Bidirectional Copy-Paste for Semi-Supervised Medical Image Segmentation[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023: 11514-11524.
[2] Wu H, Li X, Lin Y, et al. Compete to Win: Enhancing Pseudo Labels for Barely-supervised Medical Image Segmentation[J]. IEEE Transactions on Medical Imaging, 2023.
Therefore, the problem definition is new.

2. It's a bit over-claimed to unify four tasks together. For example, this paper does not provide unique designs for imbalanced learning. Do not make the claims too broad.
3. The performance is not SOTA, especially on the popular LA dataset, see [1]. Moreover, about the comparison, it is weird that there are different compared methods for different tasks.
4. Regarding the main task (semi-supervised learning), the technical insight is not clear. The domain gap should be a critical issue here. Please elaborate more on the take-home insights rather than tricks.

QUESTIONS:
See the above weaknesses

LIMITATIONS:
See the above weaknesses
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-03

comments: SUMMARY:
This work proposes a generic framework for semi-supervised learning (SSL) in medical image segmentation. They identify two obstacles that prevent the applicability of SSL in real-world scenarios. Then they propose Aggregating & Decoupling to alleviate the issues under several settings, including class imbalanced SSL, UDA, and SemiDG. They show the proposed method can mitigate overfitting to labeled data, specific domains, and classes across four settings and many datasets.

STRENGTHS:
The paper is well-structured and easy to follow. The proposed method is generalizable to four different settings, enhancing the applicability of SSL in real-world scenarios. The presentation is very clear, especially the introduction, where the obstacles, and motivations are clearly stated.

WEAKNESSES:
My biggest concern about this work is the effectiveness in more generic applications, e.g., standard computer vision benchmarks. I appreciate the authors identifying obstacles that hurt the performance of SSL in real-world applications. However, what is the relationship between the proposed approach and medical image segmentation? Is any proposed component tailored to deal with the issues with medical images? In my eye, the proposed framework is generic and should be working widely well in many applications instead of only focusing on medical images. It would be great if the authors could verify the effectiveness of the proposed methods in CV benchmarks instead of only medical imaging applications if the proposed method is not designed or tailored to medical settings. I would consider re-rating if the authors response to my concern.

QUESTIONS:
See weakness.

LIMITATIONS:
Yes.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

comments: SUMMARY:
This paper proposes an Aggregating & Decoupling framework, which unifies the SSL, Class Imbalanced SSL, UDA, and SemiDG for volumetric medical image segmentation with one generic framework.

STRENGTHS:
This paper explores a generic framework for the SSL, Class Imbalanced SSL, UDA, and SemiDG of volumetric medical image segmentation.

WEAKNESSES:
The description of the method is not clear.

QUESTIONS:
1. How are the running time and parameters of the proposed method compared to other methods?
2. The proposed approach appears to require separate training for different tasks, and whether it is possible to train only once for the SSL, Class Imbalanced SSL, UDA, and SemiDG, and truly achieve a generic framework.
3. Do ablation experiments on other datasets also demonstrate the performance of the proposed architecture and components?

LIMITATIONS:
The authors have analyzed the limitations of the proposed method.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a novel framework for semi-supervised medical segmentation that seeks to address three different scenarios simultaneously: lack of labelled data (standard semi-supervised learning -- SSL), domain shifts (unsupervised domain adaptation -- UDA) and generalization to new, unseen domains (semi-supervised domain generalization -- SemiDG). Toward this goal, they combine several strategies in their framework: sampling-based data augmentation, diffusion V-Net for learning domain invariant representations, difficulty-aware supervised training to generate class-unbiased pseudo-labels, and a reparameterize and smooth (RS) technique based on on the Gumbel softmax & ensembling strategy to generate high-quality pseudo-labels. The proposed approach is tested on four different datasets and tasks: LASeg for the SS taskL, Synapse for class imbalanced SSL, MMWHS for UDA, and M&Ms for SemiDG. Results show the method to achieve a similar or better performance on all tasks compared to SOTA methods.\n\nSTRENGTHS:\n* The proposed framework brings together three different medical image segmentation scenarios involving unlabelled data: SSL, UDA and SemiDG. This bridges a gap in the literature, as existing approaches typically focus on a single of these scenarios.\n\n* Authors do a good job in the Introduction (Figs 1-3) to situate their method with respect to existing approaches and highlight the limitations of these approaches. \n\n* Experiments to evaluate the proposed method are comprehensive. The proposed method is evaluated in four different scenarios and, fr each one, compared against the SOTA. The experimental validation also includes several relevant ablation studies that demonstrate the usefulness of the method's different components.\n\n* Results show the method to achieve SOTA performance on all tasks. These results are particularly impressive for the UDA task on the MMWHS (MR to CT) where the method achieves a performance on par with supervised training (more on this later).\n\nWEAKNESSES:\n* From my understanding, the technical novelty of the work comes from the combination of several existing techniques to address different problems in SSL. While some results are good, the proposed method feels heavily-engineered and complex.\n\n* Experiments do not fully support the main claims of the paper. For example, current experiments employ different baselines to compare the method in the SSL, UDA and SemiDG settings, instead of testing the same approaches (or some of them) across these settings. Hence, the claim that current approaches can only handle one setting is not truly validated. Likewise, the claim that having separate branches for labelled and unlabelled is not well validated in the ablation study on architectures since all tested models implement this idea (as I understand)  \n\n*  In terms of results, the method yields more limited improvements for the SSL and SemiDG settings. It is unclear if these are statistically significant.\n\nQUESTIONS:\n* p2: \"Currently, methods for these three scenarios are optimized separately, and there is no existing approach that addresses all three scenarios within a unified framework.\" As I understand, SSL and UDA are subcases of SemiDG, hence methods designed for the latter should be able to handle all three. I would relevant to compare the proposed approach against those in all three settings.\n\n* p2: \"Thus, the training process will easily get overwhelmed by the supervised training task.\" It is unclear to me why this is the case since unlabelled data is typically used to regularize the model trained with a supervised loss, and one of the main goals of regularization is to avoid overfitting. Authors should better validate this claim in their experiments.\n\n* p5: \" SVDA build upon an augmentation set and Naug operations are randomly sampled to apply to both the labelled and unlabelled data.\" This is in fact a very common strategy in SSL.\n\n* Section 3.3: Normally, diffusion models require a lot of data to train. How is it possible to achieve this using only a few labelled samples ?\n \n* Section 3.3: In the diffusion model, why do you add noise on the ground-truth segmentation maps? Assuming that the domain shift occurs in the image space (the anatomy should not vary from one modality or site to another), what can this achieve?\n\n* Figure 4: I am confused about the diffusion encoder taking both a labelled and an unlabelled image as input, since they have different shapes. Can authors explain?\n\n* Eq (2): the notation is confusing. Is 'e' the index of the sum or not? \n\n* Eq (2): Technically, in the standard definition of PSI, you should have a Delta in the sum (as it is similar to a KL)\n\n* Eq (2)-(3): As I understand, du_{k,e} is always negative while dl_{k,e} is always positive, hence w_diff should be negative?\n\n* Eq (4): Why not use a standard focal loss to address this issue?\n\n* p6: \"we found that the denoised probability map pu;ψ is too sparse, i.e., with very high confidence of each class.\" Why not use a different softmax temperature or label smoothing to avoid this ?\n\n* p7: For the class imbalanced scenario, why not used a class prior as in Bateson et al, since these priors are known from anatomical references\n\nBateson, Mathilde, Hoel Kervadec, Jose Dolz, Hervé Lombaert, and Ismail Ben Ayed. \"Source-free domain adaptation for image segmentation.\" Medical Image Analysis 82 (2022): 102617.\n\n* Table 3: the results for UDA on MMWHS are almost too good. It seems unlikely that the proposed method can achieve a Dice score similar to that of supervised trainining and a better ASD. Have authors checked for data leakage? Can they comment on this?\n\nLIMITATIONS:\nThe discussion of limitations could be expanded (only the training time of the diffusion model is mentioned). Any failure cases?", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-09"}