Bootstrapping Semi-supervised Medical
Image Segmentation with Anatomical-aware Contrastive Distillation

Framework Overview   The workflow of our proposed ACTION is illustrated in Figure 2. By default, ACTION is built on the BYOL pipeline [7] which is originally designed for image classification tasks, and for a fair comparison, we also follow the setting in [3] such as using 2D U-Net [22] as the backbone and non-linear projection heads $H$. The main differences between our proposed ACTION and [3, 9] are as follows: (1) the addition of a predictor $g(\cdot)$ to the student network to avoid collapsed solutions; (2) the utilization of a slow-moving average of the student network as the teacher network for more semantically compact representations; (3) the use of the output probability rather than logits effectively and semantically constrains the distance between the anatomical features from the imbalanced data (i.e., multi-class label imbalance cases); (4) we propose to contrast the query image features with other random image features at the global and local level, rather than only two augmented versions of the same image features; and (5) we design a novel unsupervised anatomical contrastive loss to provide additional supervision on hard pixels.

Global Contrastive Distillation Pre-Training   We follow a similar setting in [24]. Given an input query image $q\in\{x_{i}\}_{i=N+1}^{N+M}$ with the spatial size $h\times w$, we first apply two different augmentations to obtain $q_{t}$ and $q_{s}$, and randomly sample a set of augmented images $\{x_{j}\}_{j=1}^{n}$ from a set of unlabeled image slices $\{x_{i}\}_{i=N+1}^{N+M}$. We believe that such relaxation enables the model to capture more rich semantic relationships and anatomical features from its neighboring images instead of only learning from the different version of the same query image. We then feed $\{x_{j}\}_{j=1}^{n}$ to the teacher encoder $E_{t}$, and followed by the nonlinear projection head $H_{t}^{g}$ to generate their projection embeddings $\{H_{t}^{g}(E_{t}(x_{j}))\}_{j=1}^{n}$ as anchor points, and also feed $q_{t}$ and $q_{s}$ to the teacher and student (i.e., $E$ and $H$), creating $z_{t}=H_{t}^{g}(E_{t}(q_{t}))$ and $z_{s}=H_{s}^{g}(E_{s}(q_{s}))$. Here we utilize the probabilities after SoftMax instead of the feature embedding:

where $\tau_{t}$ is a temperature hyperparameter of the teacher, and $\text{sim}(\cdot,\cdot)$ is the cosine similarity. Then inspired by [7], in order to avoid collapsed solutions in an unsupervised scenario, we use a shallow multi-layer perceptron (MLP) predictor $H_{p}^{g}(\cdot)$ to obtain the prediction $z_{s}^{\ast}=H_{p}^{g}(z_{s})$. Of note, $\{a_{i}\}_{i=1}^{n}$, $z_{t}$, $z_{s}$, $z_{s}^{\ast}$ can be generated embedding from a set of randomly chosen augmented images, teacher’s projection embeddings, student’s projection embeddings, and student’s prediction embeddings in either Stage-$i$ or $ii$. Therefore, we can calculate the similarity distance between the student’s prediction and the anchor embeddings by converting them to probability distribution.

where $\tau_{s}$ refers to a temperature hyperparameter of the student. The unsupervised contrastive loss is computed as follows:

Local Contrastive Distillation Pre-Training   After training the teacher’s and student’s encoder to learn global-level image features, we attach the decoders and tune the entire models to perform pixel-level contrastive learning in a semi-supervised manner. The distinction in the training strategy between ours and [9] lies in Stage-$ii$ and $iii$: [9] only use labeled data in training, while we use both labeled and unlabeled data in training. Considering the training procedure of Stage-$ii$ is similar to Stage-$iii$, we briefly describe it here as illustrated in Figure 2. For the labeled data, we train our model by minimizing the supervised loss (the linear combination of cross-entropy loss and dice loss) in Stage-$ii$ and Stage-$iii$. As for the unlabeled input images $q$ and $\{x_{j}\}_{j=1}^{n}$, we first apply two different augmentations to $q$, creating two different versions $[q_{t}^{l},q_{s}^{l}]$, and then feed them to $F_{t}$ and $F_{s}$, and their output features $[f_{t},f_{s}]$ are fed into $H_{t}^{l}$ and $H_{t}^{l}$. The student’s projection embedding is subsequently fed into $H_{p}^{l}$ to obtain the student’s prediction embedding to enforce the similarity between the teacher and the student under the same loss as Equation 3. We also include the randomly selected images to enforce such similarity because intuitively, it may be beneficial to ensure diversity in the set of sampled examples. It is important to note that ACTION will re-use the well-trained weight of the models $F_{t}$ and $F_{s}$ as initialization for Stage-$iii$.

Anatomical Contrast Fine-Tuning   Broadly speaking, in medical images, the same tissue types may share similar anatomical information in different patients, but different tissue types often show different class, appearance, and spatial distributions, which can be described as a complicated form of imbalance and uncertainty in real clinical data, as shown in Figure 1. This motivates us to efficiently incorporate more useful features so the representations can be more balanced and better discriminated in such multi-class label imbalanced scenarios. Inspired by [15], we propose AnCo, a new unsupervised contrastive loss designed at the anatomical level. Specifically, we additionally attach a representation decoder head $H_{r}$ to the student network, parallel to the segmentation head, to decode the multi-layer hidden features by first using multiple up-sampling layers for outputting dense features with the same spatial resolution as the query image and then mapping them into high $m$-dimensional query, positive key, and negative key embeddings: $r_{q},r_{k}^{+},r_{k}^{-}$. The AnCo loss is then defined as:

where $\mathcal{C}$ is a set of all available classes in a mini-batch, and $\tau_{an}$ denotes a temperature hyperparameter for AnCo loss. $\mathcal{R}_{q}^{c}$ and $r_{k}^{c,+}$ are a set of query embeddings in class $c$ and the positive key embedding, which is the mean representation of class $c$, respectively. $\mathcal{R}_{k}^{c}$ is a set of negative key embeddings which are not in class $c$. Suppose $\mathcal{P}$ is a set including all pixel coordinates with the same resolution with $x_{i}$, these queries and keys are then defined as:

In addition, we note that contrastive learning usually benefits from a large collection of positive and negative pairs, but it is usually bounded by the size of GPU memory. Therefore, we introduce two novel active hard sampling methods. To address the uncertainty on the most challenging pixels among all available classes (i.e., close anatomical or semantic relationship), we non-uniformly sample negative keys based on relative similarity distance between the query class and each negative key class. For each mini-batch, we build a graph $G$ to measure the pair-wise class relationship to dynamically update $G$.

where $G\in\mathbb{R}^{|\mathcal{C}|\times|\mathcal{C}|}$. Note that this process may be hard to allocate more samples. Thus, to learn a more accurate decision boundary, we first apply SoftMax function by normalizing the pair-wise relationships among all negative classes $n$ from each query class $c$, yielding a distribution: $\exp(G[c,v])/\sum_{n\in\mathcal{C},n\neq c}\exp(G[c,n])$. Then we adaptively sample negative keys from each class $v$ to help learn the corresponding query class $c$. To alleviate the imbalance issue, we sample hard queries based on a defined threshold, to better discriminate the rare classes. The easy and hard queries are computed as follows:

where $\hat{y}_{q}$ is the predicted confidence of label $c$ corresponding to $r_{q}$ after SoftMax function, and $\theta_{s}$ is the user-defined confidence threshold.

Experimental Setup   We experiment on two benchmark datasets: ACDC 2017 dataset [1] and MICCAI 2017 Liver Tumor Segmentation Challenge (LiTS) [2].

The ACDC dataset includes 200 cardiac cine MRI scans from 100 patients with annotations including three segmentation classes (i.e., left ventricle (LV), myocardium (Myo), and right ventricle (RV)). Following [16, 27], we use 140, 20, and 60 scans for training, validation, and testing, respectively.

The LiTS dataset includes 131 contrast-enhanced 3D abdominal CT volumes with annotations of two segmentation classes (i.e., liver and tumor). Following [13], we use the first 100 volumes for training, and the rest 31 for testing. For pre-processing, we follow the setting in [3] to normalize the intensity of each 3D scans, resample all 2D slices and the corresponding segmentation maps to a fixed spatial resolution (i.e., 256×256 pixels). To quantitatively assess the performance of our proposed method, we report two popular metrics: Dice coefficient (DSC) and Average Surface Distance (ASD) for 3D segmentation results.

Implementation Details   All our models are implemented in PyTorch [19]. We train all methods with SGD optimizer (learning rate=$0.01$, momentum=$0.9$, weight decay=$0.0001$, batch size=$6$). All models are trained with two NVIDIA GeForce RTX 3090 GPUs. Stage-$i$ and $ii$ are trained with 100 epochs, and Stage-$iii$ is with 200 epochs. We use the temperature of teacher and student as $\tau_{t}\!=\!0.01$ and $\tau_{s}\!=\!0.1$. The teacher is updated using the following rule $\theta_{t}\leftarrow m\theta_{t}+(1-m)\theta_{s}$, where $\theta$ refers to the model’s parameters and the the momentum hyperparameter $m$ is $0.99$. The memory bank size is 36. We follow the standard augmentation strategies in [7]. In Stage-$i$, we train $E_{s}$, $E_{t}$, $H_{t}^{g}$, $H_{s}^{g}$, and $H_{p}^{g}$ on the unlabeled data with global-level $\mathcal{L}_{\text{contrast}}$ in Equation 3. We follow [9] to use a MLP as heads, and the setting of the predictors is similar to [7], which has a feature dimension of $512$. In Stage-$ii$, we train $F_{s}$, $F_{t}$, $H_{t}^{l}$, $H_{s}^{l}$, and $H_{p}^{l}$ on the labeled and unlabeled data. We train with the supervised loss [36] on labeled data, and local-level $\mathcal{L}_{\text{contrast}}$ in Equation 3 on unlabeled data. Given the logits output $\hat{y}\in\mathbb{R}^{C\times h\times w}$, we use the $1\times 1$ convolutional layer to project all pixels into the latent space with the feature dimension of $512$, and the output feature dimension of $G$ is also $512$. As for Stage-$iii$, we train $F_{s}$, $F_{t}$, $H_{t}$, $H_{s}$, and $H_{r}$ on the labeled and unlabeled data. We use the supervised segmentation loss on labeled data, unsupervised cross-entropy loss (on pseudo-labels generated by a confidence threshold $\theta_{s}$), and $\mathcal{L}_{\text{anco}}$ in Equation 4 on unlabeled data. We then adaptively sample 256 query samples and 512 key samples for each mini-batch, and temperature for the student and confidence thresholds are set to $\tau_{s}=0.5$ and $\theta_{s}=0.97$, respectively. Of note, the projection heads, the predictor, and the representation decoder head are only utilized during the training, and will be removed during the inference.

Main Results   We compare our proposed method to previous state-of-the-art SSL methods using 2D Unet [22] as backbone, including UNet trained with full/limited supervisions (UNet-F/UNet-L), EM [26], CCT [18], DAN [37], URPC [17], DCT [21], ICT [25], MT [23], UAMT [36], CPS [6], SCS [9], and GCL [3]. Table 1 shows the evaluation results on ACDC dataset under two unlabeled settings (3 or 7 labeled cases). ACTION can substantially improve results on two unlabeled settings, greatly outperforming the previous state-of-the-art SSL methods. Specifically, our ACTION, trained on 3 labeled cases, dramatically improves the previous best averaged Dice score from 73.6% to 87.5% by a large margin, and even matches previous SSL methods using 7 labeled cases. When using 7 labeled cases, ACTION further pushes the state-of-the-art results to 89.7% in Dice. We observe that the gains are more pronounced on the two categories(i.e., RV and Myo), and our ACTION achieves 89.8% and 86.7% in terms of Dice, performing competitive or even better than the supervised baseline (89.2% and 86.7%). As shown in Figure 3, we can see the clear advantage of ACTION, where the boundaries of different regions are clearly sharper and more accurate such as RV and Myo regions. Table 2 also shows the evaluation results on LiTS dataset under two unlabeled settings (5% or 10% labeled cases). On both two labeled settings, ACTION significantly outperforms all the state-of-the-art methods by a significant margin. As shown in Figure 4, ACTION achieves consistently sharp and accurate object boundaries compared to other SSL methods.

Ablation on Different Components   We investigate the impact of different components in ACTION. All reported results in this section are based on the ACDC dataset under the 3 labeled setting. Table 3 shows the ablation result of our model. Upon our choice of architecture, we first consider a naïve baseline (BYOL) without any random sampled images (RSI), stage-$ii$, and stage-$iii$, denoted by (1) Vanilla. Then, we consider a wide range of different settings for improved representation learning: (2) incorporating other random sampled images; (3) no stage-$ii$; (4) no other random sampled images and stage-$ii$; (5) no stage-$iii$; since stage-$iii$ includes two losses, (6) no $\mathcal{L}_{\mathrm{anco}}$, (7) no $\mathcal{L}_{\mathrm{unsup}}$, and (8) our proposed ACTION. As shown in Table 3, it is notable that ACTION performs generally better than other evaluated baselines. We find that only applying any single component of ACTION often comes at the cost of performance degradation. The intuitions behind are as follows: (1) incorporating other random sampled images will enforce the diversity of the sampled data, preventing redundant anatomically and semantically similar samples; (2) using stage-$ii$ leads to worse performance without considering local context; (3) using stage-$iii$ enables a robust segmentation model to learn better representations with few human annotations. Using the above components confers a significant advantage at representation learning, and further illustrates the benefit of each component.

Ablation on Different Augmentations   We investigate the impact of using weak or strong augmentations for ACTION on the ACDC dataset under 3 labeled setting. We summarize the effects of different data augmentation strategies in Table 4. We apply weak augmentation to the teacher’s input, including rotation, cropping, flipping, and strong augmentation to the student’s input, including rotation, cropping, flipping, random contrast, and brightness changes [20]. Empirically, we find that when using weak and strong augmentation strategies on the teacher and student network, the network performance is optimal.