arXiv:2410.06806cs.CV2024-10NeurIPS被引 36

用四叉树自适应划分图像块,提升视觉状态空间模型的局部建模能力。

QuadMamba: Learning Quadtree-based Selective Scan for Visual State Space Model

  • 基于四叉树动态划分图像区域,保留2D空间邻近性。
  • 在多个视觉任务上达到当前最优性能,如分类与分割。
  • 适合关注高效视觉建模与长程依赖捕捉的研究者。

近期状态空间模型(如Mamba)在降低计算复杂度方面超越了主流Transformer,但将Mamba应用于视觉任务仍面临挑战:图像具有空间局部性和不同粒度的信息分布。现有方法要么采用光栅扫描打平像素,破坏空间邻接性;要么手动分块,限制长程建模能力。为此,我们提出QuadMamba,通过四叉树自适应分区与扫描机制,有效捕捉多粒度局部依赖。其轻量级四叉树扫描模块根据特征估计每个视觉令牌的局部性得分,并自适应地将其划分为窗口四象限。同时引入全向窗口移位策略,以捕获更完整和信息丰富的局部特征。为实现离散四叉树划分的端到端可训练性,我们设计基于Gumbel-Softmax的序列掩码策略及其直通梯度估计器。大量实验表明,QuadMamba在图像分类、目标检测、实例分割和语义分割等多个视觉任务中均取得领先性能。

原文摘要 · Abstract (English)

Recent advancements in State Space Models, notably Mamba, have demonstrated superior performance over the dominant Transformer models, particularly in reducing the computational complexity from quadratic to linear. Yet, difficulties in adapting Mamba from language to vision tasks arise due to the distinct characteristics of visual data, such as the spatial locality and adjacency within images and large variations in information granularity across visual tokens. Existing vision Mamba approaches either flatten tokens into sequences in a raster scan fashion, which breaks the local adjacency of images, or manually partition tokens into windows, which limits their long-range modeling and generalization capabilities. To address these limitations, we present a new vision Mamba model, coined QuadMamba, that effectively captures local dependencies of varying granularities via quadtree-based image partition and scan. Concretely, our lightweight quadtree-based scan module learns to preserve the 2D locality of spatial regions within learned window quadrants. The module estimates the locality score of each token from their features, before adaptively partitioning tokens into window quadrants. An omnidirectional window shifting scheme is also introduced to capture more intact and informative features across different local regions. To make the discretized quadtree partition end-to-end trainable, we further devise a sequence masking strategy based on Gumbel-Softmax and its straight-through gradient estimator. Extensive experiments demonstrate that QuadMamba achieves state-of-the-art performance in various vision tasks, including image classification, object detection, instance segmentation, and semantic segmentation. The code is in https://github.com/VISION-SJTU/QuadMamba.

视觉建模四叉树状态空间模型Mamba

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。