SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To better leverage the domain- and class-unbiased features, we also transmit with knowledge distillation strategy:  $\theta = w_{ema} \times \theta + (1 - w_{ema}) \times (\xi + \psi) / 2$ ,  $w_{ema} = 0.99$ . The overall training function of the A&G framework is:
$$
\mathcal {L} = \mathcal {L} _ {\text {d e n o}} + \mathcal {L} _ {\text {d i f f}} + \mu \mathcal {L} _ {u} \tag {9}
$$
where  $\mu$  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};\xi)$  and  $D(x^{u};\theta)$  are used to generate the predictions.
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
Table 1: Results on Synapse dataset with  $20\%$  labeled data for class imbalanced SSL task. 'Common' or 'Imbalance' indicates whether the methods consider the imbalance issue or not. Notably, our method does not introduce many additional parameters compared with the existing SSL methods. Results of 3-times repeated experiments are reported in 'mean±std' format. Best results are boldfaced.
Methods	Avg. Dice	Avg. ASD	Dice of Each Class
Sp	RK	LK	Ga	Es	Li	St	Ao	IVC	PSV	PA	RAG	LAG
General	V-Net (fully)	62.09±1.2	10.28±3.9	84.6	77.2	73.8	73.3	38.2	94.6	68.4	72.1	71.2	58.2	48.5	17.9	29.0
UA-MT [8]	20.26±2.2	71.67±7.4	48.2	31.7	22.2	0.0	0.0	81.2	29.1	23.3	27.5	0.0	0.0	0.0	0.0
URPC [10]	25.68±5.1	72.74±15.5	66.7	38.2	56.8	0.0	0.0	85.3	33.9	33.1	14.8	0.0	5.1	0.0	0.0
CPS [3]	33.55±3.7	41.21±9.1	62.8	55.2	45.4	35.9	0.0	91.1	31.3	41.9	49.2	8.8	14.5	0.0	0.0
SS-Net [17]	35.08±2.8	50.81±6.5	62.7	67.9	60.9	34.3	0.0	89.9	20.9	61.7	44.8	0.0	8.7	4.2	0.0
DST [6]	34.47±1.6	37.69±2.9	57.7	57.2	46.4	43.7	0.0	89.0	33.9	43.3	46.9	9.0	21.0	0.0	0.0
DePL [5]	36.27±0.9	36.02±0.8	62.8	61.0	48.2	54.8	0.0	90.2	36.0	42.5	48.2	10.7	17.0	0.0	0.0
Imbalance	Adsh [71]	35.29±0.5	39.61±4.6	55.1	59.6	45.8	52.2	0.0	89.4	32.8	47.6	53.0	8.9	14.4	0.0	0.0
CReST [26]	38.33±3.4	22.85±9.0	62.1	64.7	53.8	43.8	8.1	85.9	27.2	54.4	47.7	14.4	13.0	18.7	4.6
SimiS [28]	40.07±0.6	32.98±0.5	62.3	69.4	50.7	61.4	0.0	87.0	33.0	59.0	57.2	29.2	11.8	0.0	0.0
Basak et al. [33]	33.24±0.6	43.78±2.5	57.4	53.8	48.5	46.9	0.0	87.8	28.7	42.3	45.4	6.3	15.0	0.0	0.0
CLD [15]	41.07±1.2	32.15±3.3	62.0	66.0	59.3	61.5	0.0	89.0	31.7	62.8	49.4	28.6	18.5	0.0	5.0
DHC [25]	48.61±0.9	10.71±2.6	62.8	69.5	59.2	66.0	13.2	85.2	36.9	67.9	61.5	37.0	30.9	31.4	10.6
A&D (ours)	60.88±0.7	2.52±0.4	85.2	66.9	67.0	52.7	62.9	89.6	52.1	83.0	74.9	41.8	43.4	44.8	27.2
Table 2: Results on two settings of LASeg dataset for SSL task.
5% labeled data (labeled:unlabeled=4:76)
Method	Metrics
Dice	Jaccard	95HD	ASD
V-Net (fully)	91.47	84.36	5.48	1.51
V-Net (5%)	52.55	39.60	47.05	9.87
UA-MT [8]	82.26	70.98	13.71	3.82
SASSNet [9]	81.60	69.63	16.16	3.58
DTC [11]	81.25	69.33	14.90	3.99
URPC [10]	82.48	71.35	14.65	3.65
MC-Net [12]	83.59	72.36	14.07	2.70
SS-Net [17]†	86.33	76.15	9.97	2.31
BCP [76]†	88.02	78.72	7.90	2.15
A&D (ours)	89.93	81.82	5.25	1.86
$10\%$  labeled data (labeled:unlabeled  $= 8:72$
Method	Metrics
Dice	Jaccard	95HD	ASD
V-Net (10%)	82.74	71.72	13.35	3.26
UA-MT [8]	87.79	78.39	8.68	2.12
SASSNet [9]	87.54	78.05	9.84	2.59
DTC [11]	87.51	78.17	8.23	2.36
URPC [10]	86.92	77.03	11.13	2.28
LMISA-3D [77]*	86.06	76.53	12.99	2.41
vMFNet [55]*	73.88	62.56	16.81	5.04
MC-Net [12]	87.62	78.25	10.03	1.82
SS-Net [17]†	88.55	79.62	7.49	1.90
BCP [76]†	89.62	81.31	6.81	1.76
A&D (ours)	90.31	82.40	5.55	1.64
${}^{ \dagger  }$  use test set for validation,we use labeled data instead.
* SOTA methods tailored for UDA and SemiDG, respectively.
Table 3: Results on two settings of MMWHS dataset for UDA task.
MR to CT
Method	Dice	ASD Average
AA	LAC	LVC	MYO	Average
Supervised Training	92.7	91.1	91.9	87.8	90.9	2.2
PnP-AdaNet [38]	74.0	68.9	61.9	50.8	63.9	12.8
AdaOutput [37]	65.2	76.6	54.4	43.6	59.9	9.6
CycleGAN [34]	73.8	75.7	52.3	28.7	57.6	10.8
CycADA [36]	72.9	77.0	62.4	45.3	64.4	9.4
SIFA [41]	81.3	79.5	73.8	61.6	74.1	7.0
DSFN [42]	84.7	76.9	79.1	62.4	75.8	N/A
DSAN [43]	79.9	84.8	82.8	66.5	78.5	5.9
LMISA-3D [77]	84.5	82.8	88.6	70.1	81.5	2.3
A&D (ours)	93.2	89.5	91.7	86.2	90.1	1.7
CT to MR
Method	Dice	ASD Average
AA	LAC	LVC	MYO	Average
Supervised Training	82.8	80.5	92.4	78.8	83.6	2.9
PnP-AdaNet [38]	43.7	68.9	61.9	50.8	63.9	8.9
AdaOutput [37]	60.8	39.8	71.5	35.5	51.9	5.7
CycleGAN [34]	64.3	30.7	65.0	43.0	50.7	6.6
CyCADA [36]	60.5	44.0	77.6	47.9	57.5	7.9
SIFA [41]	65.3	62.3	78.9	47.3	63.4	5.7
DSAN [43]	71.3	66.2	76.2	52.1	66.5	5.4
LMISA-3D [77]	60.7	72.4	86.2	64.1	70.8	3.6
SS-Net [17]*	62.1	58.4	68.9	51.4	60.2	5.9
BCP [76]*	63.6	63.7	70.9	58.0	64.1	4.5
A&D (ours)	62.8	87.4	61.3	74.1	71.4	7.9
* SOTA methods on semi-supervised segmentation.
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
Table 4: Results on two settings of M&Ms dataset for SemiDG task.
Method	2% Labeled data	5% Labeled data
Domain A	Domain B	Domain C	Domain D	Average	Domain A	Domain B	Domain C	Domain D	Average
nnUNet [69]	52.87	64.63	72.97	73.27	65.94	65.30	79.73	78.06	81.25	76.09
SDNet+Aug [53]	54.48	67.81	76.46	74.35	68.28	71.21	77.31	81.40	79.95	77.47
LDDG [78]	59.47	56.16	68.21	68.56	63.16	66.22	69.49	73.40	75.66	71.29
SAML [79]	56.31	56.32	75.70	69.94	64.57	67.11	76.35	77.43	78.64	74.88
BCP [76]*	71.57	76.20	76.87	77.94	75.65	73.66	79.04	77.01	78.49	77.05
DGNet [52]	66.01	72.72	77.54	75.14	72.85	72.40	80.30	82.51	83.77	79.75
vMFNet [55]	73.13	77.01	81.57	82.02	78.43	77.06	82.29	84.01	85.13	82.12
A&D (ours)	79.62	82.26	80.03	83.31	81.31	81.71	85.44	82.18	83.9	83.31
* SOTA method on semi-supervised segmentation.
(a)
Figure 5: Ablation study on different architectures. (e) is the final framework
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.
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
Table 5: Ablation study on the MMWHS MR to CT setting of different architectures with respect to those in Figure 5.
Arch.	a	b	c	d	e
Dice	76.14	50.53	82.19	45.63	90.14
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), MMWHS MR to CT setting (UDA) and  $2\%$  labeled  $2\%$  labeled M&Ms setting (SemiDG).
Methods	A&D	w/o SVDA	w/o Diffusion	w/o DRS	w/o RS
IBSSL	60.9	55.3	56.7	52.0	58.7
UDA	90.1	84.6	79.2	85.8	87.0
SemiDG	80.6	77.9	74.9	78.2	78.6
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.
Time step	100	500	1000	Aug #	2	3	4
Dice	69.7	70.2	71.4	Dice	68.9	71.4	70.2
Figure 6: Visualization of the UDA task in 2D and 3D views.
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.
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.
We thank the anonymous NeurIPS reviewers for providing us with valuable feedback that greatly improved the quality of this paper! This work was supported in part by the Hong Kong Innovation and Technology Fund under Project ITS/030/21 and in part by Foshan HKUST Projects under Grants FSUST21-HKUST10E and FSUST21-HKUST11E and in part by the Project of Hetao Shenzhen-Hong Kong Science and Technology Innovation Cooperation Zone (HZQB-KCZYB-2020083).
Figure 7: Visualization of the class imbalance SSL task in 2D and 3D views. The yellow arrows denote minority classes detected.
A.1 More Details of Datasets and Implementation
The hyper-parameters for different datasets are shown in Table 8.
Table 8: Hyper-parameters for different datasets.
Datasets	patch size	learning rate	batch size	feature size F
LASeg	112 × 112 × 80	1e-2	4	32
Synapse	64 × 128 × 128	3e-2	4	32
MMWHS	128 × 128 × 128	5e-3	2	32
M&Ms	32 × 128 × 128	1e-2	4	32
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.
MMWHS Dataset 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.
Table 9: Comparison of computational costs of state-of-the-art methods on different settings with the performances in terms of Dice score.
Methods	Param.(M)	Inference time (s/iter)	SSL	UDA
SS-Net [17]	75.672	21.20	88.55	78.2
CLD [15]	75.554	39.58	85.37	75.4
vMFNet [55]	20.423	4.89	73.88	72.3
EPL [54]	80.939	5.52	76.49	71.9
A&D (ours)	57.894	3.69	90.31	90.1
Figure 8: Visualization of the RS process of the foreground class on LASeg dataset. The probability map  $p^{u;\psi}$  of the difficulty-aware decoder may have low confidence in the inner region (red box), whereas the probability map  $p^{u;\xi}$  of the diffusion decoder may have inaccurate boundaries but with very high confidence (red arrows).
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\mathrm{M}$  parameters, their performances  $(73.88\%$  and  $72.3\%)$  on SSL and UDA tasks are significantly lower than ours  $(89.40\%$  and  $90.0\%)$ .
Figure 9: Visualization of the RS process of the myocardium of the left ventricle (MYO) class which is the class with worst performance on MR to CT setting of MMWHS dataset. In this case, the probability map  $p^{u;\xi}$  of the diffusion decoder contains more error regions due to the ambiguous boundaries and noise but also with very high confidence.
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;\xi}$  of diffusion with DDIM  $D(\xi)$  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;\psi}$  generated by the V-Net decoder  $D(\psi)$ , the error regions (e.g., upper right corner) with high confidence will surpass some correct regions of  $p^{u;\psi}$  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};\theta)$  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 MMWHS dataset. 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.