Robust Noisy Pseudo-label Learning for Semi-supervised Medical Image Segmentation Using Diffusion Model

Visual encoder. The visual encoder takes an image $\bm{I}$ as input and extracts multi-level features $\bm{F}_{i}$$\in$$\mathbb{R}^{C\times H\times W}$ ($i$$\in$$\{1,2,...,N\}$) for diffusion decoder as visual condition to guide denoising. $H$, $W$, and $C$ denote the height, width, and channels of one-level feature, respectively. $N$ indicates the total number of levels of visual encoder. Specifically, we first utilize a backbone ViT [6] for visual branch to extract visual features. We select features from different layers of the backbone, denoted as $\bm{X}_{l}$ ($l$$\in$$\{l_{1},l_{2},...,l_{N}\}$), where $l$ is the layer index. The $N$-level features are fed into visual-specific branches to generate the visual-specific multi-level features $\bm{F}_{i}$. The visual-specific branch $\varphi(\cdot)$ contains two stacked convolutional blocks, each of which consists of a 3$\times$3 convolution followed by a batch normalization layer and a GeLU activation layer and a 1$\times$1 convolution. Moreover, to learn discriminative multi-level visual-specific features, we also design visual-specific auxiliary heads to generate intermediate predictions. These intermediate predictions are supervised by the corresponding semantic ground-truth and pseudo labels $M_{j}$ ($j$$\in$$\{\textrm{gt},\textrm{pseudo}\}$) and the visual-specific branches will be updated by the gradient from segmentation prediction. The output of each visual-specific branch acts as the conditions for the diffusion decoder.

Label and latent encoder. We first encode the discrete, dense semantic label map into a latent feature space using a label encoder to capture semantic representations. Given the discrete nature of semantic segmentation labels, we begin by converting them into one-hot representations. The label encoder consists of a 1$\times$1 convolutional layer and a 3$\times$3 convolutional layer that transforms the dense labels into label-specific encoded features $z$. This output $z$ is then normalized to the range [$-1,+1$], and a scaling factor $s$ is applied to control the signal-to-noise ratio. This scaling increases the difficulty of the denoising task, thereby encouraging the diffusion decoder to learn more informative representations. To simulate the forward diffusion process, Gaussian noise is added to the label-specific encoded features $z$, producing the corrupted representation $z_{t}$. In the forward noising process, the corruption intensity is controlled by $\gamma(t)$$\in$[$0,1$], which decreases over time $t$. Following previous works [9, 23], we adopt a cosine schedule for $\gamma(t)$ to modulate the noise level throughout the diffusion process.

Diffusion decoder The decoder takes the noisy label-specific encoded features $z_{t}$ as input and the visual-specific multi-level features $\bm{F}_{i}$ as conditions. We also use the features $\bm{X}_{l}$ from the backbone different layers to help model the semantic space. The noisy label map $z_{t}$ is concatenated with the features $\bm{X}_{l}$ and sent to the diffusion decoder to perform level interaction successively with $\bm{F}_{i}$. Specifically, in the level-interaction phase, we perform the interaction at each level. We utilize 2 convolutional blocks (Conv-BatchNorm-GeLU-Conv) and a convolutional layer to map the concatenated feature map $z_{t}$ and $\bm{X}_{l}$ to each feature level and the channel number is transformed from $C$ to $1$. The multi-level mask maps are generated by applying a Sigmoid function to the output features of the convolutional block at each level, denoted as $\bm{M}_{i}$$\in$$\mathbb{R}^{1\times H\times W}$ ($i$$\in$$\{1,2,...,N\}$). Then, we use $\bm{M}_{i}$ to generate the final aggregated features by $\bm{F}^{\prime}$=$\sum_{i=1}^{N}\bm{M}_{i}\cdot\bm{F}_{i}$. The aggregated feature $\bm{F}^{\prime}$ is sent to the prediction branch to generate the final prediction, where the prediction branch contains three convolutional blocks. The predictions are then encoded as mentioned in $\S$ 2.1 to generate the predicted $\tilde{z}_{0}$.

Inter-class compactness. We begin by grouping the latent embeddings $\bm{Z}$ into $K$ prototypes $\{\bm{p}_{c,k}\}_{k=1}^{K}$ for each semantic class $c$ in an online clustering manner. After processing all samples in the current batch, each pixel embedding $\bm{i}$$\in$$\bm{Z}$ is assigned to the $k$-th prototype of class $c$, where the assignment is determined by $k$=$\mathrm{arg}\max_{k}\{l\}_{k=1}^{k}$ with $l$=$\left\langle\bm{i},\bm{p}_{c,k}\right\rangle$ and $\left\langle\cdot,\cdot\right\rangle$ denotes a similarity metric, e.g., the cosine similarity. This prototype assignment leads to a training objective aimed at maximizing the posterior probability of the correct prototype assignment. This can be viewed as a inter-class prototype contrastive learning loss:

where $\mathcal{P}^{-}$ denotes the set of prototypes excluding those belonging to the target class $c$, and the temperature parameter $\tau$ regulates the concentration of the similarity distribution.

Intra-class separability. To further reduce intra-class variation, i.e., to encourage pixel features assigned to the same prototype to form compact clusters, we introduce a loss that regularizes the learned representations by directly minimizing the distance between each embedded pixel and its assigned prototype:

Notably, both $\bm{i}$ and $\bm{p}_{c,k}$ are $\ell_{2}$-normalized. This normalization ensures that the training objective focuses on angular similarity, minimizing intra-class variation while preserving separation between features assigned to different prototypes.

Datasets. We evaluate our model using two datasets. First, we employ the public Endoscapes2023 dataset [22], which includes three sub-tasks: surgical scene segmentation, object detection, and critical view of safety assessment. For our experiments, we use the segmentation subset, which comprises 493 annotated frames extracted from 50 laparoscopic cholecystectomy videos. The dataset is officially split into training, validation, and test sets in a 3:1:1 ratio, respectively.

Evaluation metric. Following conventions [25, 24], mean intersection-over-union (mIoU) is adopted for evaluation.

Implementation Details. We utilize ViT-Large [6] as the visual backbone for our main experiments and ViT-Base for all ablation studies, both producing features with a stride of 16. All models are trained for 40,000 iterations with a batch size of 8. The network is optimized using the Adam optimizer ($\beta_{1}$=0.9, $\beta_{2}$=0.999), with an initial learning rate of 5e-4 and a weight decay of 1e-6 for both the Endoscapes2023 and MOSXAV datasets and a polynomial learning rate scheduler is employed. The semantic segmentation task is supervised using the cross-entropy loss. To generate pseudo masks for Endoscapes2023, we first use object detection bounding boxes as prompts for the MedSAM [16]. The initial masks produced by MedSAM are then refined using SAMRefiner [12] to obtain higher-quality pseudo labels. All the experiments are trained with 2 NVIDIA RTX 6000 GPUs for 40,000 iterations.

Endoscapes2023 dataset. Our method significantly outperforms baseline semantic segmentation approaches under the semi-supervised learning setting (i.e., GT + Pseudo), notably surpassing the diffusion-based model MedSegDiff [32] by a margin of +5.15%. Furthermore, our method also achieves the best performance when trained solely on ground truth or pseudo labels.

MOSXAV val and test sets. Table 1 shows that our method outperforms all baseline semantic segmentation methods on the val set. More importantly, on the challenging test set, it consistently surpasses MedSegDiff, achieving an mIoU improvement from 39.93% to 43.04%.

Effectiveness of Contrastive Loss. To verify the effectiveness of our overall training objective, we progressively incorporate contrastive learning components defined in Eqs. 2 and 3 into the final loss function. As shown in Table 2, the model with $\mathcal{L}_{\mathrm{CE}}$ alone achieves an mIoU score of 53.04%. Combing all the losses together leads to the best performance, yielding an mIoU score of 54.47%.

Prototype Number Per Class $K$. Table 2 presents the performance of our approach concerning the number of prototype per class. For $K$=1, we directly represent each class as the mean embedding of its pixel samples. As a result, we set $K$=10 for a better trade-off between accuracy and computation cost.