Source-linked AI summary
Faster Segment Anything: Towards Lightweight SAM for Mobile Applications
Chaoning Zhang, Dongshen Han, Yu Qiao, Jung Uk Kim, Sung-Ho Bae, Seungkyu Lee, Choong Seon Hong
TL;DR
SAM’s heavyweight image encoder limits deployment of versatile segmentation applications on resource-constrained mobile devices, while naive retraining with a smaller encoder performs unsatisfactorily under limited training sources. MobileSAM addresses this through decoupled distillation, producing a lightweight model that performs on par with original SAM while substantially reducing size and inference time.
Problem
SAM supports versatile segmentation applications, but its heavyweight image encoder makes the pipeline unsuitable for resource-constrained mobile devices, and straightforward smaller-encoder retraining performs unsatisfactorily with limited training sources.
Method
MobileSAM distills knowledge from the original SAM’s ViT-H image encoder into a lightweight encoder and separates image-encoder distillation from mask-decoder finetuning.
Results
More than 60 times smaller, MobileSAM performs on par with original SAM; it reduces encoder parameters by 100 times and total parameters by 60 times, with around 10ms single-GPU inference.
Takeaways & Limitations
MobileSAM provides a lightweight, faster alternative to original SAM for pushing prompt-guided segmentation toward mobile applications.
Abstract
from arXiv · showhide
Segment Anything Model (SAM) has attracted significant attention due to its impressive zero-shot transfer performance and high versatility for numerous vision applications (like image editing with fine-grained control). Many of such applications need to be run on resource-constraint edge devices, like mobile phones. In this work, we aim to make SAM mobile-friendly by replacing the heavyweight image encoder with a lightweight one. A naive way to train such a new SAM as in the original SAM paper leads to unsatisfactory performance, especially when limited training sources are available. We find that this is mainly caused by the coupled optimization of the image encoder and mask decoder, motivated by which we propose decoupled distillation. Concretely, we distill the knowledge from the heavy image encoder (ViT-H in the original SAM) to a lightweight image encoder, which can be automatically compatible with the mask decoder in the original SAM. The training can be completed on a single GPU within less than one day, and the resulting lightweight SAM is termed MobileSAM which is more than 60 times smaller yet performs on par with the original SAM. For inference speed, With a single GPU, MobileSAM runs around 10ms per image: 8ms on the image encoder and 4ms on the mask decoder. With superior performance, our MobileSAM is around 5 times faster than the concurrent FastSAM and 7 times smaller, making it more suitable for mobile applications. Moreover, we show that MobileSAM can run relatively smoothly on CPU. The code for our project is provided at \href{https://github.com/ChaoningZhang/MobileSAM}{\textcolor{red}{MobileSAM}}), with a demo showing that MobileSAM can run relatively smoothly on CPU.
1 Introduction
SAM combines a ViT-based image encoder with a prompt-guided mask decoder to support versatile, label-free segmentation, but its heavyweight encoder limits mobile deployment. MobileSAM addresses this bottleneck through decoupled distillation, achieving substantial size and speed reductions while preserving performance.
- SAM overview: SAM uses a ViT-based image encoder followed by a prompt-guided mask decoder.The decoder generates masks from prompts such as points or boxes.
- SAM overview: SAM enables label-free segmentation and supports applications including text-guided segmentation and fine-grained image editing.Its segmentation capability can be combined with other models for advanced vision applications.
- Motivation: Naively retraining SAM with a smaller encoder is difficult because the image encoder and mask decoder are coupled during optimization.The original training approach can require 128 GPUs for multiple days with ViT-L or ViT-B encoders.
- Proposed approach: Decoupled distillation transfers knowledge from ViT-H to a lightweight encoder that is compatible with the original mask decoder.The resulting approach can be reproduced on a single GPU in less than one day.
- Results: MobileSAM reduces total parameters by 60 times and runs around 10ms per image while performing on par with original SAM.The reported timing is 8ms for the image encoder and 4ms for the mask decoder.
2 Related work
Related work positions SAM as a versatile foundation model while examining its performance across applications and efforts to make vision transformers lightweight and efficient.
- SAM generalization and versatility: Research has extended SAM to medical, camouflaged, and transparent objects, where reported performance is weaker than in general setups.These studies examine SAM in challenging real-world conditions.
- SAM generalization and versatility: SAM has been combined with grounding, CLIP, and other models for text-guided or semantic segmentation, image editing, inpainting, tracking, and 3D reconstruction.Grounding DINO can generate a box prompt that SAM uses to produce a mask.
- ViT: lightweight and efficient: Lightweight vision-transformer research includes DeiT-Small, DeiT-Tiny, MobileViT, and EfficientFormer, often combining transformer representations with convolutional locality.MobileViT is described as combining ViT with standard convolutions to improve efficiency.
3 Mobile-Friendly SAM
Mobile-Friendly SAM replaces costly coupled retraining with a decoupled approach that distills the image encoder separately and optionally fine-tunes the lightweight mask decoder. The method is evaluated against coupled distillation and is designed for efficient mobile deployment.
- SAM background: SAM combines a ViT-based image encoder with a prompt-guided mask decoder that generates object masks from points or boxes.The encoder produces an image embedding consumed by the decoder, which can generate multiple masks for ambiguous prompts.
- Coupled distillation: The original retraining strategy couples optimization of the smaller image encoder and mask decoder, making both modules harder to improve when initially weak.The proposed alternative addresses this coupled optimization problem by separating encoder distillation from decoder finetuning.
- Decoupled distillation: Decoupled distillation first transfers knowledge from the ViT-H encoder to a smaller encoder while retaining the original mask-decoder architecture.The original decoder can serve as a readily usable component rather than being trained from scratch.
- Decoder finetuning: Mask-decoder finetuning is optional because the student encoder’s image embeddings can be sufficiently close to the teacher’s embeddings for compatibility with the frozen decoder.Joint or decoder-only finetuning may still improve performance, but it is not required by the reported empirical finding.
- Preliminary evaluation: Decoupled distillation performs better than coupled distillation while requiring less than 1% of its computation resources in the ViT-B comparison.The comparison evaluates mask prediction using mIoU between teacher and student masks generated from the same prompt point.
4 Experiments
Experiments evaluate MobileSAM with a lightweight ViT-Tiny encoder across promptable and segment-everything settings, comparing its efficiency and mask quality with SAM and FastSAM. MobileSAM produces masks similar to original SAM while using substantially fewer resources and outperforming FastSAM in reported comparisons.
- 4.1 Experimental Setup: ViT-Tiny is adopted as the lightweight image encoder for demonstrating efficient MobileSAM training and inference.The encoder is selected as a proof of concept because it performs better than DeiT-Tiny on ImageNet-1K while having similar parameters.
- 4.1 Experimental Setup: mIoU is computed between MobileSAM and original-SAM masks, treating the original SAM prediction as ground truth.The evaluation reports results for point and box prompts, while text prompts are omitted because pretrained text-guided mask-decoder models are unavailable.
- 4.2 MobileSAM performs on par with the orignal SAM: MobileSAM produces satisfactory masks similar to original SAM for both single-point and box prompts.The point-prompt results are shown in Figure 4, and the box-prompt results are shown in Figure 5.
- 4.2 MobileSAM performs on par with the orignal SAM: Increasing batch size and training epochs improves MobileSAM performance under the ablation settings, with all experiments conducted on a single GPU.The ablation varies training computation through batch size and update iterations.
- 4.3 MobileSAM outperforms FastSAM: 68M parameters and 40ms per image make FastSAM substantially larger and slower than MobileSAM, which has fewer than 10M parameters and takes 10ms.The comparison is measured on a single GPU and reports MobileSAM as four times faster than FastSAM.
- 4.3 MobileSAM outperforms FastSAM: MobileSAM aligns more closely with original SAM than FastSAM in segment-anything and segment-everything comparisons.FastSAM shows lower mIoU, can miss objects when prompt points are close, and often produces less interpretable masks with non-smooth boundaries.
5 Conclusion
The paper proposes MobileSAM by replacing SAM’s heavyweight image encoder with a lightweight one trained through decoupled distillation. The resulting model remains comparable to original SAM while being substantially smaller and faster than FastSAM, supporting mobile deployment.
- 5 Conclusion: MobileSAM replaces SAM’s heavyweight image encoder with a lightweight encoder trained through decoupled distillation from ViT-H.The distilled encoder is automatically compatible with the original SAM mask decoder.
- 5 Conclusion: More than 60 times smaller, MobileSAM performs on par with original SAM and is reported as four times faster and seven times smaller than FastSAM.The model preserves the original SAM pipeline and replaces only the image encoder, enabling plug-and-play use in existing SAM-based projects.