Prompt tuning for parameter-efficient medical image segmentation.

Area of Science:

• Medical Imaging
• Computer Vision
• Machine Learning

Background:

• Self-supervised learning is standard for data-rich, annotation-scarce environments.
• Parameter-efficient fine-tuning is crucial for adapting pre-trained models to new tasks like semantic segmentation.
• Adapting models for new classes in semantic segmentation demands effective yet efficient methods.

Purpose of the Study:

• To develop a parameter-efficient yet effective adaptation method for semantic segmentation in medical imaging.
• To introduce a prompt-able UNETR (PUNETR) architecture adaptable via class-dependent learnable prompt tokens.
• To investigate the effectiveness of prompt tuning for semantic segmentation on CT imaging datasets.

Main Methods:

• Proposed a prompt-able UNETR (PUNETR) architecture with a frozen pre-trained backbone.
• Employed class-dependent learnable prompt tokens for network adaptation.
• Utilized a dense self-supervision scheme (contrastive prototype assignment, CPA) with a student-teacher model for pre-training.
• Incorporated an additional segmentation loss for a subset of classes during pre-training.

Main Results:

• PUNETR significantly reduced the performance gap between fully fine-tuned and parameter-efficient models.
• Achieved mean Dice Similarity Coefficient (DSC) differences of 7.81 pp (TCIA/BTCV) and 5.37-6.57 pp (TotalSegmentator subsets).
• Parameter efficiency was achieved by only adjusting prompt tokens (0.51% of the backbone parameters).

Conclusions:

• Prompt tuning offers a highly parameter-efficient adaptation strategy for medical semantic segmentation.
• PUNETR effectively adapts pre-trained models, minimizing performance loss compared to full fine-tuning.
• The proposed method demonstrates strong performance on CT imaging datasets with minimal parameter updates.