SimCVD: Simple Contrastive Voxel-Wise Representation Distillation for Semi-Supervised Medical Image Segmentation

A. Semi-Supervised Medical Image Segmentation

In recent years, substantial efforts [14]–[17], [20], [42]–[47] have been devoted to incorporating unlabeled data to improve network performance due to limited annotations. Yu et al. [20] investigated an uncertainty map based on the mean-teacher framework [32] to guide the student network to capture better features. Li et al. [21] proposed to use signed distance fields for boundary prediction to improve the performance. Also, Luo et al. [23] proposed a dual-task-consistency (DTC) model for semi-supervised medical image segmentation by jointly predicting the pixel-wise segmentation maps and the global-level level set representations on the unlabeled data. Our method aims at a more practical and challenging scenario: we train our model in a more extreme few-annotation setting that relies only on a small number of annotations, while achieving superior segmentation accuracy.

B. Contrastive Learning

Self-supervised learning (SSL) [46], [48]–[50] has provided robust benefits to vision tasks by learning effective visual representations from unlabeled data in an unsupervised setting. It is based on a commonly-held belief that superior performance gains can be achieved through improved representation learning. Recently, contrastive learning, a type of self-supervised learning, has received a lot of interests [24], [26], [29], [48], [51]–[56]. The key idea of contrastive learning is to learn powerful representations that optimize similarity constraints to discriminate similar pairs (positive) and dissimilar pairs (negative) within a dataset. The primary stream of subsequent work focuses on the choice of dissimilar pairs, which is critical to the quality of learned representations. The loss function used to quantify the contrast is chosen from several options, such as InfoNCE [57], Triplet [58], and so on. Recent studies [51], [53] introduced memory bank or momentum contrast to use more negative samples for contrast computation. In the context of medical imaging, Chaitanya et al. [24] extended a contrastive learning framework to extract global and local cues in a stage-wise way, which requires human intervention and extensive training time. In contrast to Chaitanya et al. [24], our unified work focuses on explicit modeling of the intrinsic geometric structure of the semantic objects in an end-to-end manner, and hence is able to recognize object boundaries more effectively and efficiently.

C. Knowledge Distillation

The idea of knowledge distillation is to minimize the KL-divergence between the output distributions of the teacher model and the student model, and thus avoid over-fitting. KD has been applied to a variety of tasks [59]–[62], including image classification [37], [63]–[65] and semantic segmentation [39], [66]. Recent works [63], [64] found that the student model can outperform the teacher model when they share the same network architecture. Zhang et al. [67] proposed to collaboratively train multiple student models with co-distillation, which improves performance of those individual models. At the same time, in the context of medical imaging, among those existing state-of-the-art KD methods, the self-ensemble mean-teacher framework [32] is widely explored for image segmentation. Different from the existing methods that separately exploit class probabilities for each voxel, we consider knowledge distillation as a structured prediction problem by matching the relational similarity among all pairs of voxels from the encoded feature maps of the meanteacher model. We have found that our approach significantly improves learning better voxel-wise representations.

A. Overview

We aim to construct an end-to-end voxel-wise contrastive distillation algorithm to learn boundary-aware representations in the setting of extremely few annotations for volumetric medical imaging segmentation. Although the accuracy of supervised models is usually higher than that of semi-supervised models, the former requires much more labeled data than the latter. In many clinical situations, we only have few annotated data but a large amount of unlabeled data. This situation necessitates a semi-supervised segmentation algorithm that can utilize the unlabeled data to improve the segmentation performance.

To this end, we propose a novel contrastive distillation framework to advance state-of-the-art voxel-wise representation learning. In particular, our base multi-task segmentation network tackles two tasks simultaneously: classification and regression. Specifically, the segmentation network takes the input volume batch and jointly predicts the probability maps (classification) and the SDMs of the object (regression). To obtain better representations, we propose to perform structured distillation in the latent feature space, followed by contrasting the boundary-aware features in the prediction space, to learn more effective boundary-aware representations from 3D unlabeled data by regularizing the embedding space and exploring the geometric and spatial context of training voxels. At test time, we remove the mean teacher and two projection heads, and only the student network is deployed for the medical segmentation tasks.

B. Task Formulation

In this work, we consider a set of training data (3D images) including $N$ labeled data and $M$ unlabeled data, where $N\ll M$. For simplicity, we denote the small set of labeled data as ${𝒟}_l={\left\{\left({\mathbf{X}}_i,{\mathbf{Y}}_i,{\mathbf{Y}}_i^{\text{sdm}}\right)\right\}}_{i=1}^N$, and abundant unlabeled data as ${𝒟}_u={\left\{{\mathbf{X}}_i\right\}}_{i=N+1}^{N+M}$, where ${\mathbf{X}}_i\in {\mathbb{R}}^{H\times W\times D}$ is the volume input, ${\mathbf{Y}}_i\in \{\mathrm{0,1}{\}}^{H\times W\times D}$ is the ground-truth label, and ${\mathbf{Y}}_i^{\text{sdm}}\in {\mathbb{R}}^{H\times W\times D}$ is the computed ground truth SDMs from ${\mathbf{Y}}_i$, which measures the distance from each voxel to the object boundary. Every 3D image ${\mathbf{X}}_i$ consists of a set of 2D image slices ${\mathbf{X}}_i=\left[{\mathbf{x}}_{i,1},\dots ,{\mathbf{x}}_{i,D}\right]$ where ${\mathbf{x}}_{i,j}\in {\mathbb{R}}^{H\times W}$.

Our proposed SimCVD framework consists of a mean-teacher network, ${ℱ}_t\left(\mathbf{X};{\theta }_t\right)$, and a student network, ${ℱ}_s\left(\mathbf{X};{\theta }_s\right)$. Inspired by recent work [14], [32], the optimization of these two networks can be achieved with an exponential moving average (EMA) which uses a weighted combination of the parameters of the student network and the parameters of the teacher network to update the latter. This strategy has been widely shown to improve training stability and the model’s final performance. Motivated by this idea, our training strategy is divided into two steps. At each iteration, we first optimize the student network ${ℱ}_s$ by stochastic gradient descent. Then we update the teacher weights ${\theta }_t$ using an exponential moving average of the student weights ${\theta }_s$.

The inputs to the two networks are perturbed versions of the same image. That is, given a volume input ${\mathbf{X}}_i$, we first add different perturbations (i.e., affine transformation and random crop) to generate two different images ${\mathbf{X}}_i^t$ and ${\mathbf{X}}_i^s$. We then feed ${ℱ}_t$ and ${ℱ}_s$ with these two corresponding augmented images to obtain two confidence score (probability) maps ${\mathbf{Q}}_i^t$ and ${\mathbf{Q}}_i^s$.

Before we present our proposed SimCVD in detail, we first describe our base architecture below.

C. Base Architecture

Our base architecture adopts V-Net [20] as the network backbone, which consists of an encoder network $e_t:{\mathbb{R}}^{H\times W\times D}\to {\mathbb{R}}^{H^{\text{'}}\times W^{\text{'}}\times D^{\text{'}}\times D_e}$ and a decoder network $d_t:{\mathbb{R}}^{H^{\text{'}}\times W^{\text{'}}\times D^{\text{'}}\times D_e}\to [\mathrm{0,1}{]}^{H\times W\times D}\times [-\mathrm{1,1}{]}^{H\times W\times D}$ for the teacher network, and similarly $e_s,d_s$ for the student network, i.e., ${ℱ}_t=d_t\circ e_t$ and ${ℱ}_s=d_s\circ e_s.H^{\text{'}},W^{\text{'}},D^{\text{'}}$ are the size of the hidden pattern and $D_e$ is the encoded feature dimension. Inspired by previous work on medical imaging segmentation [12], [21], we incorporate multi-task learning into $ℱ$ to jointly perform both classification and regression tasks.

Given input ${\mathbf{X}}_i$, the classification branch is designed to generate the probability map ${\mathbf{Q}}_i^s\in [\mathrm{0,1}{]}^{H\times W\times D}$, and the regression branch is designed to predict the $\text{SDM}{\mathbf{Q}}_i^{s,\text{sdm}}\in$ $[-\mathrm{1,1}{]}^{H\times W\times D}$. The design of the regression branch is simple yet effective, only including the hyperbolic tangent function. This design brings two clear benefits: (1) we can eventually encode rich geometric structure information to improve segmentation accuracy, and (2) we can implicitly enforce continuity and smoothness terms for better segmentation maps. Similarly, we have outputs ${\mathbf{Q}}_i^t$ and ${\mathbf{Q}}_i^{t,\text{sdm}}$ from the teacher network.

D. Boundary-Aware Contrastive Distillation

Many prior methods distill knowledge merely in the shared prediction space by delivering the student network that matches the accuracy of the teacher network. However, this strategy is not robust for the following reasons: (1) the learned voxel-wise representations from the mean-teacher model are usually not robust due to the lack of geometric information; (2) the segmentation model still suffers from generalization issues; and (3) the network performance needs to be further improved. Therefore, we propose to perform boundary-aware contrastive distillation to train our model for better segmentation accuracy.

Our method differs from previous state-of-the-art methods in three aspects: (1) SimCVD imposes the global consistency in object boundary contours to capture more effective geometric information; (2) previous methods follow the standard setting in considering the relations of local patches, while SimCVD aims to exploit correlations among all pairs of voxels to improve robustness; and (3) due to the computational cost, SimCVD does not use a large memory bank. SimCVD trains the contrastive objective as an auxiliary loss during the volume batch updates. To specify our voxel-wise contrastive distillation algorithm on unlabeled sets, we define two discrimination terms: boundary-aware contrastive loss and pair-wise distillation loss.

A. Dataset and Pre-Processing

We evaluated our approach on two popular benchmark datasets: the Left Atrium (LA) MR dataset from the Atrial Segmentation Challenge,¹ and the NIH pancreas CT dataset [72]. For the Left Atrium dataset, it comprises 100 3D gadoliniumenhanced MR imaging scans (GE-MRIs) with expert annotations, with an isotropic resolution of 0.625 × 0.625 × 0.625mm³. Following the experimental setting in [20], we use 80 scans for training, and 20 scans for evaluation. We employ the same pre-processing methods by cropping all the scans at the heart region and normalizing the intensities to zero mean and unit variance. All the training sub-volumes are augmented by random cropping to 112 × 112 × 80mm³. For the pancreas dataset, it contains 82 contrast-enhanced abdominal CT scans. Following the experimental settings in [23], we randomly select 62 scans for training, and 20 scans for evaluation. In the pre-processing, we first truncate the intensities of the CT images into the window [−125, 275] HU [73], and then resample all the data into a fixed isotropic resolution of 1.0 × 1.0 × 1.0mm³. Finally, we crop all the scans centering at the pancreas region, and normalize the intensities to zero mean and unit variance. All the training sub-volumes are augmented by random cropping to 96 × 96 × 96mm³. In this study, we compare all the methods on LA and the pancreas dataset with respect to 20% labeled ratio. To emphasize the effectiveness of SimCVD, we further validate all the methods with respect to 10% labeled ratio on LA dataset.

B. Implementation Details

In this study, all evaluated methods are implemented in PyTorch, and trained for 6000 iterations on an NVIDIA 1080Ti GPU with a batch size of 4. For data augmentation, we use standard data augmentation techniques (i.e., random rotation, flipping, and cropping). We set the hyper-parameters $\alpha ,\lambda ,\beta ,\gamma ,\tau$ as 0.1, 0.5, 0.1, 0.1, 0.5, respectively. For the projection head, we set $p=0.1$ in the AlphaDropout layer, and output size 128 × 128 for AdaptiveAvgPool2d. We use SGD optimizer with a momentum of 0.9 and a weight decay of 0.0005 to optimize network parameters. The initial learning rate is set as 0.01 and divided by 10 every 3000 iterations. For EMA updates, we follow the experimental setting in [20], where the EMA decay rate $\alpha$ is set to 0.999. We use the time-dependent Gaussian warming-up function ${\Psi }_{\text{con}}\left(t\right)=\text{exp}\left(-5{\left(1-t/t_{\text{max}}\right)}^2\right)$ to ramp up parameters, where $t$ and $t_{\text{max}}$ denote the current and the maximum training step, respectively. For fairness, we do not adopt any post-processing step.

In the testing stage, we adopt four metrics to evaluate the segmentation performance: Dice coefficient (Dice), Jaccard Index (Jaccard), 95% Hausdorff Distance (95HD), and Average Symmetric Surface Distance (ASD). Following [20] and [23], we adopt a sliding window strategy, which uses a stride with 18 × 18 × 4 for the LA and 16 × 16 × 16 for the pancreas.

A. Experiments: Left Atrium

We compare SimCVD with published results from previous state-of-the-art semi-supervised segmentation methods, including V-Net [7], MT [32], DAN [15], CPS [68], Entropy Mini [69], UA-MT [20], ICT [70], SASSNet [21], DCT [23], and Chaitanya et al. [24] on the LA dataset in two labeled ratio settings (i.e., 10% and 20%).

The quantitative results on the LA dataset are shown in Table I. SimCVD substantially improved the segmentation accuracy in both 10% and 20% labeled cases. The results are visualized in Fig 2. Specifically, in the setting of 20% labeled ratio, our proposed SimCVD raises the previous best average results from 89.94% to 90.85% and from 81.82% to 83.80% in terms of Dice and Jaccard, even achieving comparable performance to the fully supervised baseline. Using the 10% labeled ratio, SimCVD further advances the state-of-the-art results from 87.49% to 89.03% in Dice. The gains in Jaccard, ASD, and 95HD are also substantial, achieving 80.34%, 2.59, and 8.34, respectively. This suggests that: (1) taking voxel samples with a contrastive objective yields better voxel embeddings; (2) incorporating pair-wise spatial labeling consistency can boost the performance by accessing more structural knowledge; and (3) utilizing a geometric constraint (i.e., SDM) is capable of helping identify more accurate boundaries. Leveraging all these aspects, we can observe consistent performance gains.

B. Experiments: Pancreas

To further evaluate the effectiveness of SimCVD, we compare our model on the pancreas CT dataset. Experimental results on the pancreas CT dataset are summarized in Table II. We observe that our model consistently outperforms all previous methods, achieving up to 6.72% absolute improvements in Dice. As shown in Figs. 2 and 3, our method is capable of predicting high-quality object segmentation, considering the fact that the improvement in such a setting is difficult. This demonstrates: (1) the necessity of comprehensively considering both boundary-aware contrast and pair-wise distillation; and (2) the efficacy of global shape information. Compared to the previous strong models, our approach achieves large improvements on all the datasets, demonstrating its effectiveness.

A. Analysis on Boundary-Aware Contrastive Distillation

B. Analysis on Projection Head

To further understand how different aspects of our projection head contribute to the superior model performance, we conduct extensive experiments and discuss our findings below.