Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SAM-Lightening: A Lightweight Segment Anything Model with Dilated Flash Attention to Achieve 30 times Acceleration

Yanfei Song, Bangzheng Pu, Peng Wang, Hongxu Jiang, Dong Dong, Yongxiang Cao, Yiqing Shen

TL;DR

This work tackles the bottleneck of high computational cost in Segment Anything Model (SAM) due to attention by introducing SAM-Lightening, which replaces standard attention with Dilated Flash Attention and enables efficient knowledge transfer via Dynamic Layer-Wise Distillation (DLD). The approach achieves substantial runtime and memory gains (e.g., 7 ms per 1024×1024 image and 244 MB memory) while preserving segmentation accuracy comparable to vanilla SAM on COCO and LVIS across box and point prompts. The key innovations are a sparsified, parallelizable attention mechanism with output reassembly, and a progressively weighted distillation process that focuses on deeper teacher layers and aligns the decoder, enabling training from the vanilla SAM without full retraining. These improvements yield a practical, deployment-friendly SAM variant that can run on resource-constrained devices and be complementary to pruning or quantization techniques, broadening SAM’s real-world applicability.

Abstract

Segment Anything Model (SAM) has garnered significant attention in segmentation tasks due to their zero-shot generalization ability. However, a broader application of SAMs to real-world practice has been restricted by their low inference speed and high computational memory demands, which mainly stem from the attention mechanism. Existing work concentrated on optimizing the encoder, yet has not adequately addressed the inefficiency of the attention mechanism itself, even when distilled to a smaller model, which thus leaves space for further improvement. In response, we introduce SAM-Lightening, a variant of SAM, that features a re-engineered attention mechanism, termed Dilated Flash Attention. It not only facilitates higher parallelism, enhancing processing efficiency but also retains compatibility with the existing FlashAttention. Correspondingly, we propose a progressive distillation to enable an efficient knowledge transfer from the vanilla SAM without costly training from scratch. Experiments on COCO and LVIS reveal that SAM-Lightening significantly outperforms the state-of-the-art methods in both run-time efficiency and segmentation accuracy. Specifically, it can achieve an inference speed of 7 milliseconds (ms) per image, for images of size 1024*1024 pixels, which is 30.1 times faster than the vanilla SAM and 2.1 times than the state-of-the-art. Moreover, it takes only 244MB memory, which is 3.5\% of the vanilla SAM. The code and weights are available at https://anonymous.4open.science/r/SAM-LIGHTENING-BC25/.

SAM-Lightening: A Lightweight Segment Anything Model with Dilated Flash Attention to Achieve 30 times Acceleration

TL;DR

This work tackles the bottleneck of high computational cost in Segment Anything Model (SAM) due to attention by introducing SAM-Lightening, which replaces standard attention with Dilated Flash Attention and enables efficient knowledge transfer via Dynamic Layer-Wise Distillation (DLD). The approach achieves substantial runtime and memory gains (e.g., 7 ms per 1024×1024 image and 244 MB memory) while preserving segmentation accuracy comparable to vanilla SAM on COCO and LVIS across box and point prompts. The key innovations are a sparsified, parallelizable attention mechanism with output reassembly, and a progressively weighted distillation process that focuses on deeper teacher layers and aligns the decoder, enabling training from the vanilla SAM without full retraining. These improvements yield a practical, deployment-friendly SAM variant that can run on resource-constrained devices and be complementary to pruning or quantization techniques, broadening SAM’s real-world applicability.

Abstract

Segment Anything Model (SAM) has garnered significant attention in segmentation tasks due to their zero-shot generalization ability. However, a broader application of SAMs to real-world practice has been restricted by their low inference speed and high computational memory demands, which mainly stem from the attention mechanism. Existing work concentrated on optimizing the encoder, yet has not adequately addressed the inefficiency of the attention mechanism itself, even when distilled to a smaller model, which thus leaves space for further improvement. In response, we introduce SAM-Lightening, a variant of SAM, that features a re-engineered attention mechanism, termed Dilated Flash Attention. It not only facilitates higher parallelism, enhancing processing efficiency but also retains compatibility with the existing FlashAttention. Correspondingly, we propose a progressive distillation to enable an efficient knowledge transfer from the vanilla SAM without costly training from scratch. Experiments on COCO and LVIS reveal that SAM-Lightening significantly outperforms the state-of-the-art methods in both run-time efficiency and segmentation accuracy. Specifically, it can achieve an inference speed of 7 milliseconds (ms) per image, for images of size 1024*1024 pixels, which is 30.1 times faster than the vanilla SAM and 2.1 times than the state-of-the-art. Moreover, it takes only 244MB memory, which is 3.5\% of the vanilla SAM. The code and weights are available at https://anonymous.4open.science/r/SAM-LIGHTENING-BC25/.
Paper Structure (10 sections, 6 equations, 4 figures, 3 tables)

This paper contains 10 sections, 6 equations, 4 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related work
  3. Methods
  4. Dilated Flash Attention
  5. Dynamic Layer-Wise Distillation (DLD)
  6. Experiment
  7. Experimental Setups
  8. Results
  9. Ablation study
  10. Conclusion

Figures (4)

  • Figure 1: The overall framework of SAM-Lightening along with the dynamic layer-wise distillation that can efficiently transfer knowledge from the vanilla SAM without training from scratch.
  • Figure 2: Representative image segmentation results between SAM-Lightening and the vanilla SAM in prompt mode.
  • Figure 3: Representative samples under anything mode.
  • Figure 4: Impacts of inference time with FlashAttention over input size, where we select two embedding dimensions, namely 768 and 384, for comparison.