VSSD: Vision Mamba with Non-Causal State Space Duality

📄 arXiv: 2407.18559v2 📥 PDF

作者: Yuheng Shi, Minjing Dong, Mingjia Li, Chang Xu

分类: cs.CV

发布日期: 2024-07-26 (更新: 2024-08-04)

备注: 16 pages, 5 figures, 7 tables

🔗 代码/项目: GITHUB

💡 一句话要点

提出VSSD：一种非因果状态空间对偶视觉Mamba模型，提升视觉任务性能。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 视觉状态空间对偶 非因果模型 状态空间模型 视觉Mamba 图像分类 目标检测 图像分割

📋 核心要点

  1. 现有视觉Transformer计算量大，难以处理长序列，而传统SSM的因果性限制了其在非因果视觉任务中的应用。
  2. VSSD通过放弃隐藏状态和tokens之间交互的幅度，保留相对权重，解除了token贡献对先前token的依赖，实现非因果性。
  3. 实验表明，VSSD在图像分类、检测和分割等任务上超越了现有的基于SSM的最先进模型，提升了性能和效率。

📝 摘要（中文）

视觉Transformer在计算机视觉领域取得了显著进展，提供了强大的建模能力和全局感受野。然而，其高计算需求限制了它们在处理长序列中的应用。为了解决这个问题，状态空间模型（SSM）因其线性计算复杂度而在视觉任务中获得了突出地位。最近，状态空间对偶（SSD）作为SSM的改进变体，在Mamba2中被引入，以提高模型性能和效率。然而，SSD/SSM固有的因果性质限制了它们在非因果视觉任务中的应用。为了解决这个限制，我们引入了视觉状态空间对偶（VSSD）模型，它具有SSD的非因果形式。具体来说，我们建议放弃隐藏状态和tokens之间交互的幅度，同时保留它们的相对权重，这减轻了token贡献对先前token的依赖。结合多扫描策略的应用，我们表明扫描结果可以被整合以实现非因果性，这不仅提高了SSD在视觉任务中的性能，而且提高了其效率。我们在包括图像分类、检测和分割在内的各种基准上进行了广泛的实验，其中VSSD超越了现有的最先进的基于SSM的模型。

🔬 方法详解

问题定义：论文旨在解决状态空间模型（SSM）及其改进版本（如SSD）在视觉任务中因其固有的因果性而受限的问题。现有方法无法有效处理需要全局上下文信息的非因果视觉任务，例如图像分割和目标检测，因为它们依赖于序列的顺序处理，忽略了未来信息对当前token的影响。

核心思路：VSSD的核心思路是通过解除token贡献对先前token的依赖，实现非因果性。具体来说，VSSD放弃了隐藏状态和tokens之间交互的幅度，而只保留它们的相对权重。这样，每个token的贡献不再受到其之前token的影响，从而允许模型利用全局上下文信息。

技术框架：VSSD模型基于状态空间对偶（SSD）框架，并对其进行了修改以实现非因果性。其整体架构包括输入tokens的嵌入层、VSSD模块以及输出层。VSSD模块是核心，它使用非因果状态空间更新机制来处理tokens序列。此外，VSSD还采用了多扫描策略，即从不同方向扫描输入序列，并将扫描结果整合，以进一步增强非因果性。

关键创新：VSSD最重要的技术创新点在于其非因果状态空间更新机制。与传统的因果SSM不同，VSSD通过放弃隐藏状态和tokens之间交互的幅度，保留相对权重，从而解除了token贡献对先前token的依赖。这种设计使得模型能够利用全局上下文信息，从而更好地处理非因果视觉任务。

关键设计：VSSD的关键设计包括：1) 非因果状态空间更新公式，该公式放弃了隐藏状态和tokens之间交互的幅度，保留相对权重；2) 多扫描策略，通过从不同方向扫描输入序列，并将扫描结果整合，增强非因果性；3) 损失函数的设计，可能包括交叉熵损失（用于分类任务）或IoU损失（用于分割任务），以优化模型性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

VSSD在图像分类、目标检测和图像分割等多个基准数据集上进行了实验，结果表明VSSD超越了现有的最先进的基于SSM的模型。具体的性能数据和提升幅度在论文中详细给出，证明了VSSD在视觉任务中的有效性和优越性。例如，VSSD在ImageNet图像分类任务上取得了显著的精度提升。

🎯 应用场景

VSSD模型具有广泛的应用前景，可应用于图像分类、目标检测、图像分割等多种视觉任务。其非因果建模能力使其在需要全局上下文信息的场景中表现出色，例如医学图像分析、遥感图像处理等。VSSD的未来发展方向包括探索更有效的非因果建模方法、优化模型结构以提高效率，以及将其应用于更多复杂的视觉任务。

📄 摘要（原文）

Vision transformers have significantly advanced the field of computer vision, offering robust modeling capabilities and global receptive field. However, their high computational demands limit their applicability in processing long sequences. To tackle this issue, State Space Models (SSMs) have gained prominence in vision tasks as they offer linear computational complexity. Recently, State Space Duality (SSD), an improved variant of SSMs, was introduced in Mamba2 to enhance model performance and efficiency. However, the inherent causal nature of SSD/SSMs restricts their applications in non-causal vision tasks. To address this limitation, we introduce Visual State Space Duality (VSSD) model, which has a non-causal format of SSD. Specifically, we propose to discard the magnitude of interactions between the hidden state and tokens while preserving their relative weights, which relieves the dependencies of token contribution on previous tokens. Together with the involvement of multi-scan strategies, we show that the scanning results can be integrated to achieve non-causality, which not only improves the performance of SSD in vision tasks but also enhances its efficiency. We conduct extensive experiments on various benchmarks including image classification, detection, and segmentation, where VSSD surpasses existing state-of-the-art SSM-based models. Code and weights are available at \url{https://github.com/YuHengsss/VSSD}.