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

Segment Anything Model (SAM)因其零样本泛化能力在分割任务中受到了广泛关注。然而，由于其注意力机制导致的推理速度低和计算内存需求高，限制了SAM在实际应用中的更广泛应用。现有工作主要集中在优化编码器，但并没有充分解决注意力机制本身的低效率问题，即使将其压缩到较小的模型中，也仍有改进的空间。因此，我们引入了SAM-Lightening，这是SAM的一个变体，具有重新设计的注意力机制，称为Dilated Flash Attention。它不仅促进了更高的并行性，提高了处理效率，而且还保持了与现有的FlashAttention的兼容性。相应地，我们提出了一种渐进式蒸馏方法，使得可以从原始的SAM中进行高效的知识转移，而无需昂贵的从头开始训练。在COCO和LVIS上的实验证明，SAM-Lightening在运行时效率和分割精度方面显著优于现有的最先进方法。具体而言，它可以在每张大小为1024 * 1024像素的图像上实现7毫秒（ms）的推理速度，比原始SAM快30.1倍，比最先进的方法快2.1倍。此外，它只需要244MB的内存，仅为原始SAM的3.5％。代码和权重可在https://anonymous.4open.science/r/SAM-LIGHTENING-BC25/上获得。

SAM模型在分割任务中具有零样本泛化能力，但由于其注意机制导致推理速度慢，计算内存需求高，限制了其在实际应用中的广泛应用。本论文旨在解决SAM模型的注意机制效率问题。

本论文提出了一种改进的注意机制，称为Dilated Flash Attention，可以提高并行性，增强处理效率，并且与现有的FlashAttention兼容。此外，还提出了渐进蒸馏方法，使得从原始SAM模型中进行高效的知识迁移成为可能。

SAM-Lightening在COCO和LVIS数据集上的实验结果表明，其在运行时效率和分割精度方面均显著优于现有方法。它可以实现每张1024*1024像素大小的图像的推理速度为7毫秒，比原始SAM快30.1倍，比现有最先进的方法快2.1倍。此外，它只需要244MB的内存，仅为原始SAM的3.5％。作者还提供了代码和权重。

最近的相关研究包括：《Segmentation Transformer: What If Transformers Consume the Image》、《Fastformer: Additive Attention Can Be All You Need》等。