[Paper Review] Learning to Prompt Segment Anything Models

Jiaxing Huang, Kai Jiang+5 more|arXiv (Cornell University)|Jan 9, 2024

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes Spatial-Semantic Prompt Learning (SSPrompt), a method that learns optimal spatial and semantic prompts for Segment Anything Models (SAMs) by optimizing in high-dimensional embedding space and selectively leveraging pre-trained prompt encoders. SSPrompt achieves state-of-the-art performance across multiple few-shot segmentation benchmarks, outperforming baselines by up to 16.0 mIoU on ADE20K under 16-shot settings.

ABSTRACT

Segment Anything Models (SAMs) like SEEM and SAM have demonstrated great potential in learning to segment anything. The core design of SAMs lies with Promptable Segmentation, which takes a handcrafted prompt as input and returns the expected segmentation mask. SAMs work with two types of prompts including spatial prompts (e.g., points) and semantic prompts (e.g., texts), which work together to prompt SAMs to segment anything on downstream datasets. Despite the important role of prompts, how to acquire suitable prompts for SAMs is largely under-explored. In this work, we examine the architecture of SAMs and identify two challenges for learning effective prompts for SAMs. To this end, we propose spatial-semantic prompt learning (SSPrompt) that learns effective semantic and spatial prompts for better SAMs. Specifically, SSPrompt introduces spatial prompt learning and semantic prompt learning, which optimize spatial prompts and semantic prompts directly over the embedding space and selectively leverage the knowledge encoded in pre-trained prompt encoders. Extensive experiments show that SSPrompt achieves superior image segmentation performance consistently across multiple widely adopted datasets.

Motivation & Objective

• Address the limited search space in spatial prompt optimization by shifting from 2D coordinates to high-dimensional embedding space.
• Mitigate side effects from pre-trained text prompt encoders that are biased toward foreground object descriptions.
• Develop a unified framework that learns both spatial and semantic prompts simultaneously for improved segmentation performance.
• Enable effective prompt learning under few-shot settings by leveraging knowledge from pre-trained prompt encoders selectively.
• Improve training efficiency by avoiding full fine-tuning of large prompt encoders while maintaining high performance.

Proposed method

• Introduce SpaPrompt: learnable weights fuse default spatial prompt embeddings (from fixed spatial prompt encoder) with learnable spatial prompt embeddings in the 512D embedding space.
• Introduce SemPrompt: learnable weights combine default semantic prompt embeddings (from fixed text prompt encoder) with learnable semantic prompt embeddings to form optimized semantic prompts.
• Optimize spatial and semantic prompts jointly in the embedding space, enabling a larger search space than 2D coordinate optimization.
• Use selective fusion via learnable weights to leverage well-learned foreground knowledge from prompt encoders while minimizing influence from poorly-learned background knowledge.
• Maintain inference efficiency by avoiding fine-tuning of large pre-trained prompt encoders, only updating small learnable prompt embeddings.
• Train using only a few annotated samples per class, enabling few-shot adaptation across diverse datasets.

Experimental results

Research questions

• RQ1 How does optimizing spatial prompts in high-dimensional embedding space improve segmentation performance compared to 2D coordinate optimization?
• RQ2 To what extent do pre-trained text prompt encoders introduce bias due to foreground-centric pre-training data, and how can this be mitigated?
• RQ3 Can joint learning of spatial and semantic prompts lead to complementary improvements in segmentation accuracy?
• RQ4 How does the proposed method perform under low-data regimes, particularly in few-shot settings?
• RQ5 What is the trade-off between performance gain and training efficiency when using selective prompt learning versus full fine-tuning?

Key findings

• SSPrompt achieves 55.2 mIoU on ADE20K under 16-shot settings, outperforming the baseline SEEM-T by 16.0 mIoU.
• On Cityscapes, SSPrompt improves mIoU from 39.2 (SEEM-T) to 55.2 under 16-shot, demonstrating consistent gains across datasets.
• The method maintains superior performance even with only 4-shot data, showing strong generalization in low-data regimes.
• Training time is reduced by 36.0% and memory by 53.5% compared to SEEM-T, due to avoiding full encoder fine-tuning.
• Visualized attention weights show higher values for foreground classes and lower for background, confirming selective use of well-learned knowledge.
• Ablation studies confirm that learnable weights in SemPrompt and SpaPrompt are essential for performance, as they enable selective knowledge utilization and avoid negative effects from poorly-learned background knowledge.