[논문 리뷰] Mamba YOLO: A Simple Baseline for Object Detection with State Space Model

Zeyu Wang, Chen Li외 3명|arXiv (Cornell University)|2024. 06. 09.

Computer Science and Engineering인용 수 37

한 줄 요약

Mamba-YOLO는 글로벌 의존성을 선형 복잡도로 포착하기 위해 YOLO에 상태공간모델(SSM)을 통합하고, 로컬 특성 모델링을 강화하기 위해 LSBlock과 RGBlock을 추가하여 다양한 크기의 COCO/VOC에서 강력한 성능을 달성합니다.

ABSTRACT

Driven by the rapid development of deep learning technology, the YOLO series has set a new benchmark for real-time object detectors. Additionally, transformer-based structures have emerged as the most powerful solution in the field, greatly extending the model's receptive field and achieving significant performance improvements. However, this improvement comes at a cost as the quadratic complexity of the self-attentive mechanism increases the computational burden of the model. To address this problem, we introduce a simple yet effective baseline approach called Mamba YOLO. Our contributions are as follows: 1) We propose that the ODMamba backbone introduce a extbf{S}tate extbf{S}pace extbf{M}odel ( extbf{SSM}) with linear complexity to address the quadratic complexity of self-attention. Unlike the other Transformer-base and SSM-base method, ODMamba is simple to train without pretraining. 2) For real-time requirement, we designed the macro structure of ODMamba, determined the optimal stage ratio and scaling size. 3) We design the RG Block that employs a multi-branch structure to model the channel dimensions, which addresses the possible limitations of SSM in sequence modeling, such as insufficient receptive fields and weak image localization. This design captures localized image dependencies more accurately and significantly. Extensive experiments on the publicly available COCO benchmark dataset show that Mamba YOLO achieves state-of-the-art performance compared to previous methods. Specifically, a tiny version of Mamba YOLO achieves a extbf{7.5}\% improvement in mAP on a single 4090 GPU with an inference time of extbf{1.5} ms. The pytorch code is available at: \url{https://github.com/HZAI-ZJNU/Mamba-YOLO}

연구 동기 및 목표

• YOLO에서 글로벌 컨텍스트 이해를 개선하기 위해 State Space Models(SSMs)를 활용한 경량 실시간 탐지기를 고안한다.
• SSM 백본을 기반으로 로컬 이미지 국지성과 채널 특징을 다루는 새로운 블록(ODSSBlock, LSBlock, RGBlock)을 설계한다.
• COCO와 VOC 데이터셋에서 Tiny/Base/Large 등 여러 스케일에서 강력한 baseline으로 Mamba-YOLO를 확립한다.
• 객체 탐지에서 속도-정확도 트레이드오프에서 SSM 기반 백본이 전통적인 CNN/트랜스포머 하이브리드보다 우수할 수 있음을 입증한다.

제안 방법

• ODMamba 백본 내에 ODSSBlock 코어를 채택하여 SS2D 기반의 글로벌 모델링과 로컬 컨볼루션을 융합한다.
• 깊이별 분리 합성곱과 잔여 융합을 통해 미세한 로컬 정보를 포착하기 위한 LocalSpatial Block(LSBlock)을 도입한다.
• 잔류 연결과 1x1 합성을 통한 게이트화된 집계를 구현하는 ResGated Block(RGBlock)을 도입하여 효율적인 글로벌-로컬 특징 혼합을 수행한다.
• 목과(네크)에서 C2f/표준 다운샘플링을 ODSSBlock 유도 다운샘플링과 Vision Clue Merge로 대체하여 더 풍부한 SS2D 가이던스를 제공한다.
• 연속 상태 SSM(A,B)을 Zero-Order Hold를 통해 이산화된 Ā,B̄로 모델링하여 효율적인 학습/추론을 구현한다.
• Tiny/Base/Large 계열의 모델을 제공하고 COCO와 VOC에서 평가하여 확장성과 경쟁력을 보여준다.

실험 결과

연구 질문

• RQ1 YOLO 백본에 State Space Models를 통합하면 실시간 추론을 희생하지 않으면서 탐지 정확도를 향상시킬 수 있는가?
• RQ2 LSBlock과 RGBlock이 SS2D 기반 아키텍처에서 잠재적인 로컬 이미지 모델링의 격차를 효과적으로 보완하는가?
• RQ3 ODSS 기반 설계가 COCO와 VOC에서 다양한 모델 규모에 걸쳐 일관된 이점을 제공하는가?

주요 결과

Method	AP^val(%)	AP50(%)	AP75(%)	AP_S(%)	AP_M(%)	AP_L(%)	#param.	FLOPs
YOLOv5-N	28.0	45.7					1.9 M	4.5 G
YOLOv5-S	37.4	56.8					7.2 M	16.5 G
YOLOv5-M	45.4	64.1					21.2 M	49.0 G
YOLOv5-L	49.0	67.3					46.5 M	109.1 G
YOLOv6-3.0-N	37.0	52.7					4.7 M	4.7 G
YOLOv6-3.0-S	44.3	61.2					4.7 M	45.3 G
YOLOv6-3.0-M	49.1	66.1					85.8 M	85.8 G
YOLOv6-3.0-L	51.8	69.2					59.6 M	150.7 G
YOLOv7-Tiny	37.4	55.2	37.3	15.7	38.0	53.4	6.2 M	13.7 G
YOLOv7	51.2	69.7	55.9	31.8	55.5	65.0	36.9 M	104.7 G
YOLOv7-X	52.9	71.7	51.4	36.9	57.7	68.6	71.3 M	189.9 G
YOLOv8-N	37.3	52.6	40.6	18.8	41.0	53.5	3.2 M	8.7 G
YOLOv8-S	44.9	61.8	48.6	26.0	49.9	61.0	11.2 M	28.6 G
YOLOv8-M	50.2	67.3	54.8	32.3	55.9	66.5	25.9 M	78.9 G
YOLOv8-L	52.9	69.8	57.7	35.5	58.5	69.8	43.7 M	165.2 G
Gold-YOLO-N	39.6	55.7		19.7	44.1	57.0	5.6 M	12.1 G
Gold-YOLO-S	45.4	62.5		25.3	50.2	62.6	21.5 M	46.0 G
Gold-YOLO-M	49.8	67.0		32.3	55.3	66.3	41.3 M	87.5 G
Gold-YOLO-L	51.8	68.9		34.1	57.4	68.2	75.1 M	151.7 G
Mamba YOLO-T	45.4	62.3	49.1	25.2	50.4	62.9	6.1 M	14.3 G
Mamba YOLO-B	49.9	67.2	54.4	30.6	55.4	67.0	21.8 M	49.7 G
Mamba YOLO-L	52.1	69.8	56.5	34.1	57.3	68.1	57.6 M	156.2 G

• Mamba YOLO-T/L은 여러 기준선 대비 파라미터 및 FLOPs 감소를 동반하며 COCOval에서 경쟁력 있는 AP를 달성한다.
• Mamba YOLO-T는 최고급 작은 경량 모델 DAMO YOLO-T/YOLO MS-XS보다 AP 3.4% 및 AP50 2.0% 향상을 보이며, YOLOv8-S에 비해 파라미터 및 FLOPs가 현저히 적다.
• Mamba YOLO-L은 AP 52.1, AP50 69.8, AP75 56.5를 달성하며, 비슷한 정확도에서 Gold-YOLO-L에 비해 파라미터와 FLOPs를 감소시킨다.
• Mamba YOLO는 COCO에서 Tiny/Base/Large 변형 간에 우수한 FLOPs-정확도 및 파라미터-정확도 트레이드오프를 보여주고, 여러 최첨단 검출기를 능가한다.