[논문 리뷰] SAM-Lightening: A Lightweight Segment Anything Model with Dilated Flash Attention to Achieve 30 times Acceleration

Yanfei Songa, Bangzheng Pua외 5명|arXiv (Cornell University)|2024. 03. 14.

Data Visualization and Analytics인용 수 5

한 줄 요약

SAM-Lightening re-engineers SAM's encoder with a Dilated Flash Attention and dynamic layer-wise distillation to deliver ~30× faster inference and dramatically lower memory while preserving segmentation quality.

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/.

연구 동기 및 목표

• 실제_world_배치에서 Segment Anything Model (SAM)의 계산 병목 현상을 해결한다.
• Dilated Flash Attention을 사용한 효율적인 이미지 인코더를 도입하여 추론 속도와 메모리 감소를 달성한다.
• Dynamic Layer-Wise Distillation (DLD)을 제안하여 처음부터 학습하지 않고 vanilla SAM으로부터 지식을 전이한다.
• SAM-Lightening이 COCO와 LVIS에서 경쟁력 있는 분할 성능을 유지하면서도 훨씬 더 효율적임을 입증한다.

제안 방법

• Dilated Flash Attention 메커니즘을 설계하여 vanilla self-attention을 대체하고 더 높은 병렬성과 희소화된 세그먼트를 구현한다.
• Dynamic Layer-Wise Distillation (DLD)을 적용하여 SAM으로부터 경량 인코더로 지식을 점진적으로 전이한다.
• 출력에 가까운 더 깊은 계층에 초점을 맞춘 decoupled feature distillation를 사용하여 표현을 교사와 정렬한다.
• 프롬프트 기반 분할 작업을 위한 경량 인코더와의 정합을 위해 디코더를 미세조정한다.
• 1% SA-1B 데이터로 학습하고, 증류를 가속하기 위해 SAM 엔코더 출력을 저장한 뒤 표준 벤치마크에서 평가한다.

실험 결과

연구 질문

• RQ1 재구성된 주의 메커니즘(Dilated Flash Attention)이 SAM의 엔코더를 정확도를 잃지 않고 가속시킬 수 있는가?
• RQ2 Dynamic Layer-Wise Distillation가 SAM으로부터 경량 인코더로 지식을 효과적으로 전이하는가?
• RQ3 COCO와 LVIS에서 SAM-Lightening의 속도, 메모리, 분할 성능 간의 트레이드오프는 어떠한가?
• RQ4 Box, 1P, 3P와 Anything 모드에서 다양한 프롬프트 하에서 SAM-Lightening은 최신 SAM 변종과 어떻게 비교되는가?

주요 결과

• SAM-Lightening은 1024×1024 입력에서 이미지당 7 ms를 달성하여 vanilla SAM 대비 30.1×, 최신 기법 대비 2.1× 빠릅니다.
• 메모리 사용량은 224 MB로, vanilla SAM의 약 3.5%에 불과합니다.
• 추론 지연시간과 메모리 효율은 프롬프트 및 데이터세트 전반에서 경쟁력 있는 경량 SAM 변종보다 우수합니다.
• SAM-Lightening은 Box 및 Point 프롬프트 하에서 COCO와 LVIS에서 vanilla SAM과 비교 가능한 세분화 성능을 유지합니다.
• Dynamic Layer-Wise Distillation은 점진적 계층 가중치 부여와 출력 근처의 심층 특징 층에 대한 집중을 통해 지식을 효율적으로 전이합니다.
• Point/Box 프롬프트로 디코더를 미세 조정하면 경량 인코더와 고정된 디코더 간의 정합이 개선됩니다.