Improving the Generalization of Segmentation Foundation Model under Distribution Shift via Weakly Supervised Adaptation

TL;DR

SAM generalizes poorly under distribution shift; this work proposes a source-free, weakly supervised self-training adaptation with anchor regularization and low-rank encoder updates to improve robustness and efficiency. By leveraging weak supervision prompts (box, points, coarse masks) aligned with SAM's prompt encoder, the method yields robust generalization across natural, corrupted, medical, camouflaged, and robotic segmentation tasks and outperforms state-of-the-art SFDA and WDASS baselines. The approach is memory-efficient and broadly task-agnostic, with strong cross-prompt generalization and practical impact for real-world segmentation.

Abstract

The success of large language models has inspired the computer vision community to explore image segmentation foundation model that is able to zero/few-shot generalize through prompt engineering. Segment-Anything(SAM), among others, is the state-of-the-art image segmentation foundation model demonstrating strong zero/few-shot generalization. Despite the success, recent studies reveal the weakness of SAM under strong distribution shift. In particular, SAM performs awkwardly on corrupted natural images, camouflaged images, medical images, etc. Motivated by the observations, we aim to develop a self-training based strategy to adapt SAM to target distribution. Given the unique challenges of large source dataset, high computation cost and incorrect pseudo label, we propose a weakly supervised self-training architecture with anchor regularization and low-rank finetuning to improve the robustness and computation efficiency of adaptation. We validate the effectiveness on 5 types of downstream segmentation tasks including natural clean/corrupted images, medical images, camouflaged images and robotic images. Our proposed method is task-agnostic in nature and outperforms pre-trained SAM and state-of-the-art domain adaptation methods on almost all downstream tasks with the same testing prompt inputs.

Figures (12)

• Figure 1: Segment Anything Model was pre-trained on a large-scale dataset but exhibits awkward performance on diverse downstream segmentation tasks. We adapt SAM through weak supervision to enhance its generalization capabilities.
• Figure 2: The proposed self-training architecture with anchor network regularization and contrastive loss regularization. Red arrows indicates the backpropagation flow.
• Figure 3: Illustration of contrastive loss between two views.
• Figure 4: qualitative results on some selected examples. Three types of prompts at testing stage are visualized for reference.
• Figure 5: Annotation cost vs. performance. 1. Performance of different numbers of weak labels on performance. 2. Performance of three weak labels under the same prompt verification.