Adapting Foundation Models for Few-Shot Medical Image Segmentation: Actively and Sequentially

TL;DR

ASAP addresses the challenge of adapting foundation models to few-shot medical image segmentation with large domain gaps by reframing FSDA as a multi-armed bandit problem. It actively and sequentially selects among auxiliary datasets during a single-round fine-tuning, guided by a gradient-informed reward $R_{a,t}$ and a trace-UCB decision policy to balance exploration and exploitation. The method yields substantial Dice score improvements over state-of-the-art FSDA baselines on MRI and CT tasks while incurring lower computational overhead, demonstrating efficient reuse of public medical data and pretrained models. This approach promises fast, scalable, source-free adaptation suitable for diverse medical imaging tasks and resources.

Abstract

Recent advances in foundation models have brought promising results in computer vision, including medical image segmentation. Fine-tuning foundation models on specific low-resource medical tasks has become a standard practice. However, ensuring reliable and robust model adaptation when the target task has a large domain gap and few annotated samples remains a challenge. Previous few-shot domain adaptation (FSDA) methods seek to bridge the distribution gap between source and target domains by utilizing auxiliary data. The selection and scheduling of auxiliaries are often based on heuristics, which can easily cause negative transfer. In this work, we propose an Active and Sequential domain AdaPtation (ASAP) framework for dynamic auxiliary dataset selection in FSDA. We formulate FSDA as a multi-armed bandit problem and derive an efficient reward function to prioritize training on auxiliary datasets that align closely with the target task, through a single-round fine-tuning. Empirical validation on diverse medical segmentation datasets demonstrates that our method achieves favorable segmentation performance, significantly outperforming the state-of-the-art FSDA methods, achieving an average gain of 27.75% on MRI and 7.52% on CT datasets in Dice score. Code is available at the git repository: https://github.com/techicoco/ASAP.

Figures (3)

• Figure 1: Illustration of our active and sequential domain adaptation (ASAP) framework. The agent defines the policy $\pi$ that determines which arm to pull. The environment includes the auxiliary data pool $\mathcal{D}_\mathcal{A}$, the target dataset $\mathcal{D}_\mathcal{T}$, and the model $f_\theta$. At each turn $t$, ASAP executes the four shown steps.
• Figure 2: Visualization of different domain adaptation methods performance of two specific target tasks: WMH segmentation on MRI images and liver segmentation on CT images, both using MambaUNet. The pixels highlighted in red represent incorrect predictions.
• Figure 3: A 3D visualization of the active and sequential training process. The figure shows the selected auxiliary datasets at turn 0, turn 200, and turn 400, for two specific target tasks. The images are presented alongside their ground truth.