[논문 리뷰] EdgeSAM: Prompt-In-the-Loop Distillation for SAM

Chong Zhou, Xiangtai Li외 2명|arXiv (Cornell University)|2023. 12. 11.

Advanced Memory and Neural Computing참고 문헌 61인용 수 11

한 줄 요약

EdgeSAM은 SAM을 CNN 기반 인코더로 증류하고 프롬프트-루프(distillation) 및 경량 그레나랄리티 프라이어 모듈을 도입해 경쟁력 있는 정확도로 실시간 기기 내 대화형 세분화를 달성합니다.

ABSTRACT

This paper presents EdgeSAM, an accelerated variant of the Segment Anything Model (SAM), optimized for efficient execution on edge devices with minimal compromise in performance. Our approach involves distilling the original ViT-based SAM image encoder into a purely CNN-based architecture, better suited for edge devices. We carefully benchmark various distillation strategies and demonstrate that task-agnostic encoder distillation fails to capture the full knowledge embodied in SAM. To overcome this bottleneck, we include both the prompt encoder and mask decoder in the distillation process, with box and point prompts in the loop, so that the distilled model can accurately capture the intricate dynamics between user input and mask generation. To mitigate dataset bias issues stemming from point prompt distillation, we incorporate a lightweight module within the encoder. As a result, EdgeSAM achieves a 37-fold speed increase compared to the original SAM, and it also outperforms MobileSAM/EfficientSAM, being over 7 times as fast when deployed on edge devices while enhancing the mIoUs on COCO and LVIS by 2.3/1.5 and 3.1/1.6, respectively. It is also the first SAM variant that can run at over 30 FPS on an iPhone 14. Code and demo are available at https://www.mmlab-ntu.com/project/edgesam.

연구 동기 및 목표

• 대규모 계산 비용 없이 엣지 디바이스에서 SAM과 유사한 기능을 가능하게 하여 기기 내 대화형 세분화를 고무한다.
• 인코더 증류 전략을 조사하고 태스크-무관 접근 방식의 한계를 입증한다.
• 동적 프롬프트-인-더-루프 메커니즘을 갖춘 프롬프트 인식 지식 증류를 제안한다.
• 데이터셋별 프롬프트 세분성을 다루기 위해 경량 그레나랄리티 프라이어 모듈을 포함한다.
• 표준 분할 벤치마크에서 제로샷 전이 및 실시간 성능을 평가한다.

제안 방법

• SAM의 ViT 기반 이미지 인코더를 에지 배치를 위한 순수 CNN 기반 아키텍처로 증류한다.
• 인코더-전용 지식 증류를 적용하고 태스크 특이적 가이던스 없이 그 한계를 보인다.
• 학생 모델이 잘못한 영역에서 프롬프트를 반복적으로 샘플링하여 마스크 디코더를 안내하는 동적 프롬프트-인-더-루프 증류를 도입한다.
• SAM의 마스크 디코더를 유지하고 디코더 출력에 초점을 맞춘 증류 타깃으로 학습한다.
• 추론 시 데이터세트 특유의 세분성 프라이어를 포함하기 위해 경량 그레나랄리티 프라이어 모듈(RPN with FPN)을 추가한다.
• 계단식 학습을 수행한다: 인코더-전용 KD, 프롬프트-인-더-루프 KD, 그리고 선택적으로 경량 RPN을 이용한 그라운딩.

실험 결과

연구 질문

• RQ1 SAM을 엣지 디바이스에 효과적으로 배포할 수 있을까, CNN 기반 인코더로 구성 요소를 증류하되 세분화 품질의 큰 손실 없이?
• RQ2 증류 중 태스크 인지형 프롬프트 가이던스가 인코더-전용 증류보다 밀집 예측 태스크에서 우수한가?
• RQ3 프롬프트 유형과 그레나랄리티 프라이어가 제로샷 전이 및 엣지 디바이스 성능에 어떤 영향을 주는가?
• RQ4 엣지 하드웨어에서 처리량과 정확도의 균형을 가장 잘 맞추는 백본/백본 수정은 무엇인가?
• RQ5 경량 RPN이 세분성 프라이어를 활용해 모호한 프롬프트의 성능을 개선할 수 있는가?

주요 결과

• EdgeSAM은 iPhone 14에서 30 FPS 이상으로 실행되며, 실시간 엣지 디바이스 성능을 갖춘 최초의 SAM 변형이다.
• 인코더-전용 증류만으로는 원래의 SAM 대비 성능 격차가 발생하여 프롬프트-인-더-루프 증류의 필요성을 입증한다.
• 프롬프트-인-더-루프 증류는 추가 정제 프롬프트와 함께 인코더-전용 KD보다 마스크 정확도를 향상시키며 향상된 정제 성능을 제공한다.
• 순수 CNN 기반 백본(RepViT-M1 with FPN)이 엣지 배포를 위한 최고의 처리량-정확도 균형을 제공한다.
• EdgeSAM은 엣지 디바이스에서 SAM보다 최대 40배, MobileSAM보다 최대 14배 속도 향상을 달성하며 COCO와 LVIS에서 경쟁력 있거나 더 나은 mIoU/IoU 이득을 보인다.
• GT 박스를 프롬프트로 사용할 때 EdgeSAM은 SAM과의 간극을 좁히고 여러 설정에서 MobileSAM을 능가하며, 그레나랄리티 프라이어 RPN은 COCO에서 중심점 성능을 추가로 높인다.