SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Haonan Wang $^{1}$ , Xiaomeng Li $^{1,2*}$
$^{2}$ HKUST Shenzhen-Hong Kong Collaborative Innovation Research Institute, Futian, Shenzhen hwanggr@connect.ust.hk, eexmli@ust.hk
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.
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
Figure 1: From (a), (b) to (c): generalizing SSL to UDA and SemiDG settings by sampling more diverse data to form the training and testing sets.
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 SSSVMIS 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 SSSVMIS 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 SSSVMIS 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.
Used for prediction
Labeled data flow
Unlabeled data flow
(a) Mean Teacher
(b) CPS
UA-MT (Teacher-Student-based)
Figure 2: Our proposed A&D framework differs in that the labeled and unlabeled data flows are separate, and only the unlabeled flow is used for predicting.
CLD (CPS-based)
Figure 3: The training loss curves of UAMT and CLD on MMWHS dataset. The overall losses (black) are dominated by the supervised losses.
Based on the similarities and the main weaknesses of the mainstream SSSVMIS 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 SSSVMIS 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.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 SSL-based 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 settings, 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.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.
Figure 4: Overview of the proposed Aggregating & Decoupling framework. Blue and orange regions denote the training process with labeled data and unlabeled data, respectively. We separate the training of the decoders using labeled data and unlabeled data, and only use the decoder trained with unlabeled data for prediction.
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_i^l,y_i)\}_{i = 1}^{N_L}$  and  $N_U$  unlabeled samples  $\{x_i^u\}_{i = 1}^{N_U}$ , where  $x_{i}\in \mathbb{R}^{D\times H\times W}$  is the input volume and  $y_{i}\in \mathbb{R}^{K\times D\times H\times W}$  is the ground-truth 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;\xi}$  difficulty-aware labeled data flow  $h^{l;\psi}$ , 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_{\text{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_{\text{aug}}$  is empirically set to 3.
Algorithm 1: Training Pipeline of A&D.
Input: Labeled samples  $\{(x_i^l,y_i)\}_{i = 1}^{NL}$  , unlabeled samples  $\{x_i^u\}_{i = 1}^{NU}$    
Output: Diffusion encoder  $E(x^{l},x^{u};\xi)$  with V-Net decoder  $D(x^{u};\theta)$  for inference   
1 Initialization;   
2 for batched data  $(x_{i}^{l},y_{i})$ $x_{i}^{u}$  do   
3 Add SVDA on  $x^l$  and  $x^{u}$  to obtain  $\hat{x}^l$  and  $\hat{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};\xi)$  and  $D(x^{l};\xi)$  with  $(\hat{x}_i^l,y_t)$  by Eq. 2;   
6 Train the diffusion encoder  $E(x^{l};\xi)$  and difficulty-aware decoder  $D(x^{l};\psi)$  by Eq. 6;   
7 Generate domain-unbiased probability map  $p^{u;\xi}$  with  $E(x^{l},x^{u};\xi) + D(x^{l};\xi)$  by DDIM [57];   
8 Generate class-unbiased probability map  $p^{u;\psi}$  with  $D(x^{l};\psi)$  with forward pass;   
9 Ensemble the two maps and obtain the pseudo label  $y^{\xi ,\psi}$  by Eq. 7;   
10 Train  $E(x^{u};\xi)$  and V-Net decoder  $D(x^{u};\theta)$  with unlabeled data  $(\hat{x}^u,y^{\xi ,\psi})$  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} \in \mathbb{R}^{D \times W \times H}$  and its label  $y \in \mathbb{R}^{D \times W \times H}$ , we first convert the label to the one-hot format  $y_{0} \in \mathbb{R}^{K \times D \times W \times H}$  and add successive  $t$  step noise  $\epsilon$  to obtain the noisy label  $y_{t} \in \mathbb{R}^{K \times D \times W \times H}$ , which is the diffusion forward process:
$$
y _ {t} = \sqrt {\bar {\alpha} _ {t}} y _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \epsilon , \epsilon \in \mathcal {N} (0, 1) \tag {1}
$$
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};\xi)$  as decoder, the diffusion encoder takes concatenation  $\text{concat}([y_t,x^l])$  and time step  $t$  as input to generate the time-step-embedded multi-scale features  $h_i^{l;\xi} \in \mathbb{R}^{i\times F\times \frac{D}{2^i}\times \frac{W}{2^i}\times \frac{H}{2^i}}$  where  $i$  is the stage and  $F$  is the basic feature size.  $h_i^{l;\xi}$  are further used by  $D(x^{l};\xi)$  to predict the clear label  $y_0$ . For the difficulty-aware training flow and the unlabeled data flow, i.e.,  $D(x^{l};\psi)$  and  $D(x^{u};\theta)$  as decoders, the encoder only takes  $x_{l}$  and  $x_{u}$  as input to obtain the multi-scale features  $h_i^{l;\psi}, h_i^u$ , respectively. Note that  $h_i^{l;\psi}, h_i^u$  are with same shapes with  $h_i^{l;\xi}$ .
3.3 Decoupling Stage
The decoupling stage consists of four steps: supervised denoising training to generate domain-unbiased 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};\xi)$  Taking  $h_i^{l;\xi}$  as inputs,  $D(x^{l};\xi)$  decodes the features to predict the clear label  $y_{0}$  as domain-unbiased pseudo label. The objective function is defined as follow:
$$
\mathcal {L} _ {d e n o} = \frac {1}{N _ {L}} \sum_ {i = 0} ^ {N _ {L}} \mathcal {L} _ {D i c e C E} \left(p _ {i} ^ {l; \xi}, y _ {i}\right) \tag {2}
$$
where  $\mathcal{L}_{DiceCE}(x,y) = \frac{1}{2} [\mathcal{L}_{CE}(x,y) + \mathcal{L}_{Dice}(x,y)]$  is the combined dice and cross entropy loss.
Supervised Difficulty-aware Training with  $D(x^{l};\psi)$  To alleviate the common class imbalance issue in SSSMIS, 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;\xi}$  produced by diffusion decoder  $D(x^{l};\xi)$ : learning speed and performance. We use Population Stability Index [70] to measure the
learning speed of each class after the  $e^{th}$  iteration:
$$
d u _ {k, e} = \sum_ {e - \tau} ^ {e} \min  (\triangle , 0) \ln \left(\frac {\lambda_ {k , e}}{\lambda_ {k , e - 1}}\right), \quad d l _ {k, e} = \sum_ {e - \tau} ^ {e} \max  (\triangle , 0) \ln \left(\frac {\lambda_ {k , e}}{\lambda_ {k , e - 1}}\right) \tag {3}
$$
where  $\lambda_{k}$  denotes the Dice score of  $p^{l;\xi}$  of  $k^{th}$  class in  $e^{th}$  iteration and  $\triangle = \lambda_{k,e} - \lambda_{k,e-1}$ .  $du_{k,e}$  and  $dl_{k,e}$  denote classes not learned and learned after  $e^{th}$  iteration.  $\tau$  is the number accumulation iterations and set to 50 empirically. Then, we define the difficulty of  $k^{th}$  class after  $e^{th}$  iteration as:
$$
d _ {k, e} = \frac {d u _ {k , e}}{d l _ {k , e}} \tag {4}
$$
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 - \lambda_{k,e}$  for  $\tau$  iterations to obtain the reversed dice weight  $w_{\lambda_{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:
$$
w _ {k} ^ {\text {d i f f}} = w _ {\lambda_ {k, e}} \cdot \left(d _ {k, e}\right) ^ {\alpha} \tag {5}
$$
where  $\alpha$  is empirically set to  $\frac{1}{5}$  in the experiments to alleviate outliers. The objective function of the supervised difficulty-aware training is defined as follow:
$$
\mathcal {L} _ {\text {d i f f}} = \frac {1}{N _ {L}} \frac {1}{K} \sum_ {i = 0} ^ {N _ {L}} \sum_ {k = 0} ^ {K} w _ {k} ^ {\text {d i f f}} \mathcal {L} _ {\text {D i c e C E}} \left(p _ {i, k} ^ {l; \psi}, y _ {i, k}\right) \tag {6}
$$
Pseudo Labeling with Reparameterize & Smooth (RS) Strategy The domain-unbiased  $p^{u;\xi}$  probability map is generated by iterating the diffusion model  $(E(x^l, x^u; \xi) + D(x^l; \xi)) t$  times with the Denoising Diffusion Implicit Models (DDIM) method [57]. The class-unbiased  $p^{u;\psi}$  probability map can be obtained by  $D(x^l; \psi)$  with stopped gradient forward pass. We ensemble  $p^{u;\xi}$  and  $p^{u;\psi}$  to generate high-quality pseudo labels. However, when combining these two maps, we found that the denoised probability map  $p^{u;\psi}$  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;\psi}$  and is not robust to noise and error. Thus, we re-parameterize  $p^{u;\psi}$  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:
$$
y ^ {\xi , \psi} = \operatorname {a r g m a x} (\text {G u m b e l - S o f m a x} \left(p ^ {u; \xi}\right) + \text {S o f t m a x} \left(p ^ {u; \psi}\right)) \tag {7}
$$
Unsupervised Training with  $D(x^{u};\theta)$  Finally, we can use the pseudo label  $y^{\xi ,\psi}$  to train  $D(x^{u};\theta)$  in an unsupervised manner. The objective function of the unsupervised training is defined as:
$$
\mathcal {L} _ {u} = \frac {1}{N _ {U}} \sum_ {i = 0} ^ {N _ {U}} \mathcal {L} _ {\text {D i c e C E}} \left(p _ {i} ^ {u; \theta}, y ^ {\xi , \psi}\right) \tag {8}
$$