Keep Your Friends Close & Enemies Farther:
Debiasing Contrastive Learning with Spatial Priors in 3D Radiology Images ^††thanks: 978-1-6654-6819-0/22/$31.00 ©2022 IEEE

Thus, given a patch $P_{i}$ from image $I_{i}$, we find its corresponding crop in $I_{t}$ with $P_{t}$=$T_{i}(P_{i})$. Further, $P_{i}$’s corresponding crop in any other image $I_{j}\in D$ can be extracted via $P_{j}$=$T_{j}^{-1}\circ T_{i}(P_{i})$. We concisely denote the corresponding patch of $P_{i}$ in image $I_{j}$ as $P_{i\rightarrow j}$. Note that we do not transform the images, rather, we only compute $\mathbb{T}$.
Contrastive Representation Learning. We denote the $L_{2}$-normalized unit embedding of an anchor patch as $v$. Borrowing notation from [3], we define a positive sample as $v^{+}\in\Lambda^{+}$, where $\Lambda^{+}$ is the set of positive embedding vectors for anchor $v$. Similarly, a negative sample is defined as $v^{-}\in\Lambda^{-}$, where $\Lambda^{-}$ is the set of negative embedding vectors. The contrastive loss (NCE [16] for short) is then defined as:

where “$\mathcal{L}^{NCE}$” is shorthand for the loss $\mathcal{L}^{NCE}(v,v^{+})$ between an anchor $v$ and a positive $v^{+}$, $\tau$ is a temperature hyperparameter, and $g$ & $l$ indicate global & local components respectively. Incorporating all positive samples, the complete contrastive loss (CON for short) is written as:

Spade pretrains a randomly initialized UNet-like network containing an encoder $f_{\theta}$, decoder $g_{\theta}$, global projection module $h_{\theta_{g}}$, local projection module $h_{\theta_{l}}$, and their corresponding momentum counterparts $f_{\epsilon}$, $g_{\epsilon}$, $h_{\epsilon_{g}}$, and $h_{\epsilon_{l}}$, respectively (see Fig. 1). For each batch, an anchor image $I_{a}$ and $p$ pairs of anchor patches $\{P_{a}^{j}$, $P_{a}^{k},\ldots\}$ are sampled where overlaps of paired patches are at least $o$% and $P\in\mathbb{R}^{1\times D_{in}\times H_{in}\times W_{in}}$. Following MoCo [5], each $P$ is randomly transformed twice via spatial augmentations $\mathcal{T}_{S}$ (e.g., flipping, rotating) and intensity noising $\mathcal{T}_{I}$ (e.g., those proposed in [24]). All these views are inputted into both the regular & momentum networks.

Spade contains three pretext tasks (i.e., global contrastive learning, local contrastive learning, and reconstruction). These tasks use global logits $z$=$f(P)$, and local logits $Z$=$g\circ f(P)$, where $f$ may refer to either $f_{\theta}$ or $f_{\epsilon}$ for brevity. For global & local contrastive learning, $v$ is a unit feature embedding obtained from the global projection module $v_{g}$=$h_{g}(z)$, $v_{g}\in\mathbb{R}^{C_{g}^{emb}}$ or the local projection module $v_{l}$=$h_{l}(Z)$, $v_{l}\in\mathbb{R}^{C_{l}^{emb}\times H_{l}^{emb}\times W_{l}^{emb}}$. Note that $v_{g}$ is a vector while $v_{l}$ has dimensionality $H_{l}^{emb}\times W_{l}^{emb}$ since we aim to preserve spatial equivariance for local representations. After each iteration, global & local embeddings from the momentum network are then enqueued into $M_{g}$ (size $\mathcal{Q}_{g}$) & $M_{l}$ (size $\mathcal{Q}_{l}$), respectively, along with their corresponding positions in the template image.

For global strategy 1 (G1), we augment positives with crops of corresponding positions from different images. This is similar in spirit to [3][20], but we compute correspondences in 3D instead of assuming alignment and sample crops instead of slices which are both richer in spatial context & more robust to alignment errors from a single axis. Concretely, $\Lambda^{-}_{\textit{G1}}$=$M_{g}$, $\Lambda^{+}_{\textit{G1}}$=$\{v_{a},v_{a\rightarrow i_{1}},v_{a\rightarrow i_{2}},\ldots,v_{a\rightarrow i_{n^{+}}}\}$, where $v_{a\rightarrow i}$=$h_{g}\circ f(P_{a\rightarrow i})$ denotes the feature embedding of $P_{a\rightarrow i}$, and $n^{+}$ is the number of sampled positives in separate volumes that spatially correspond to the anchor patch $P_{a}$.

For a large embedding queue, quantity has a quality of its own, but naively raising this capacity increases incidences of false negatives. In global strategy 2 (G2), we debias the negative cohort by removing crop embeddings in $M_{g}$ that are “close” to the anchor patch. We quantify “close” as the overlap between two patches in the template space. Given two patches $P_{a}^{j}$ & $P_{a}^{k}$, we denote their overlap region as patch $P_{a}^{j\cap k}$. Their overlap ratio or IoU is defined as $IoU(P_{a}^{j},P_{a}^{k})$=$V(P_{a}^{j\cap k})/[V(P_{a}^{j})+V(P_{a}^{k})-V(P_{a}^{j\cap k})]$, where $V()$ returns the volume of a patch. Given an overlap threshold $o\in[0,1]$, we define the cohorts of positives & negatives as $\Lambda^{-}_{\textit{G2}}$=$\{v_{i}\ |\ IoU(v_{i\rightarrow t},v_{a\rightarrow t})\leq o,\forall v_{i}\in M_{g}\}$ & $\Lambda^{+}_{\textit{G2}}$=$\{v_{a},v_{a\rightarrow i_{1}},v_{a\rightarrow i_{2}},\ldots,v_{a\rightarrow i_{n^{+}}}\}$, where $IoU(v_{i\rightarrow t},v_{a\rightarrow t})$ is notational shorthand for $IoU(P_{i\rightarrow t},P_{a\rightarrow t})$.

Although false negatives are alleviated, $\Lambda^{+}_{\textit{G1}}$ and $\Lambda^{+}_{\textit{G2}}$ still suffer from the same drawbacks as [20][3]. Reliance on intra-batch samples greatly limits the variability of positives just as it had for negatives. In light of this, we propose global strategy 3 (G3) which converts the removed negative patches in G2 to positives based on the reasoning that if the semantics of a patch are similar enough to be considered a false negative, then the mutual information present would enrich the representations with added positives. We define $\Lambda^{-}_{\textit{G3}}$=$\{v_{i}\ |\ IoU(v_{i\rightarrow t},v_{a\rightarrow t})\leq o,\forall v_{i}\in M_{g}\}$ & $\Lambda^{+}_{\textit{G3}}$=$\{v_{a},v_{a\rightarrow i_{1}},v_{a\rightarrow i_{2}},\ldots,v_{a\rightarrow i_{n}^{+}}\}\cup\{v_{i}\ |\ IoU(v_{i\rightarrow t},v_{a\rightarrow t})>o,\forall v_{i}\in M_{g}\}$.

Instead of relying on full-resolution features, we use local features from the overlapping regions of crop pairs. Given overlapping crops $P_{a}^{j},P_{a}^{k}$ from image $I_{a}$ where $IoU(P_{a}^{j},P_{a}^{k})\geq o$ , we sample corresponding overlapping crops from other images (e.g., $P_{b}^{j},P_{b}^{k}$). We introduce three improvements to local feature learning. First, only logits in the overlapping regions (e.g., $Z_{a}^{j\cap k},Z_{b}^{j\cap k}$) are used for representation learning. This strategy mitigates the deleterious effects of misalignment errors while still promoting local understanding, retaining sample diversity, and limiting computation cost. Second, we reverse the spatial transforms on logits (e.g., $Z^{\prime j\cap k}_{a}$=$\mathcal{T}_{s}^{-1}(Z_{a}^{j\cap k})$) before obtaining the final embeddings (e.g., $v^{\prime j\cap k}_{a}$=$h_{l}(Z^{\prime j\cap k}_{a})$) so that representations are contrasted in their original orientations and equivariance is maintained. Finally, spatial relations in embeddings are preserved with 1x1 convolutions in $h_{l}$ instead of an MLP. Note that $h_{l}$ also resizes logits to the local embedding size of $v_{l}\in\mathbb{R}^{C^{emb}\times H^{emb}\times W^{emb}}$.

In local strategy 1 (L1), we sample positives from the overlapping regions of the same crop pair and treat all $v_{l}\in M_{l}$ as negatives. Thus, we have $\Lambda^{-}_{\textit{L1}}$=$M_{l}$, $\Lambda^{+}_{\textit{L1}}$=$\{v^{\prime j\cap k}_{a},v^{\prime k\cap j}_{a}\}$. Note that $v_{a}^{j\cap k}$ is the overlapping region of $v_{a}^{j}$ & $v_{a}^{k}$ within $v_{a}^{j}$, while $v_{a}^{k\cap j}$ is the same corresponding area except in $v_{a}^{k}$. Local strategy 2 (L2) debiases negatives that are above an overlap threshold (the same $o$ as previously defined). More specificially, $\Lambda^{-}_{\textit{L2}}$=$\{v_{i}\ |\ IoU(v_{i\rightarrow t},v_{a\rightarrow t})\leq o,\forall v_{i}\in M_{l}\}$, $\Lambda^{+}_{\textit{L2}}$=$\{v^{\prime j\cap k}_{a},v^{\prime k\cap j}_{a}\}$. Local strategy 3 (L3) selects other corresponding overlapped regions as positives in addition to applying L2: $\Lambda^{-}_{\textit{L3}}$=$\{v_{i}\ |\ IoU(v_{i\rightarrow t},v_{a\rightarrow t})\leq o,\forall v_{i}\in M_{l}\}$, $\Lambda^{+}_{\textit{L3}}$=$\{v^{\prime j\cap k}_{a},v^{\prime k\cap j}_{a},v^{\prime j\cap k}_{b},v^{\prime k\cap j}_{b},\ldots\}$. Finally, local strategy 4 (L4) adopts the idea introduced in G3 that utilizes debiased negatives as positives: $\Lambda^{-}_{\textit{L4}}$=$\{v_{i}\ |\ IoU(v_{i\rightarrow t},v_{a\rightarrow t})\leq o,\forall v_{i}\in M_{l}\}$, $\Lambda^{+}_{\textit{L4}}$=$\{v^{\prime j\cap k}_{a},v^{\prime k\cap j}_{a},v^{\prime j\cap k}_{b},v^{\prime k\cap j}_{b},\ldots\}\cup\{v_{i}\ |\ IoU(v_{i\rightarrow t},v_{a\rightarrow t})>o,\forall v_{i}\in M_{l}\}$.

To compute $\mathbb{T}$ for volume alignment, we first crop out the background by thresholding at $-350$ Hounsfield Units, and downsample all images by a factor of 2. Affine registration is performed using the negative normalized cross correlation metric and tri-linear interpolation. Optimization is run until convergence (minimum 50 iterations) with a 0.5 learning rate.

For network components, $f$ is a 3D Res2Net-50 (scale=4, stride=[1,2,2] in the first downsampling layer) [8], while $g$ is a lightweight decoder. The global projection module $h_{g}$ follows [5] with average pooling to size 1x1x1, flattening, linear projection to 2048 channels, ReLU, and linear projection to embeddings of size $C_{g}^{emb}=128$. The local projection module resizes features to 1x3x3 via average pooling, applies 1x1x1 convolution with 1024 filters, activates with ReLU, and outputs embeddings after a 1x1x1 convolution with 64 filters ($v_{l}\in\mathbb{R}^{64\times 3\times 3},C_{l}^{emb}$=$64,H_{l}^{emb}$=$3,W_{l}^{emb}$=$3$; note the depth dimension is squeezed). In contrastive pretraining, we use $\mathcal{Q}_{g}$=$16000$, $\mathcal{Q}_{l}$=$1000$, $o$=$0.2$, $p$=$2$, $n^{+}$=$4$, and $\beta$=$0.99$. For losses, we set $\tau$=$0.2$, $\lambda$=$0.5$, and $\lambda_{r}$=$10$.

AMOS [12] (CT Abdominal Multi-Organ Segmentation) provides 240 training and 120 test abdominal CTs from diverse sources. We remove all “abnormal images” with anisotropic spacing or sizing along the right/left and anterior/inferior axes. Pretraining uses the remaining 200 training scans and is validated on 100 testing scans. Note that although masks are given for 15 organs, we do not use them in any capacity.

LUNA [13] (CT Lung Nodule Analysis) consists of 888 chest CTs across 10 folds. The first 7 folds (623 images) and the remaining 3 folds (265 images) are used for pretraining and validation, respectively.

BCV [1] (CT Beyond the Cranial Vault) contains 30 abdominal CTs with 13 anatomical annotations. The dataset is split into 21 training, 3 validation, and 6 testing images.

MMWHS [25] (CT & MR Multi-Modality Whole Heart Segmentation) presents 20 labeled CT and 20 labeled MR images with annotations for seven cardiac structures. We treat each modality as its own downstream task and evaluate them independently. For both CT & MR subsets, we split the labeled images into 14 training, 2 validation, and 4 testing.

We select state-of-the-art medical pretraining baselines like PCRL [23] TransVW [9], Rubik++ [15], PCL [20], and others. We choose PCL over [3] for its proposed improvements regarding slices near partition borders. We make comparisons as fair as possible with PCL (a 2D method) by matching the number of parameters in 2D Res2Net with our 3D model via width increases. FG [7] is omitted since global contrast coupled with reconstruction is explored by PCRL [23] and by us in § IV-A. We exclude HSSL [22] for similar reasons in addition to the fact that it only uses 2D slices.

First, we note that incorporating negatives, even without debiasing, improves over positive-only approaches (see rows 2 & 3). Also, enriching positives with spatially corresponding patches (G1, row 4) adds additional benefits. Debiasing negatives (rows 5-8), however, yields the largest increases when the overlap threshold $o$ is properly selected. A threshold that’s too low (row 5) may cause the undesirable removal of hard negatives (performing even worse than no debiasing in row 4). On the other hand, setting $o$ too high (row 7) prevents removal of damaging false negatives. Useful to contrastive learning, $o$ can modulate the amount of mutual information that’s distilled between samples (e.g., low $o$ amounts to negatives with low mutual information with the anchor). This enables pretraining flexibility in the face of wide ranges of possible tasks & data.

Next, we study the interactions between queue size and debiasing. In contrastive learning without debiasing (rows 9-11), we observe declines in performance as queue size increases. This is may be attributed to the increasing number of false negatives; surprisingly, the decline is limited possibly because the vast majority of queue entries are valid negatives. In contrast, our debiasing strategies (rows 12-14) empower larger queue sizes where performance steadily improves. This affirms the principle that appropriate selections of both positives & negatives facilitate contrastive representation learning.

In Table III, we explore the efficacy of local approaches by jointly pretraining using our best global strategy (G3 in row 1) and the indicated local strategy (note: the same $o$ & $n^{+}$ are adopted from the global strategy). To evaluate, we use the same precedure as experiments in Table II, except a single layer segmentation head (1x1x1 conv.) is appended to end of the decoder for pixel classification.

We first compare our equivariant contrastive method (L1, row 3) with a traditional reconstruction approach (row 2) and find the contrastive approach to be superior. However, we integrate reconstruction since we observe performance improvements, and also qualitatively see acceleration in pretraining convergence & speculate that it expedites initialization of sensible features which shortens feature warmup for contrastive tasks. Next, we find that local strategies (rows 3-6) behave in stark contrast to global ones in that diversifying positives (rows 5-6) impairs features. This is probably from rough correspondences yielding erroneous positives & negatives, thus, blunting both cohorts. So, our final proposed method utilizes G3, L2, and reconstruction. Comparing the training times of Spade to PCRL [23], our method takes 25.4 hours for 5.2 million patches while PCRL requires 28 hours.

To study the sensitivity of our approach to different templates, we select 3 full-torso CTs: two from within the pretraining set (in AMOS), and one outside of pretraining data (in BCV). The only selection criteria is that the entire torso is covered. From Table III, we conclude that although there’s a slight reduction in performance for the BCV template, Spade performs reasonably well for all three. This affirms the notion that approximate spatial correspondences are sufficiently effective for semantic comparisons.

From our main experiments in Table I, we make the following observations. 1) Spade generally outperforms recent state-of-the-art medical pretraining methods. This supports our assumption that spatial priors are effective for predicting semantic similarity and transferring this knowledge into useful representations. This also shows that simple sampling strategies can beat fairly complex pipelines like [23] which uses multiple regularization approaches like mixup, attention modules, and transformation embeddings. 2) Our method benefits performance the most when annotations are more limited. This is a desirable property for pretraining since its primary objective is to improve representations when annotations are most scarce due to acquisition costs. However, this is a double-edged sword in that when there’s more labels (e.g. MMWHS 50% labels), the cost of pretraining may not out-weight the benefit. 3) We affirm that effective negative samples are central to representation learning in 3D radiology images. Our studies on sampling strategies (see § IV-A, § IV-B) indicate that debiasing improves downstream performance more than enriching the diversity of positives. The task of making two patches with comparable semantics similar may be an easier task (courtesy of the curse of dimensionality) than contrasting similar patches. More studies regarding metric learning in medical images are needed to further explore this observation.