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

TL;DR

This paper introduces LLM-SegNet, a semi-supervised 3D medical image segmentation model that leverages large language models to incorporate task-specific knowledge, improving segmentation accuracy with limited annotations.

Contribution

The paper presents a novel framework that integrates LLMs into semi-supervised segmentation, along with a unified loss function to reduce errors and enhance learning from unannotated data.

Findings

Outperforms state-of-the-art methods on multiple datasets

Effectively utilizes unannotated data for improved segmentation

Ablation studies confirm the contribution of each module

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.