Momentum Contrastive Voxel-Wise Representation Learning for Semi-supervised Volumetric Medical Image Segmentation

Abstract

Contrastive learning (CL) aims to learn useful representation without relying on expert annotations in the context of medical image segmentation. Existing approaches mainly contrast a single positive vector (i.e., an augmentation of the same image) against a set of negatives within the entire remainder of the batch by simply mapping all input features into the same constant vector. Despite the impressive empirical performance, those methods have the following shortcomings: (1) it remains a formidable challenge to prevent the collapsing problems to trivial solutions; and (2) we argue that not all voxels within the same image are equally positive since there exist the dissimilar anatomical structures with the same image. In this work, we present a novel Contrastive Voxel-wise Representation Learning (CVRL) method to effectively learn low-level and high-level features by capturing 3D spatial context and rich anatomical information along both the feature and the batch dimensions. Specifically, we first introduce a novel CL strategy to ensure feature diversity promotion among the 3D representation dimensions. We train the framework through bi-level contrastive optimization (i.e., low-level and high-level) on 3D images. Experiments on two benchmark datasets and different labeled settings demonstrate the superiority of our proposed framework. More importantly, we also prove that our method inherits the benefit of hardness-aware property from the standard CL approaches.

Appendices

Appendix

Table 3. Quantitative segmentation results on the pancreas dataset. The backbone network of all evaluated methods are V-Net.

Table 4. More Ablation study on the LA dataset for voxel-wise contrastive objective $\mathcal {L}_v$ and dimensional contrastive objective $\mathcal {L}_d$ (8 labeled and 72 unlabeled). We can observe our dimensional contrastive objective can be considered as the complementary to voxel-wise contrastive objective. This demonstrates the importance of dimensional contrastive objective, which provides good performance gains.

A Hardness-aware Property of the Contrastive Losses

Recent work [26] has studies the property of standard contrastive objective (i.e., InfoNCE). For better illustration, we reproduce the contrastive objective defined in Sect. 1.2 as follows:

$$\begin{aligned} \mathcal {L}_q = -\log \frac{\exp (q \cdot k_+/\tau )}{\exp (q \cdot k_+/\tau ) + \sum _{k \in \mathcal {K}_-}\exp (q \cdot k/\tau )} \end{aligned}$$

Given query vector q, and a set of key vectors $\mathcal {K}$, the InfoNCE loss can be understood as maximizing/minimizing the positive score $\exp (q \cdot k_+/\tau )$/negative score $\exp (q \cdot k/\tau )$. Concretely, $\mathcal {L}_q = -\log P_+$ where $P_i$ is the probability of q matched to $k_i$:

$$\begin{aligned} P_i = \frac{\exp (q \cdot k_i/\tau )}{\sum _{k \in \mathcal {K}}\exp (q \cdot k/\tau )} \end{aligned}$$

Next, we show that InfoNCE has the hardness-aware property compared to a more naïve contrastive loss $\mathcal {L}'_q$ whose object is also to maximize the similarity between the query q and the positive key $k_+$:

$$\begin{aligned} \mathcal {L}'_q = -q \cdot k_+ + \lambda \sum _{k \in \mathcal {K}_-}q \cdot k \end{aligned}$$

We analyze the derivative with respect to the positive key $k_+$ and any negative key $k_-$:

$$\begin{aligned} \frac{\partial \mathcal {L}_q}{\partial k_+}&= -\frac{[\sum _{k \in \mathcal {K}}\exp (q \cdot k/\tau )] \exp (q \cdot k_+/\tau ) - [\exp (q \cdot k_+/\tau )]^2}{P_+ \, [\sum _{k \in \mathcal {K}}\exp (q \cdot k/\tau )]^2} \cdot \frac{q}{\tau },\\&= -\frac{\sum _{k \in \mathcal {K}_-}\exp (q \cdot k/\tau )}{\sum _{k \in \mathcal {K}}\exp (q \cdot k/\tau )} \cdot \frac{q}{\tau } = -\frac{q}{\tau }\sum _{k\ne +}P_k,\\ \frac{\partial \mathcal {L}_q}{\partial k_-}&= -\frac{-\exp (q \cdot k_+/\tau )\exp (q \cdot k_-/\tau )}{P_+ \, [\sum _{k \in \mathcal {K}}\exp (q \cdot k/\tau )]^2} \cdot \frac{q}{\tau } =\frac{q}{\tau }P_{-},\\ \frac{\partial \mathcal {L'}_q}{\partial k_+}&= -q, \quad \frac{\partial \mathcal {L'}_q}{\partial k_-} = \lambda q. \end{aligned}$$

We can observe that the derivative of $\mathcal {L}_q$ with respect to any negative sample $k_-$ is proportional to the exponential term $\exp (q \cdot k_{-}/\tau )$, indicating hardness-aware property. On the other hand, $\mathcal {L}'_q$ is not hardness-aware because each negative key is weighted the same (Table 4).

In Sect. 1.2, both voxel-wise and dimensional contrastive objectives use the InfoNCE loss, and thus benefit from the hardness-aware property. (1) Voxel-wise Contrastive Objective, we define queries and keys as:

$$\begin{aligned} q = f(t(x))_i,\quad k_i = g(t'(x))_i,\quad k_j = g(t'(x))_j, \end{aligned}$$

where i and j denote voxel indices. $k_j$ is a feature voxel at a different location, which is considered as a negative key. $\mathcal {L}_v$ essentially pushes the $k_j$ away from q more strongly when they are close. This encourages the representations to contain unique local information. (2) Dimensional Contrastive Objective: for brevity, here we denote i and j as dimension index. $\mathcal {L}_d$ encourages each dimension of the features to encode dissimilar information and prevents dimensional collapse [25].