Knowledge-Guided Prompt Learning for Lifespan Brain MR Image Segmentation

TL;DR

This work tackles lifespan brain MRI segmentation with limited labeled data by introducing Knowledge-Guided Prompt Learning (KGPL), a two-step framework that first pre-trains on noisily labeled data and then fine-tunes with knowledge-wise prompts derived from a BiomedCLIP text encoder. The prompts are learnable tokens that interact with image embeddings while the encoder is frozen, enabling efficient adaptation across diverse ages and backbones, notably Swin UNETR, achieving state-of-the-art performance (average $DSC$ around 95% for tissue and 94% for structure) with far fewer trainable parameters. Quantitative analyses demonstrate consistent gains across backbones, significant improvements over full fine-tuning and random-prompt baselines, and up to 3x faster convergence, while qualitative results show closer alignment to ground truth in challenging brain regions. The approach holds practical value for clinical brain disease diagnosis across the lifespan by leveraging biomedical priors to learn robust, age-spanning segmentation.

Abstract

Automatic and accurate segmentation of brain MR images throughout the human lifespan into tissue and structure is crucial for understanding brain development and diagnosing diseases. However, challenges arise from the intricate variations in brain appearance due to rapid early brain development, aging, and disorders, compounded by the limited availability of manually-labeled datasets. In response, we present a two-step segmentation framework employing Knowledge-Guided Prompt Learning (KGPL) for brain MRI. Specifically, we first pre-train segmentation models on large-scale datasets with sub-optimal labels, followed by the incorporation of knowledge-driven embeddings learned from image-text alignment into the models. The introduction of knowledge-wise prompts captures semantic relationships between anatomical variability and biological processes, enabling models to learn structural feature embeddings across diverse age groups. Experimental findings demonstrate the superiority and robustness of our proposed method, particularly noticeable when employing Swin UNETR as the backbone. Our approach achieves average DSC values of 95.17% and 94.19% for brain tissue and structure segmentation, respectively. Our code is available at https://github.com/TL9792/KGPL.

Figures (3)

• Figure 1: Overview of our proposed two-step segmentation framework, incorporating Knowledge-Guided Prompt Learning (KGPL) for brain MRI across the lifespan. The top part shows the vision pre-training on the source domain, and the bottom part shows the refinement of the pre-trained models with knowledge-wise prompts on the target domain. Here, we adopt the weights learned from the source domain to initialize the model of the target domain, by only updating the learnable parameters and the decoder while freezing the encoder.
• Figure 2: Brain tissue segmentation results from different backbones are showcased in the transverse view. The first and second rows correspond to subjects aged 20 and 21, respectively. Each panel displays three segmentation results obtained by refining backbones using full, random prompts, and knowledge-wise prompts. Regions enhanced by our method are highlighted and enlarged within yellow boxes.
• Figure 3: Parcellation results of brain structure from three different backbones in the coronal view. The first and second rows correspond to subjects aged 17 and 24, respectively. Each panel displays three segmentation results achieved by refining backbones using full, random prompts, and knowledge-wise prompts. Regions improved by our proposed method are enclosed and magnified within yellow boxes.