Leveraging Task-Specific Knowledge from LLM for Semi-Supervised 3D Medical Image Segmentation

TL;DR

This work tackles the high cost of voxel-level annotations in 3D medical image segmentation by proposing LLM-SegNet, a semi-supervised framework that injects task-specific knowledge from a large language model into a co-training CNN setup. It introduces a Unified Segmentation Loss (USL) that combines negative log-likelihood and mean-squared error in a confidence-aware scheme to better exploit unlabeled data. The approach uses textual descriptors generated by an LLM to augment image representations, fused at the encoder bottleneck via a learnable balance with text features, enabling efficient learning without requiring 3D image-text pairs. Empirical results on Left Atrium, Pancreas-CT, and Brats-19 demonstrate consistent improvements over state-of-the-art SSL methods, supported by ablative analyses that highlight the contributions of the LLM, textual guidance, supervised losses, and USL to overall performance.

Abstract

Traditional supervised 3D medical image segmentation models need voxel-level annotations, which require huge human effort, time, and cost. Semi-supervised learning (SSL) addresses this limitation of supervised learning by facilitating learning with a limited annotated and larger amount of unannotated training samples. However, state-of-the-art SSL models still struggle to fully exploit the potential of learning from unannotated samples. To facilitate effective learning from unannotated data, we introduce LLM-SegNet, which exploits a large language model (LLM) to integrate task-specific knowledge into our co-training framework. This knowledge aids the model in comprehensively understanding the features of the region of interest (ROI), ultimately leading to more efficient segmentation. Additionally, to further reduce erroneous segmentation, we propose a Unified Segmentation loss function. This loss function reduces erroneous segmentation by not only prioritizing regions where the model is confident in predicting between foreground or background pixels but also effectively addressing areas where the model lacks high confidence in predictions. Experiments on publicly available Left Atrium, Pancreas-CT, and Brats-19 datasets demonstrate the superior performance of LLM-SegNet compared to the state-of-the-art. Furthermore, we conducted several ablation studies to demonstrate the effectiveness of various modules and loss functions leveraged by LLM-SegNet.

Figures (6)

• Figure 1: Illustration of the process for creating specific word embeddings. Initially, we provide prompt input to LLM, which generates a textual response. This response is then fed into the BERT tokenizer to produce the corresponding word embeddings.
• Figure 2: Overview of our LLM-SegNet approach. We apply the supervised loss to align the labeled predictions from model $\mathcal{A}$ with the ground truth and do the same for model $\mathcal{B}$. To learn from unlabeled data, we utilize the proposed unified segmentation loss (USL) between the pseudo-label from the other model and the model's own unlabeled predictions. Furthermore, our co-training framework uses both text and image features jointly to train the model.
• Figure 3: Qualitative results of our method compared to recent state-of-the-art methods on the LA dataset with a 10 % labeled ratio.
• Figure 4: Qualitative results of our method compared to recent state-of-the-art methods on the pancreas dataset with a 20 % labeled ratio.
• Figure 5: Qualitative results of our method compared to recent state-of-the-art methods on the Brats-2019 dataset with a 10 % labeled ratio.