Neighbor Does Matter: Density-Aware Contrastive Learning for Medical Semi-supervised Segmentation

TL;DR

This work addresses semi-supervised multi-organ segmentation under limited labels and low tissue contrast by exploiting the geometry of the feature space. It introduces Density-Aware Contrastive Learning (DACL), which uses density-aware neighbor graphs and a memory bank to enhance intra-class compactness, and Soft Density-Guided Contrastive Learning (SDCL) to align sparse anchors with high-density positives. The approach samples low-density anchors, assembles high-density positives from both batch and memory, and weights their contributions with adaptive positiveness scores to form a center-based contrastive objective. Experiments on the ACDC and Synapse datasets demonstrate state-of-the-art performance and clear ablation-driven gains, underscoring the method's robustness and potential for improving clinical segmentation tasks.

Abstract

In medical image analysis, multi-organ semi-supervised segmentation faces challenges such as insufficient labels and low contrast in soft tissues. To address these issues, existing studies typically employ semi-supervised segmentation techniques using pseudo-labeling and consistency regularization. However, these methods mainly rely on individual data samples for training, ignoring the rich neighborhood information present in the feature space. In this work, we argue that supervisory information can be directly extracted from the geometry of the feature space. Inspired by the density-based clustering hypothesis, we propose using feature density to locate sparse regions within feature clusters. Our goal is to increase intra-class compactness by addressing sparsity issues. To achieve this, we propose a Density-Aware Contrastive Learning (DACL) strategy, pushing anchored features in sparse regions towards cluster centers approximated by high-density positive samples, resulting in more compact clusters. Specifically, our method constructs density-aware neighbor graphs using labeled and unlabeled data samples to estimate feature density and locate sparse regions. We also combine label-guided co-training with density-guided geometric regularization to form complementary supervision for unlabeled data. Experiments on the Multi-Organ Segmentation Challenge dataset demonstrate that our proposed method outperforms state-of-the-art methods, highlighting its efficacy in medical image segmentation tasks.

Figures (5)

• Figure 1: Left (a): Qualitative comparison of abdominal organs. As observed in the red boxes, the contrastive learning method that incorporates geometric neighborhood information achieves a more plausible appearance and a more complete structure compared to the conventional SSL method. Right (b): t-SNE van2008visualizing reveals that our method produces more compact class clusters with clearer separation between each class compared to the baseline. The clustering performance is evaluated using Silhouette rousseeuw1987silhouettes, Davies-Boulding davies1979cluster and V-Measure rosenberg2007v score. Our method outperforms the baseline in clustering performance metrics.
• Figure 2: Overview of the proposed unified learning framework. (a) shows the feature density-aware module. (b) shows our Soft Density-guided Contrastive Learning strategy. For labeled image $x_i^l$, we apply the commonly-used supervised loss $\mathcal{L}_{sup}$ to update model parameters. For unlabeled image $x_i^u$, the model is optimized with the cross-supervised consistency loss $\mathcal{L}_{cross}$ and one manifold constraint $\mathcal{L}^{soft}_{CL}$ to explore the geometry of the feature space.
• Figure 3: Visualization of the segmentation results from different methods on the Synapse dataset.
• Figure 4: Qualitative results on Synapse dataset using different components, corresponding to experiments I, II, III, V, and VI in the table \ref{['22ab']}.
• Figure 5: Mean Dice performances on the Synapse dataset with different hyperparameters. Performances with various (a) $\lambda_{cross}$ and $\lambda_{CL}$ in Eq. \ref{['eq1']}, (b) Memory bank size, (c) threshold $\phi$ for generating the 0-1 seed mask in Eq. \ref{['tau']}, and (c) temperature hyperparameters $\tau$ in Eq. \ref{['13']}.