SUGFW+: An Uncertainty-guided Feature Weighting Framework for Cold Start Active Adaptation of SAM in Medical Image Segmentation
Authors: Xiaochuan Ma, Ning Zhu, Jia Fu, Lanfeng Zhong, Hanyu Jiang, Bin Song, Kang Li, Guotai Wang
Organizations: School of Mechanical and Electrical Engineering, University of Electronic Science and Technology of China, Chengdu, China. · Glassgow College, University of Electronic Science and Technology of China, Chengdu, China. · Shanghai Artificial Intelligence Laboratory, Shanghai, China. · Department of Radiology, West China Hospital, Sichuan University, Chengdu, China.
Cold Start Active Learning (CSAL) is important in improving the performance of a medical image segmentation model with low annotation budget by querying a small subset for annotation from an unlabeled training set. Existing CSAL methods typically rely on inefficient dataset-specific Self-Supervised Learning (SSL) to map the unlabeled images into a feature space for sample selection. Recently, the advent of foundation models such as the Segment Anything Model (SAM) offer a promising alternative as the pre-trained model can provide strong generalizable feature embeddings, and allow high performance in downstream tasks after fine-tuning (adaptation). However, how to systematically exploit SAM's inherent embeddings for cold-start sample selection during adaptation with low annotation budget remains underexplored. To address this, we propose an extended SAM-based Uncertainty-guided Feature Weighting (SUGFW+) framework for CSAL and adaptation of SAM. Specifically, it leverages the SAM for Patch-level Feature and Uncertainty Calculation (PFUC), and introduces a Patch-based Global Distinct Representation (PGDR) module that aggregates patch-level embeddings into highly discriminative, uncertainty-aware image-level features. These features are then utilized by a Greedy Selection with Cluster and Uncertainty (GSCU) strategy to combine diversity and uncertainty during sample selection. Unlike prior CSAL methods that decouple sample selection from model training, SUGFW+ tightly integrates these two stages via an Uncertainty-Prompted Fine-Tuning (UPFT) process of SAM in model training. Extensive experiments on four public datasets demonstrate that SUGFW+ achieves state-of-the-art performance against existing CSAL methods. Code is available at https://github.com/HiLab-git/SUGFW-plus.
Figures & tables
Figure 1: Overview of our proposed SUGFW+ framework for CSAL in medical image segmentation. It consists of two core components: 1) SAM-based Uncertainty-Guided Feature Weighting (SUGFW) for selecting a small set of query samples from the unannotated target dataset, and 2) Uncertainty-Prompted Fine-Tuning (UPFT) of SAM with the query set for model training. SUGFW consists of three steps: 1) SAM-based Patch-level Feature and Uncertainty Calculation (PFUC), 2) Patch-based Global Distinct Representation (PGDR) and 3) Greedy Selection with Cluster and Uncertainty (GSCU).
Metric
Method
Annotation Budget
3.0%
4.0%
5.0%
6.0%
7.0%
8.0%
DSC (%) ↑
Random
84.03±5.87
85.72±4.86
85.74±4.33
86.24±6.79
86.55±6.27
86.64±6.60
ALPS [ 44 ]
81.43±8.39
82.46±10.14
84.97±7.09
86.44±5.37
87.92±4.48
88.39±5.65
CALR [ 22 ]
82.89±5.74
84.44±6.90
85.72±5.59
86.63±5.56
86.84±7.80
87.53±7.08
FPS [ 23 ]
84.24±6.15
85.40±5.48
86.95±6.31
87.16±4.48
87.32±5.50
88.21±3.56
Probcover [ 43 ]
83.08±4.82
86.22±4.82
86.49±5.26
86.70±4.15
88.53±3.56
88.88±3.52
Table 1: Quantitative comparison of different CSAL methods on the Promise12 dataset. The best results are in bold, and the second-best are underlined. ∗ indicates that the p-value < 0.05 in paired t-test compared to the second-best results.
Metric
Method
Annotation Budget
0.10%
0.15%
0.20%
0.25%
0.50%
0.75%
DSC (%) ↑
Random
70.96±11.96
75.17±5.32
81.51±4.84
82.51±4.06
85.00±3.58
87.67±3.23
ALPS [ 44 ]
73.80±6.48
77.24±5.69
80.72±5.19
83.73±4.58
85.68±5.08
87.61±4.21
CALR [ 22 ]
73.78±5.59
74.96±5.40
81.28±4.47
80.72±4.89
82.84±5.72
88.66±3.44
FPS [ 23 ]
46.97±12.54
69.74±13.30
70.16±14.79
71.58±10.94
71.76±7.62
82.57±7.74
Probcover [ 43 ]
62.72±12.83
74.78±5.21
76.83±5.17
76.90±4.78
78.41±4.62
83.30±4.13
Table 2: Quantitative comparison of different CSAL methods on the UTAH dataset. The best results are in bold, and the second-best are underlined. ∗ indicates that the p-value < 0.05 in paired t-test compared to the second-best results.
Metric
Method
Annotation Budget
0.05%
0.06%
0.08%
0.10%
0.20%
0.30%
DSC (%) ↑
Random
68.99±9.51
70.09±6.85
74.32±6.12
77.86±5.54
84.34±4.99
90.37±4.71
ALPS [ 44 ]
63.79±9.30
78.90±4.70
83.17±4.28
85.64±4.02
90.99±3.80
92.57±3.31
CALR [ 22 ]
62.19±11.36
67.35±9.91
71.24±8.43
75.49±7.16
84.24±5.00
91.97±5.29
FPS [ 23 ]
72.15±5.07
75.58±6.01
76.13±5.89
76.82±5.67
77.98±5.45
89.58±4.06
Probcover [ 43 ]
58.70±9.22
77.13±4.18
78.64±5.12
80.37±6.33
83.03±7.66
84.15±6.23
Table 3: Quantitative comparison of different CSAL methods on the MSD Liver dataset. The best results are in bold, and the second-best are underlined. ∗ indicates that the p-value < 0.05 in paired t-test compared to the second-best results.
Metric
Method
Annotation Budget
0.5 %
1.0%
1.5%
2.0%
2.5%
3.0%
DSC (%) ↑
Random
83.67±15.75
83.99±15.91
84.08±15.18
84.46±16.31
86.23±16.18
86.42±16.06
ALPS [ 44 ]
82.60±21.79
82.52±15.00
86.56±17.02
87.21±13.08
87.67±14.23
88.98±13.07
CALR [ 22 ]
84.00±15.00
85.26±14.55
85.23±17.16
86.74±12.59
87.18±12.56
88.45±14.38
FPS [ 23 ]
82.91±15.19
83.86±17.82
85.10±14.50
86.14±12.37
87.19±13.91
87.50±11.74
Probcover [ 43 ]
79.95±15.14
80.65±18.31
84.47±17.82
84.57±18.87
84.92±17.33
87.22±15.86
Table 4: Quantitative comparison of different CSAL methods on the ISIC 2018 dataset. The best results are in bold, and the second-best are underlined. ∗ indicates that the p-value < 0.05 in paired t-test compared to the second-best results.
Figure 2: Visual comparison of our method and other CSAL methods. All datasets are presented in axial views; additionally, the sagittal views and 3D reconstructions are provided for the left atrium and liver segmentations.
Method
UPFUC
FPGDR
GSCU
Promise12
ISIC 2018
3.0%
4.0%
5.0%
0.5%
1.0%
1.5%
Random
×
×
×
84.03±5.87
85.72±4.86
85.74±4.33
83.67±15.75
83.99±15.91
84.08±15.18
Un only
✓
×
×
85.73±5.05
86.60±4.07
87.71±4.00
85.94±12.40
86.09±13.31
86.94±13.56
w/o GSCU
✓
✓
×
81.22±6.18
85.32±6.78
87.55±5.24
82.46±17.38
83.72±14.85
86.26±12.26
w/o PGDR
✓
×
✓
85.13±4.66
85.30±5.25
87.49±5.61
83.85±16.23
86.19±13.51
86.58±11.47
SUGFW
✓
✓
✓
86.38±5.60
87.35±3.90
87.86±5.86
86.36±13.50
87.08±13.74
87.59±14.22
Table 5: Ablation study on the sample selection strategy SUGFW on Promise12 and ISIC 2018 datasets. UPFUC means unceratinty estiamtion based on PFUC, and FPGDR means image-level feature extraction based on PGDR. UPFT of SAM is used for model training after querying.
Method
Trainable Params (M)
Promise12
ISIC 2018
3.0%
4.0%
5.0%
0.5%
1.0%
1.5%
SwinUNet
62.19
60.59±2.01
64.47±5.43
72.63±2.86
57.85±27.76
59.26±27.40
61.44±26.10
nnUNet
32.42
63.07±10.51
68.41±7.31
76.05±5.03
70.87±25.80
74.03±26.12
76.96±23.83
DINOUNet
11.18
85.72±4.60
86.34±4.08
86.75±3.95
85.11±15.75
85.94±13.49
86.77±13.37
Ours w/o Un
28.67
85.01±5.67
86.72±4.23
87.52±5.38
86.02±13.80
86.88±14.93
87.02±13.54
Ours
29.86
86.38±5.60
87.35±3.90
87.86±5.86
86.36±13.50
87.08±13.74
87.59±14.22
Table 6: Ablation study on model training methods on Promise12 and ISIC 2018 datasets. The query samples are selected by our SUGFW strategy.
Figure 3: Hyper-parameter analysis of λ and αmax on the Promise12 dataset.
Medical image segmentation delineates organs, lesions, and other structures in scans to support measurement, diagnosis, treatment planning, and research.
Deep learning for 3D medical image segmentation requires extensive manual annotations, a major bottleneck in volumetric medical imaging. Active learning aims to reduce this burden by selecting informative samples for annotation, but most methods assume that an initial labeled set is already available. This leaves the cold-start problem largely unresolved: how to select the first volumes from a fully unlabeled pool before any task-specific model is trained. We propose CSCS, a Curriculum-Stratified Cold-Start framework that adapts initial sample selection to the structure of the unlabeled dataset. CSCS combines two self-supervised, label-free signals: local typicality, measuring representativeness in the embedding space, and reconstruction-based uncertainty, used as a proxy for sample difficulty. These signals are combined through a weighted geometric score, where the weighting is determined by a closed-form pacing rule based on the effective annotation budget and the Difficulty-Coverage Ratio, a pool-level statistic measuring the alignment between difficulty and representativeness. We evaluate CSCS on four 3D medical image segmentation benchmarks: BraTS, FeTA, Spleen, and an in-house fetal MRI dataset. Using nnU-Net as downstream segmentation model, CSCS shows consistently competitive performance across datasets and annotation budgets, with the strongest gains in low-to-mid annotation regimes. These results suggest that dataset-aware cold-start initialization can improve the robustness of active learning for 3D medical image segmentation by adapting sample selection to the geometry of the unlabeled pool.
Semi-supervised learning (SSL) has become a promising solution to alleviate the annotation burden of deep learning-based medical image segmentation models. While recent advances in foundation model-driven SSL have pushed the boundary to extremely limited annotation scenarios, they fail to maintain robust competitive performance in complex imaging modalities. In this paper, we propose SemiSAM-O1, an annotation-efficient framework using only one annotated template image for segmentation. SemiSAM-O1 extends the specialist-generalist collaborative learning framework to the extreme one-label setting by fully exploiting the foundation model's feature representation capability beyond its prompting interface. SemiSAM-O1 operates in two stages. In the first stage, the foundation model's encoder extracts dense features from all volumes, and class prototypes derived from the single annotated template are propagated to the unlabeled pool via feature similarity to produce coarse initial pseudo-labels. In the second stage, an iterative training-and-refinement loop progressively improves both the segmentation model and the pseudo-labels over multiple rounds, where each round trains the model from scratch on current pseudo-labels and generates updated predictions with voxel-wise uncertainty estimates. An uncertainty-guided refinement step further leverages the foundation model's global feature space to correct high-uncertainty regions by aggregating labels from their most similar confident neighbors, establishing a virtuous cycle of mutual improvement. Extensive experiments on a wide range of segmentation tasks across different modalities and anatomical targets demonstrate that SemiSAM-O1 significantly narrows the performance gap between one-label semi-supervised learning and full supervision, while significantly reducing the computational overhead of online foundation model inference.
Yichi Zhang, Le Xue, Bichun Xu +6
Artificial Intelligence Innovation and Incubation Institute, Fudan University, Shanghai, China. · Shanghai Academy of Artificial Intelligence for Science, Shanghai, China. · Department of Radiotherapy, Tongji Hospital, School of Medicine, Tongji University, Shanghai, China. +2
Segmentation models such as Segment Anything Model (SAM) and SAM2 achieve strong prompt-driven zero-shot performance. However, their training on natural images limits domain transfer to medical data. Consequently, accurate segmentation typically requires extensive fine-tuning and expert-designed prompts. We propose DiffuSAM, a diffusion-based adaptation of SAM2 for prompt-free medical image segmentation. Our framework synthesizes SAM2-compatible segmentation mask-like embeddings via a lightweight diffusion-prior from off-the-shelf frozen SAM2 image features. The generated embeddings are integrated into SAM2's mask decoder to produce accurate segmentations, thereby eliminating the need for user prompts. The diffusion prior is further conditioned on previously segmented slices, enforcing spatial consistency across volumes. Evaluated on the BTCV and CHAOS datasets for CT and MRI under Source-Free Unsupervised Domain Adaptation (SF-UDA) and Few-Shot settings, DiffuSAM achieves competitive performance with efficient training and inference. Code is available upon request from the corresponding author.
Tal Grossman, Noa Cahan, Lev Ayzenberg +1
School of Electrical Engineering, Tel Aviv University, Israel · School of Biomedical Engineering, Tel Aviv University, Israel