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

In plain words

Labelling 3D medical scans, such as heart MRI or abdominal CT, takes expert time, so a segmentation model often has only a few labelled scans to learn from. This paper trains a segmenter on a small labelled set plus a larger set of unlabelled scans. A "teacher" copy of the model, updated slowly from the "student" copy, gives stable targets.

The paper's main idea is a simple way to learn from the unlabelled scans. The model predicts, besides the object mask, a map of each voxel's distance to the object's edge. Two copies of the same scan go through the network with different random dropout masks. The two results of the same slice are treated as a matching pair, and other slices are treated as non-matching. The model is trained to pull matching pairs together. A second term asks each voxel to resemble the matching voxel in the other copy, so that the spatial layout agrees. Only the student is used at test time.

TL;DR

• A semi-supervised 3D segmenter (V-Net, mean-teacher setup) that predicts a probability map and a signed distance map (SDM) §3.3.
• Unlabelled volumes train a symmetric InfoNCE loss on per-slice features of the image plus SDM. The two views differ only by dropout masks §3.4 eq. 2 eq. 3.
• A voxel pair-wise distillation loss on encoder features is added §3.4 eq. 4.
• Author-reported: Dice 90.85% with 20% labels and 89.03% with 10% labels on the Left Atrium MR dataset tab. 1. No code is released.

Problem

Supervised 3D segmentation needs many labelled volumes. With few labels, a mean-teacher semi-supervised model is usually less robust than its supervised counterpart, and it has no explicit term for object geometry §1 §3.1. The authors also state that prior semi-supervised work mostly distils knowledge in the prediction space and that its local-patch contrastive objectives do not capture global boundary shape §3.4.

The setting is a small labelled set and a large unlabelled set: N labelled and M unlabelled volumes with N ≪ M §3.2. The experiments use 20% and 10% labelled ratios on LA and 20% on the pancreas §4.1.

Method

Overview: a student network and a teacher network, each a V-Net encoder and decoder, are trained together. The student is updated by SGD and the teacher by an exponential moving average of the student weights §3.1 §3.2. For each volume, two perturbed copies (affine transform and random crop) are created, one for each network §3.1.

Base architecture §3.3. Each decoder has two heads. The classification head outputs a probability map in [0,1]. The regression head outputs the SDM in [-1,1] through a hyperbolic tangent only. The authors say this encodes geometric structure and encourages smooth maps §3.3.

Supervised loss §3.3 eq. 1. On labelled volumes the loss is a segmentation loss (Dice plus cross-entropy) plus α times an MSE loss between the predicted SDM and the SDM computed from the ground truth.

\mathcal{L}_{\mathrm{sup}} = \frac{1}{N} \sum_{i=1}^{N} \mathcal{L}_{\mathrm{seg}}(\mathbf{Q}_i^{s}, \mathbf{Y}_i) + \frac{\alpha}{N} \sum_{i=1}^{N} \mathcal{L}_{\mathrm{mse}}(\mathbf{Q}_i^{s,\mathrm{sdm}}, \mathbf{Y}_i^{\mathrm{sdm}}) \qquad [\text{eq. 1}]

Boundary-aware features §3.4. For each unlabelled volume, the image is added to the student SDM and to the teacher SDM. This gives two boundary-aware feature volumes. The authors chose addition over concatenation, which they say needs a more complex head that over-fits more easily §3.4.

Projection head §3.4 §4.2. Each 2D slice of a boundary-aware volume passes through an alpha dropout layer (p = 0.1), an adaptive average pool to 128 × 128, and a 3-layer MLP that gives one D_h-dimensional vector per slice. The student and teacher heads use independent dropout masks z^s and z^t. The output size D_h is not stated.

Boundary-aware contrastive loss §3.4 eq. 2 eq. 3. The same slice from the two views is the positive. Slices from other locations or other volumes in the mini-batch are negatives. The InfoNCE term with temperature τ is

\mathcal{L}(\mathbf{h}^{t}_{i,j}, \mathbf{h}^{s}_{i,j}) = -\log \frac{\exp(\mathbf{h}^{t}_{i,j} \cdot \mathbf{h}^{s}_{i,j} / \tau)} {\sum_{k,l} \exp(\mathbf{h}^{t}_{i,j} \cdot \mathbf{h}^{s}_{k,l} / \tau)} \qquad [\text{eq. 2}]

and the loss is averaged over all unlabelled slices, in both directions:

\mathcal{L}_{\mathrm{contrast}} = \frac{1}{|\mathcal{N}^{+}|} \sum_{(i,j) \in \mathcal{N}^{+}} \left[ \mathcal{L}(\mathbf{h}^{t}_{i,j}, \mathbf{h}^{s}_{i,j}) + \mathcal{L}(\mathbf{h}^{s}_{i,j}, \mathbf{h}^{t}_{i,j}) \right] \qquad [\text{eq. 3}]

Pair-wise distillation loss §3.4 eq. 4. The encoder features of the student and teacher are flattened over their three spatial dimensions, giving H′W′D′ vectors of size D_e per volume. Cosine similarity s is computed between voxel vectors. For each student voxel the matching teacher voxel should get the highest softmax score among all teacher voxels of the same volume:

\mathcal{L}_{\mathrm{pd}} = -\frac{1}{M} \sum_{i=N+1}^{N+M} \sum_{j=1}^{H'W'D'} \log \frac{\exp(s(\mathbf{v}^{s}_{i,j}, \mathbf{v}^{t}_{i,j}))} {\sum_{k} \exp(s(\mathbf{v}^{s}_{i,j}, \mathbf{v}^{t}_{i,k}))} \qquad [\text{eq. 4}]

Consistency loss §3.4 eq. 5. An MSE term between the student output on a noisy view and the teacher output on another noisy view, over all unlabelled volumes. The noise distribution is not stated.

\mathcal{L}_{\mathrm{con}} = \frac{1}{M} \sum_{i=N+1}^{N+M} \mathcal{L}_{\mathrm{mse}}\left( \mathcal{F}_s(\mathbf{X}^{s}_{i} + \eta^{s}_{i}), \mathcal{F}_t(\mathbf{X}^{t}_{i} + \eta^{t}_{i})\right) \qquad [\text{eq. 5}]

Overall objective §3.4 eq. 6. The total loss is the supervised term plus three weighted unsupervised terms. The weights are called λ, β and γ.

\mathcal{L} = \mathcal{L}_{\mathrm{sup}} + \lambda \mathcal{L}_{\mathrm{contrast}} + \beta \mathcal{L}_{\mathrm{pd}} + \gamma \mathcal{L}_{\mathrm{con}} \qquad [\text{eq. 6}]

The contrastive loss is computed only on unlabelled indices (N⁺), which the authors say also affects the teacher and student §3.4.

# SimCVD training iteration  [§3.2] [§3.4] [eq. 1-6]
# Input:  labelled volumes x_l with masks y and SDMs y_sdm; unlabelled x_u
# Output: student updated by SGD; teacher updated by EMA

def training_step(x_l, y, y_sdm, x_u, step, max_step):
    # Two perturbed views per volume: affine transform and random crop  [§3.1]
    s_in_l, t_in_l = perturb(x_l)
    s_in_u, t_in_u = perturb(x_u)
    # Probability map and SDM from each network  [§3.3]
    prob_s, sdm_s = student(s_in_l)          # prob in [0,1], sdm in [-1,1]
    prob_t, sdm_t = teacher(t_in_u)
    # Supervised loss: Dice + cross-entropy, plus alpha * SDM MSE  [eq. 1]
    loss_sup = seg_loss(prob_s, y) + alpha * mse(sdm_s, y_sdm)
    # Boundary-aware features: image plus SDM, for unlabelled volumes  [§3.4]
    ba_s = s_in_u + sdm_s_unlabelled
    ba_t = t_in_u + sdm_t_unlabelled
    # Projection head with independent dropout masks (alpha dropout p = 0.1),
    # adaptive average pool to 128 x 128, then a 3-layer MLP per slice  [§4.2]
    h_s = projection_head(ba_s, mask_s)      # (slices, D_h)
    h_t = projection_head(ba_t, mask_t)      # (slices, D_h)
    # Symmetric InfoNCE over slices; same slice is the positive  [eq. 2] [eq. 3]
    loss_contrast = info_nce(h_t, h_s, tau) + info_nce(h_s, h_t, tau)
    # Voxel-to-voxel distillation on encoder features  [eq. 4]
    # Each volume: flatten encoder features to (H'W'D', D_e); cosine similarity
    v_s = flatten(encoder_student(x_u))      # (H'W'D', D_e)
    v_t = flatten(encoder_teacher(x_u))      # (H'W'D', D_e)
    # For voxel j of the student, score the teacher voxel j against all voxels
    # k of the same volume. No temperature is stated in the paper.
    loss_pd = pair_distill(v_s, v_t)
    # Student-teacher consistency on noisy views  [eq. 5]
    loss_con = mse(student(s_in_u + noise_s), teacher(t_in_u + noise_t))
    # Total objective: lambda = 0.5, beta = 0.1, gamma = 0.1  [eq. 6] [§4.2]
    loss = loss_sup + 0.5 * loss_contrast + 0.1 * loss_pd + 0.1 * loss_con
    # Ramp-up Psi(t); the weight it scales is not stated  [§4.2]
    ramp = exp(-5 * (1 - step / max_step) ** 2)
    # Student by SGD; teacher as an EMA of the student  [§3.2] [§4.2]
    sgd_update(student, loss)
    ema_update(teacher, student, decay=0.999)
    return loss

Hyperparameters

Results

Author-reported. Setup: V-Net backbone for all methods, Dice, Jaccard, ASD and 95HD (ASD and 95HD in voxels). The comparisons are against the published numbers listed in the paper §5.1 §5.2. The paper reports one set of numbers per setting and does not state the number of seeds or any variance §5 §4.2.

Ablations, LA with 10% labels (Dice %, Jaccard %, ASD, 95HD) tab. 3 tab. 4 tab. 5:

Headline claims in the text: on LA, SimCVD raises the best prior Dice from 89.94 to 90.85 at 20% and from 87.49 to 89.03 at 10% §5.1. On the pancreas, the authors report up to 6.72 Dice points over the weakest listed baseline (DAN) §5.2 tab. 2. The paper's own explanation of the gains is that contrastive learning, pair-wise distillation and SDM supervision each add to them §5.1 §6.1.

Paper vs code

No official code.

Search, in order: links in the paper and its footnotes (only the LA dataset site); a project page (none linked); GitHub repository search for "SimCVD", "simcvd", "contrastive voxel distillation", "semi-supervised segmentation dropout signed distance contrastive", "dropout contrastive distillation medical segmentation" in the Yale org, and "semi-supervised left atrial segmentation pancreas"; a GitHub user lookup for the first author's likely handle, which returned 404. All returned nothing relevant.

Pitfalls

• Abstract vs Table 1 (10% LA). The abstract says the 10% result is a 2.22 point gain. Table 1 and §5.1 give 89.03 against the best prior of 87.49, a gain of 1.54. The 2.22 figure equals the gain over SASSNet (86.81), which is not the best listed prior at 10%. Use Table 1.
• Mismatched "best previous" gains in §6.1. The text says the full model beats the best previous method by +0.71 Dice, +1.08 Jaccard, +1.20 ASD and +6.35 95HD. Those numbers are the differences to the "adaptive max pooling" ablation row (88.32, 79.26, 3.79, 14.69) in Table 3. Against the best prior (DTC) the Dice gain is +1.54.
• Baseline row mislabelled in Table 3. The row labelled "Baseline (UA-MT)" has values 84.24 / 73.26 / 2.71 / 19.41, which are the MT row of Table 1. The UA-MT row of Table 1 is 84.25 / 73.48 / 3.36 / 13.84. Which baseline the ablation uses is unclear.
• The "w/o L_sdm" ablation has no matching term. Eq. 6 has no L_sdm term. The SDM regression appears only as the α-weighted MSE inside L_sup eq. 1. An older commented-out SDM loss is also in the source. What was removed for this row is not stated.
• Dropout ASD sign. The text says "no dropout" (p = 0) worsens ASD by 1.71 voxels compared with p = 0.1. Table 4 gives ASD 2.49 for p = 0 and 2.59 for p = 0.1, so p = 0 is slightly better on ASD. The text's −1.71 does not match any row.
• Notation clash. The symbol α is the SDM-loss weight (0.1) in eq. 1 and the EMA decay (0.999) in §4.2. Keep them separate in code.
• Eq. 4 has no temperature, and the paper gives no value. Cosine similarities lie in [−1, 1], so the softmax over all voxels of a volume is nearly flat. The loss will stay large unless the implementation adds a temperature. The paper does not say whether one was used.
• Eq. 2 denominator includes the positive. The sum over (k, l) includes the matching pair, while the text calls only the other slices negatives. This is the usual InfoNCE convention, but the text should not be read as excluding the positive.
• Pair-wise distillation scales with the square of voxel count. Eq. 4 sums over all pairs of encoder positions per volume. The paper gives no memory or time cost for this term, nor for the per-slice contrastive batch.
• Dropout is the only view difference for the contrastive positives. The network inputs also get affine and crop perturbations §3.1, so a reader might expect those to make the positives. The contrastive loss compares the two dropout-masked projection outputs. Adding noise-type augmentations (local shuffle, intensity, in-painting, out-painting) on top of dropout did not beat the dropout-only setting in the paper's runs tab. 5.
• Dropout sensitivity. p = 0 and p = 0.5 both degrade results, and p = 0.2 is worse than 0.1 by about 0.9 Dice tab. 4. The paper does not say whether the p values were chosen on the evaluation split.
• Pooling size. Results grow with pool size up to 128 × 128, then drop at 256 × 256 tab. 4. The paper picks 128 with no stated variance.
• Ramp-up weight not stated. Ψ(t) is given §4.2, but which of the unsupervised weights it multiplies is not. The three unsupervised terms have fixed weights in eq. 6.
• Crop units. The input crop sizes are written as "mm³" (112 × 112 × 80, 96 × 96 × 96). The LA voxel size is 0.625 mm and the pancreas voxel size is 1.0 mm, so these are probably voxel counts. The paper does not say.
• Significance tests. Table 3 has a p-value column headed "[%]", while the text describes p-values near or below 0.05. Some entries are 2.47, 5.14 or 2.02, so the unit is unclear. The t-test is per-case Dice; the paper does not say whether the test is paired.
• Metrics not always best. At 10% LA, DTC has lower ASD (2.37 vs 2.59), and its 95HD is higher (9.06 vs 8.34). SimCVD is not best on every metric, although the conclusion says it obtained new state-of-the-art results.
• Single setting per dataset. Only the LA and pancreas datasets were tested, and the paper's own conclusion lists multi-class segmentation as future work §7. Results on other organs or modalities are not reported.

Lineage

• Builds on: the mean-teacher framework (Tarvainen and Valpola, "Mean teachers are better role models", NeurIPS 2017), which supplies the EMA teacher and the consistency term §3.2 §3.4. The V-Net backbone and the EMA decay follow UA-MT (Yu et al., MICCAI 2019) §3.3 §4.2. The SDM regression follows Xue et al. (AAAI 2020, shape-aware organ segmentation with signed distance maps) and Li et al. (MICCAI 2020, shape-aware semi-supervised 3D segmentation) §3.3. The contrastive loss follows Chaitanya et al. (NeurIPS 2020) §3.4. Pair-wise distillation is described as a structured-prediction variant of knowledge distillation §2.
• Compared against: V-Net, MT, DAN, CPS, Entropy Minimisation, UA-MT, ICT, SASSNet, DTC and Chaitanya et al. on LA tab. 1. On the pancreas: MT, DAN, CPS, Entropy Minimisation, UA-MT, ICT, SASSNet, DTC and Chaitanya et al. tab. 2.
• Superseded by: not known from this wiki.