[Paper Review] RepViT-SAM: Towards Real-Time Segmenting Anything

Ao Wang, Hui Chen+3 more|arXiv (Cornell University)|Dec 10, 2023

Advanced Neural Network Applications4 citations

TL;DR

RepViT-SAM replaces the heavy ViT image encoder in Segment Anything Model (SAM) with RepViT-M2.3, a mobile-optimized vision transformer that uses efficient convolutional architectures and structural reparameterization, achieving nearly 10× faster inference than MobileSAM while maintaining or exceeding its zero-shot segmentation performance across diverse benchmarks.

ABSTRACT

Segment Anything Model (SAM) has shown impressive zero-shot transfer performance for various computer vision tasks recently. However, its heavy computation costs remain daunting for practical applications. MobileSAM proposes to replace the heavyweight image encoder in SAM with TinyViT by employing distillation, which results in a significant reduction in computational requirements. However, its deployment on resource-constrained mobile devices still encounters challenges due to the substantial memory and computational overhead caused by self-attention mechanisms. Recently, RepViT achieves the state-of-the-art performance and latency trade-off on mobile devices by incorporating efficient architectural designs of ViTs into CNNs. Here, to achieve real-time segmenting anything on mobile devices, following MobileSAM, we replace the heavyweight image encoder in SAM with RepViT model, ending up with the RepViT-SAM model. Extensive experiments show that RepViT-SAM can enjoy significantly better zero-shot transfer capability than MobileSAM, along with nearly $10 imes$ faster inference speed. The code and models are available at \url{https://github.com/THU-MIG/RepViT}.

Motivation & Objective

• To enable real-time inference of the Segment Anything Model (SAM) on resource-constrained mobile devices.
• To address the high memory and computational overhead of self-attention mechanisms in lightweight SAM variants like MobileSAM.
• To improve zero-shot transfer performance of lightweight SAM models across diverse segmentation tasks.
• To leverage knowledge distillation to transfer knowledge from the large ViT-H encoder to a smaller, efficient RepViT-based encoder.
• To establish a new mobile-optimized baseline for edge deployment of segmentation models with strong performance and low latency.

Proposed method

• Replace the original ViT-H image encoder in SAM with RepViT-M2.3, a lightweight, convolutional architecture with structural reparameterization and depthwise separable convolutions.
• Apply decoupled knowledge distillation: directly distill the RepViT-M2.3 encoder from the ViT-H encoder using a mean squared error (MSE) loss on feature maps.
• Set the stride of the final downsampling layer in RepViT to 1 to ensure output resolution compatibility with SAM's prompt-guided mask decoder.
• Precompute and cache image embeddings from the ViT-H encoder during distillation to accelerate training without re-running the heavy encoder.
• Use the same prompt-guided mask decoder as in the original SAM, maintaining consistency in inference pipeline.
• Employ a 16×16 grid of foreground point prompts for zero-shot edge detection and apply standard post-processing including Sobel filtering and non-maximum suppression.

Figure 1 : Mask predictions of SAM, MobileSAM, and RepViT-SAM with point prompts (top) and box prompts (bottom).

Experimental results

Research questions

• RQ1 Can a lightweight, convolutional vision transformer like RepViT achieve real-time inference in SAM while preserving strong zero-shot generalization?
• RQ2 How does the performance of RepViT-SAM compare to MobileSAM and ViT-B-SAM across diverse zero-shot segmentation benchmarks?
• RQ3 To what extent can knowledge distillation from a large ViT-H encoder transfer high-quality features to a smaller RepViT-based encoder?
• RQ4 Does the pure convolutional design of RepViT enable better latency and memory efficiency than attention-based lightweight encoders on mobile platforms?
• RQ5 Can RepViT-SAM outperform even the original ViT-H-SAM in specific downstream tasks like anomaly detection due to better feature distillation?

Key findings

• RepViT-SAM achieves 48.9ms inference latency on iPhone 12, while MobileSAM fails due to out-of-memory (OOM) issues, demonstrating its feasibility for mobile deployment.
• On Macbook M1 Pro, RepViT-SAM runs in 44.8ms, which is nearly 10× faster than MobileSAM (482.2ms) and 140× faster than ViT-H-SAM (6249.5ms).
• In zero-shot edge detection, RepViT-SAM achieves 0.764 ODS, 0.786 OIS, and 0.773 AP, outperforming MobileSAM (0.756, 0.768, 0.746) and matching ViT-H-SAM (0.768, 0.786, 0.794) in OIS.
• For zero-shot instance segmentation on COCO, RepViT-SAM achieves 46.1 mean AP, surpassing MobileSAM (43.9) and ViT-B-SAM (44.8), and approaching ViT-H-SAM (48.7).
• In zero-shot video object segmentation on DAVIS 2017, RepViT-SAM achieves 73.5 mean J&F, outperforming MobileSAM (71.1) and ViT-B-SAM (71.3), with a 2.4-point gain over MobileSAM.
• On MVTec-AD for zero-shot anomaly detection, RepViT-SAM achieves 37.96 F1-pixel score, exceeding ViT-H-SAM (37.65) and outperforming MobileSAM (36.44), indicating superior feature transfer in specialized tasks.

Figure 2 : Visualization results of SAM, MobileSAM, and RepViT-SAM for zero-shot edge detection on BSDS500.