Fast SAM2 with Text-Driven Token Pruning

TL;DR

The paper tackles the computational bottleneck in SAM2-based video object segmentation caused by dense token propagation across time. It introduces a post-encoder, text-guided token pruning framework that fuses semantic prompts, predictive uncertainty via Monte Carlo Dropout, and local context to selectively retain informative tokens before memory propagation. This approach yields substantial efficiency gains—up to 42.5% faster inference and 37.4% GPU memory reduction—while maintaining competitive segmentation quality across multiple benchmarks. The method is training-free at deployment, modular, and robust to prompt variations, enhancing the practicality of transformer-based VOS in real-time and resource-constrained scenarios.

Abstract

Segment Anything Model 2 (SAM2), a vision foundation model has significantly advanced in prompt-driven video object segmentation, yet their practical deployment remains limited by the high computational and memory cost of processing dense visual tokens across time. The SAM2 pipelines typically propagate all visual tokens produced by the image encoder through downstream temporal reasoning modules, regardless of their relevance to the target object, resulting in reduced scalability due to quadratic memory attention overhead. In this work, we introduce a text-guided token pruning framework that improves inference efficiency by selectively reducing token density prior to temporal propagation, without modifying the underlying segmentation architecture. Operating after visual encoding and before memory based propagation, our method ranks tokens using a lightweight routing mechanism that integrates local visual context, semantic relevance derived from object-centric textual descriptions (either user-provided or automatically generated), and uncertainty cues that help preserve ambiguous or boundary critical regions. By retaining only the most informative tokens for downstream processing, the proposed approach reduces redundant computation while maintaining segmentation fidelity. Extensive experiments across multiple challenging video segmentation benchmarks demonstrate that post-encoder token pruning provides a practical and effective pathway to efficient, prompt-aware video segmentation, achieving up to 42.50 percent faster inference and 37.41 percent lower GPU memory usage compared to the unpruned baseline SAM2, while preserving competitive J and F performance. These results highlight the potential of early token selection to improve the scalability of transformer-based video segmentation systems for real-time and resource-constrained applications.

Figures (9)

• Figure 1: Qualitative visualisation of retained tokens with the text-driven token pruning approach (on UVO dataset) wang2021unidentifiedvideoobjectsbenchmark before passing into SAM2 decoder
• Figure 2: Qualitative visualization steps of segmentation with our text-driven token pruning approach atop SAM2
• Figure 3: End-to-end pipeline of our text-driven token pruning framework atop SAM2. Semantic cues from text prompts, visual context from ViT tokens, along with token uncertainty are fused by a lightweight router to retain only task-relevant tokens before SAM2 decoding.
• Figure 4: Monte Carlo (MC) Dropout applied to ViT-Hiera layers L3–L5 for uncertainty prediction per token
• Figure 5: Qualitative segmentation results on a UVO dataset wang2021unidentifiedvideoobjectsbenchmark sequence with our text-driven token pruning approach atop SAM2. We achieve mean $\mathcal{J\&F}$ of 97.95% with no other refinement clicks required other than seed click to serve as input to SAM2's prompt encoder. The image shows the dense tokens in the input frame, followed by semantically relevant tokens at 30% retention, and the segmentation mask by SAM2, with the ground truth also shown for comparison.