Mamba YOLO: A Simple Baseline for Object Detection with State Space Model (2406.05835v2)

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 \textbf{S}tate \textbf{S}pace \textbf{M}odel (\textbf{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 \textbf{7.5}\% improvement in mAP on a single 4090 GPU with an inference time of \textbf{1.5} ms. The pytorch code is available at: \url{https://github.com/HZAI-ZJNU/Mamba-YOLO}

• The paper introduces Mamba YOLO, an object detection framework using a state-space model (SSM) backbone for linear complexity, addressing efficiency issues of transformer models.
• Empirical evaluation shows Mamba YOLO-Tiny achieves a 7.5% mAP improvement on MSCOCO with a 1.5ms inference time on an Nvidia 4090 GPU, demonstrating real-time efficiency.
• Mamba YOLO exhibits competitive performance against established models like YOLOv8 and DETR, making it suitable for resource-constrained and real-time environments.

Overview of "Mamba YOLO: A Simple Baseline for Object Detection with State Space Model"

The paper presents Mamba YOLO, a novel object detection framework that effectively combines the advantages of the YOLO series and the state-space modeling approach, specifically utilizing the structured state-space model (SSM) concept. The motivation behind this work stems from the challenges presented by the quadratic complexity of transformer-based models, which, although powerful, suffer from computational inefficiencies. Mamba YOLO tackles this issue by employing a state-space model with linear complexity, offering an efficient alternative without requiring pre-training on large datasets.

Key Contributions

1. Linear Complexity with SSM: The proposed Mamba YOLO introduces a backbone that harnesses the state-space model (SSM) to address the non-linear complexity issues associated with conventional self-attention mechanisms in transformers. This transition aims to maintain performance while reducing computational demands.
2. Real-Time Structure Design: The designing phase involved optimizing the macro architecture of ODMamba, which includes determining the optimal stage ratios and scaling sizes suitable for real-time object detection scenarios.
3. Residual Gated (RG) Block: A multi-branch structure named RG Block is devised to boost the model's ability to capture localized dependencies in images while overcoming the limitations that SSM may present in sequence modeling. This design enhances the integration and processing of channel-based information within the network.

Empirical Evaluation

The evaluation is based on the MSCOCO dataset, where the Mamba YOLO model demonstrates state-of-the-art performance metrics. Notably, the Mamba YOLO-Tiny variant achieves a 7.5% improvement in mean Average Precision (mAP) on this dataset, coupled with an inference time of 1.5 milliseconds on a single Nvidia 4090 GPU, underscoring its efficiency for real-time applications.

Comparison to State-of-the-Art Models

Mamba YOLO competes impressively against established models, such as YOLOv6, YOLOv8, and DETR-based models like DINO. The results indicate Mamba YOLO's competitive advantage in both performance and computational efficiency, highlighting its potential application in scenarios where resource constraints are pertinent.

Theoretical and Practical Implications

Theoretically, Mamba YOLO advances the understanding of integrating state-space models within the domain of object detection, positioning itself as a viable alternative to the computationally intensive transformer frameworks. Practically, its deployment is particularly suitable for environments with stringent real-time processing requirements and limited computational resources.

Future Directions

While the paper establishes a solid foundation, future work could explore the extension of the Mamba framework to other tasks within computer vision, such as instance segmentation and keypoint detection. Additionally, further tuning and exploration of the RG Block can potentially unlock even more robust performance gains across varied datasets.

In conclusion, Mamba YOLO sets a new benchmark in the field of object detection by introducing an effective baseline that balances model complexity with performance, leveraging the state-space model concept. This approach not only promises enhancements in speed and accuracy but also signals a shift towards more computationally efficient methodologies within the broader AI research community.