通过空间维度池化加速视觉Mamba，推理速度提升72.5%
Fast Vision Mamba: Pooling Spatial Dimensions for Accelerated Processing
- 交替对图像维度进行池化，减少Mamba块的并行步骤数
- 高分辨率图像下推理速度最高提升72.5%，性能不降
- 适合需要高速处理高分辨图像的任务场景
状态空间模型（SSM）结合选择性扫描（Mamba）已被用于构建高效的视觉模型。与视觉变压器不同，Mamba通过递归隐状态过程实现线性复杂度的标记交互。其序列化处理通过并行扫描算法优化，将递归步骤从$ L $个串行步骤降至$ \log(L) $个并行步骤。本文提出Fast Vision Mamba（FastVim），在保持模型性能的前提下，进一步减少视觉Mamba模型中递归步骤数量。通过在Mamba块间交替对图像维度进行池化，使SSM块的并行步骤数减少2倍。在2048×2048高分辨率图像上，相比基线视觉Mamba模型，推理速度最高提升72.5%。实验表明，该模型在图像分类、细胞扰动预测、分割和目标检测等任务中达到顶尖性能，且吞吐量显著提升。代码已开源于https://github.com/insitro/FastVim。
原文摘要 · Abstract (English)
State Space Models (SSMs) with selective scan (Mamba) have been adapted into efficient vision models. Mamba, unlike Vision Transformers, achieves linear complexity for token interactions through a recurrent hidden state process. This sequential processing is enhanced by a parallel scan algorithm, which reduces the computational time of recurrent steps from $L$ sequential steps to $log(L)$ parallel steps with respect to the number of input tokens ($L$). In this work, we propose Fast Vision Mamba (FastVim), that further reduces the computational time of the SSM block by reducing the number of recurrent steps in Vision Mamba models while still retaining model performance. By alternately pooling tokens along image dimensions across Mamba blocks, we obtain a 2$\times$ reduction in the number of parallel steps in SSM block. Our model offers up to $72.5\%$ speedup in inference speed compared to baseline Vision Mamba models on high resolution (2048$\times$2048) images. Our experiments demonstrate state-of-the-art performance with dramatically improved throughput in a range of tasks such as image classification, cell perturbation prediction, segmentation, and object detection. Code is made available at https://github.com/insitro/FastVim
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。