[Paper Review] Faster Segment Anything: Towards Lightweight SAM for Mobile Applications

Chaoning Zhang, Dongshen Han+5 more|arXiv (Cornell University)|Jun 25, 2023

Advanced Neural Network ApplicationsComputer Science141 citations

TL;DR

The paper presents MobileSAM, a lightweight, mobile-friendly variant of Segment Anything by decoupled distillation to replace the heavyweight image encoder, achieving on-par segmentation with much smaller size and faster inference than the original SAM and FastSAM.

ABSTRACT

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}{ extcolor{red}{MobileSAM}}), with a demo showing that MobileSAM can run relatively smoothly on CPU.

Motivation & Objective

• Motivate mobile deployment of SAM for resource-constrained devices.
• Reduce model size and improve inference speed by replacing the heavy image encoder.
• Maintain compatibility with the original SAM mask decoder without extensive retraining.
• Demonstrate that distillation-based training can yield a lightweight yet accurate segmentation model.

Proposed method

• Replace the ViT-H image encoder in SAM with a lightweight encoder via knowledge distillation from the teacher ViT-H.
• Apply decoupled distillation to transfer knowledge to a small student encoder while keeping the original mask decoder frozen or lightly finetuned.
• Use an MSE loss for image embedding alignment instead of the combined focal and dice losses used in related work.
• Optionally finetune the mask decoder, but show that decoupled distillation already yields good alignment between student encoder and decoder.
• Evaluate MobileSAM against the original SAM and FastSAM using mIoU and inference speed metrics.

Experimental results

Research questions

• RQ1 Can a lightweight image encoder distilled from the heavy SAM encoder achieve comparable segmentation quality to the original SAM?
• RQ2 Does decoupled distillation outperform coupled or semi-coupled distillation for training a lightweight SAM?
• RQ3 How does MobileSAM compare to FastSAM in terms of accuracy (mIoU) and efficiency (parameters, speed) on segment-anything tasks?
• RQ4 Is MobileSAM capable of running efficiently on CPU for on-device applications?

Key findings

• MobileSAM reduces encoder parameters by about 100x and total parameters by about 60x, while achieving comparable performance to the original SAM.
• On a single GPU, MobileSAM processes an image in about 10 ms (8 ms for the encoder and 4 ms for the decoder).
• MobileSAM is approximately 5x faster and 7x smaller than FastSAM, while delivering superior performance in the segment-anything setting.
• Decoupled distillation (training the lightweight encoder directly from the heavy teacher without joint decoder training) yields better mIoU (0.75) than coupled distillation (0.72) in preliminary experiments.
• MobileSAM can run relatively smoothly on CPU, enabling mobile device deployment.