论文学习笔记 VMamba: Visual State Space Model

概览

这篇论文的动机源于在计算机视觉领域设计计算高效的网络架构的持续需求。当前的视觉模型如卷积神经网络（CNNs）和视觉Transformer（ViTs）在处理大规模视觉任务时展现出良好的表现，但都存在各自的局限性。特别是，ViTs尽管在处理大规模数据上具有优势，但其自注意力机制的二次复杂度对高分辨率图像处理时的计算成本极高。因此，研究者希望通过引入新的架构来降低这种复杂度，并提高视觉任务的效率。

现有方法的不足：

1. CNNs虽然在提取局部特征方面表现优秀，但在捕捉长距离依赖时能力不足。
2. ViTs引入了自注意力机制，可以捕捉全局依赖，但自注意力的二次计算复杂度会带来显著的计算开销，尤其在处理高分辨率输入时。
3. 尽管一些方法（如线性注意力）试图通过降低自注意力的复杂度来优化性能，但往往导致感受野的缩小，或者在不同任务中的性能下降。

拟解决的问题：
该论文旨在提出一种新的视觉主干网络架构（VMamba），通过引入基于状态空间模型（SSM）的模块来高效处理视觉数据，减少计算复杂度，并在保持模型性能的同时提高推理速度。

主要贡献：

1. 提出了VMamba，这是一个基于SSM的视觉主干网络，能够在视觉表示学习中实现线性时间复杂度。
2. 引入了2D选择扫描（SS2D）模块，使得选择性SSM能够有效地处理视觉数据，弥合一维扫描与二维视觉数据之间的鸿沟。
3. 在多个视觉任务上展示了VMamba的优异性能，包括图像分类、目标检测和语义分割。尤其在处理大尺寸输入时，VMamba展现了线性增长的计算复杂度，具有显著的输入扩展性。

创新点：

1. 2D选择扫描（SS2D）模块：这是VMamba的核心创新，通过沿着四个扫描路径遍历图像的方式，使得每个图像块可以通过扫描路径上计算出的压缩隐状态来获取上下文信息，从而降低了计算复杂度。
2. 状态空间模型与视觉任务的结合：VMamba成功地将SSM应用于视觉任务，突破了SSM原本在自然语言处理中的一维顺序数据处理的局限，将其扩展到二维空间的视觉数据中。
3. 线性计算复杂度的实现：相比于ViTs的二次复杂度，VMamba通过SS2D和SSM的结合，实现了视觉数据处理中的线性计算复杂度，大幅降低了计算成本，特别是在高分辨率图像处理任务中具有优势。

总的来说，VMamba通过创新性的架构设计，提出了一种高效且性能强大的视觉主干网络，在保持全局感受野的同时，大幅降低了计算开销。

2D选择扫描（SS2D）模块是否属于一种注意力机制？

2D选择扫描（SS2D）模块虽然在功能上与注意力机制有相似之处，但它并不是传统意义上的注意力机制。

在传统的自注意力机制中，每个输入（如图像中的每个patch）会计算与所有其他输入的相关性权重，通过全局注意力获取信息。这种操作的计算复杂度通常是二次的（与输入的token数量有关），这也是视觉Transformer中高计算成本的主要原因。

而2D选择扫描（SS2D）模块的设计初衷是为了降低这种二次复杂度。SS2D通过沿着四个扫描路径遍历图像，构建一种线性时间复杂度的上下文信息获取方式。在这个过程中，每个图像patch通过这些路径的累积压缩状态来获取信息，而不是像注意力机制那样直接计算与所有其他patch的相关性。虽然这种方式也能够建立全局的感受野，但它通过更高效的扫描算法和递归状态更新来实现，而非显式的注意力矩阵计算。

因此，SS2D不属于严格意义上的注意力机制，但它确实在一定程度上实现了类似于注意力机制的功能，即从全局范围内收集上下文信息，同时显著降低了计算复杂度。

如何理解SS2D的selective？

SS2D中的**“selective”（选择性）**机制核心在于它如何通过扫描路径中的累积压缩状态来动态地选择和聚合相关的上下文信息。具体来说，SS2D的“选择性”体现在以下几个方面：

1. 基于路径的选择性信息聚合：

SS2D模块通过沿着四个不同的扫描路径（例如，从左到右、从上到下等）遍历图像的每个patch。在每个路径上，它逐步聚合扫描路径上先前位置的隐状态（压缩信息）。这种聚合是选择性的，因为它并不是简单地累积所有的信息，而是动态选择了与当前patch最相关的上下文信息。这与注意力机制中的权重选择相似，只不过这里的选择性聚合通过累积的隐状态来实现，而非显式计算每个patch之间的注意力权重。

2. 动态权重更新：

在每个扫描路径中，SS2D使用了类似状态空间模型（SSM）的递归机制。每个patch的输出不仅依赖于自身的信息，还取决于路径中之前所有位置的压缩隐状态。这个递归关系使得系统可以根据不同的输入数据，动态调整隐状态的累积过程，从而实现选择性信息传递。这种隐状态随着路径的推进不断更新，而这种更新本质上是选择性的，因为每一步只保留了当前步中最重要的信息。

3. 跨路径的选择性融合：

SS2D的最后一步是通过**跨路径融合（cross-merge）**将沿四个路径获得的信息整合到一起。不同路径上的信息提供了不同方向的上下文信息，最终输出是从多个方向上“选择”相关的信息进行融合，从而构建一个全局上下文。这一步的选择性体现在，不同的路径可能聚合了不同的局部和全局信息，最终的融合输出是这些选择性信息的结合。

总结：

SS2D中的选择性主要体现在其动态的上下文信息聚合过程。通过多个扫描路径和状态空间模型中的隐状态递归，系统能够选择性地保留最相关的上下文信息，并通过跨路径融合，最终得到丰富的、全局化的特征表示。相比传统的自注意力机制，SS2D的选择性机制更隐式、更高效，同时避免了二次复杂度的计算。

Vmamba不同stage输出的feature maps是否类似卷积神经网络具有浅层捕获低级特征能力强，深层捕获语义信息强的特点？

是的，VMamba的不同stage输出的feature maps与卷积神经网络（CNN）在浅层捕获低级特征、深层捕获高级语义特征的能力非常相似。具体来说：

1. 浅层（Stage 1和Stage 2）：
   浅层通常负责捕捉图像的低级特征，例如边缘、纹理、颜色、形状等。这些特征较为局部，关注图像中的细节部分。在VMamba的第一、第二个阶段（Stage 1和Stage 2），输入图像被划分成patch，经过VSS block的处理后，这些block能够通过选择性扫描机制聚合局部信息，从而捕捉低级的视觉特征。这与卷积神经网络的浅层操作相似，都是关注局部的、细粒度的特征。

2. 深层（Stage 3和Stage 4）：
   随着网络的加深，VMamba中的后续stage（Stage 3和Stage 4）经过更多次下采样，feature map的分辨率逐渐降低，而通道数量增加，这意味着模型具备了更大的感受野。此时，模型可以聚合更多的上下文信息，从而逐渐捕获图像中的高级语义信息，例如对象的整体形状、类别等。在这些深层中，特征的语义信息更丰富，类似于卷积神经网络深层提取高级抽象特征的过程。

总结来说，VMamba的结构在特征提取上确实与CNN类似，前几层主要捕获低级特征，后几层逐步提取更加抽象和语义丰富的特征。这种设计使得VMamba能够在不同层级上捕获从局部到全局、从低级到高级的多尺度特征。

Vision Transformer是通过self-attention来实现注意力机制的，那Vmamba是如何实现注意力机制的？

VMamba虽然没有采用传统的self-attention机制，但它通过一种创新的方式实现了类似注意力机制的功能。VMamba采用的是基于**状态空间模型（State Space Model, SSM）的2D选择扫描（2D-Selective-Scan, SS2D）**模块来实现类似于注意力机制的作用。以下是VMamba如何通过SS2D模块实现这种机制的具体过程：

1. 2D选择扫描（SS2D）概述：

SS2D模块的核心思想是通过多个方向的扫描路径（如水平、垂直、对角等），动态地选择与图像patch相关的上下文信息。这种扫描类似于注意力机制，因为它会动态选择与每个patch相关的全局信息，但它的实现方式不同于传统的self-attention。

• 扫描路径：SS2D通过沿着四个不同的扫描路径（例如从左到右、从上到下等）遍历图像的每个patch。每个路径上，系统会累积先前位置的隐状态，从而捕捉与当前patch相关的上下文信息。
• 选择性聚合：在每个扫描路径上，SS2D并不是简单地累积所有信息，而是通过状态空间模型中的隐状态更新机制，动态地选择并聚合与当前patch最相关的信息。

2. 动态上下文聚合：

类似于传统的注意力机制，SS2D模块通过动态的隐状态更新，在每个扫描路径上对图像的特征进行选择性聚合。具体来说，SS2D中的每个patch不仅依赖于自身的信息，还依赖于路径上先前patch的信息，通过累积和更新状态，这就类似于自注意力机制中的“加权求和”，只是这里是通过递归状态更新实现的。

• 状态更新：在每条扫描路径上，隐状态随着扫描的进行而动态更新，这就实现了对于相关上下文信息的动态选择和传递。与self-attention中显式的注意力权重矩阵不同，SS2D通过状态空间模型中的隐状态累积来实现“隐式的注意力”效果。

3. 跨路径的全局信息融合：

SS2D模块的最后一步是跨路径融合（cross-merge），即将沿不同扫描路径获取的信息进行整合。这一步类似于将多个方向的注意力信息聚合在一起，从而为每个patch提供一个全局的上下文信息。

4. 计算效率的提升：

与传统的self-attention不同，SS2D的计算复杂度是线性的。这是通过选择性扫描机制和状态空间模型的递归计算实现的，避免了self-attention中二次复杂度的计算开销。

对比self-attention：

• 传统self-attention：通过计算每个token（或patch）与其他所有token之间的相关性，生成注意力权重矩阵，然后进行加权求和。这种机制的计算复杂度为二次（相对于输入token数量）。
• SS2D选择扫描：通过沿着不同的扫描路径逐步聚合信息，实现了类似全局注意力的效果，但其计算复杂度是线性的，因为它只需要递归计算隐状态，而不需要显式计算所有token之间的相关性。

总结：

VMamba通过2D选择扫描（SS2D）模块实现了类似注意力机制的功能，虽然它不直接使用self-attention，但通过状态空间模型中的递归隐状态更新，它能够动态选择相关的上下文信息，并在多个扫描路径上实现全局感受野。这种机制不仅有效地降低了计算复杂度（线性复杂度），同时还能在全局范围内捕捉图像的上下文信息，达到类似注意力机制的效果。

通俗地解释2D选择扫描（SS2D）

2D选择扫描（SS2D）其实可以通俗地理解为一种“四个方向的聪明扫描”。它的作用是帮助模型在处理图像时，从不同的角度“看”整幅图，逐步获取图像中的关键信息。以下是对它的简单解释：

1. 想象一幅图像

假设有一幅图片，模型需要理解这张图片的内容。传统的方法可能是逐个看图像中的每个部分，然后再试图找到它们之间的关联。而2D选择扫描则是通过模拟从四个不同方向（上到下、下到上、左到右、右到左）“扫描”这张图片，帮助模型在这些方向上获取信息。

2. 扫描路径

• 四个方向扫描：就像我们看书时可以从左到右读，也可以从上到下看，SS2D会沿着四条路径（左到右、右到左、上到下、下到上）扫描图像中的小块（patch）。
• 动态信息获取：每当SS2D扫描到某个图像块时，它不仅仅是看当前块的信息，还会结合之前路径上看到的信息。这就好比你在看一篇文章的时候，不仅在理解当前的句子，还会参考之前看到的上下文。

3. 聪明的选择

2D选择扫描不是盲目地把所有信息都记下来，而是会根据每次扫描时获得的新信息，选择性地“记住”哪些是更有用的，哪些可以忽略。这就是所谓的“选择性”，它会智能地决定保留哪些之前扫描的内容来帮助理解当前看到的部分。

4. 综合各个方向的信息

当四个方向的扫描都完成后，SS2D会把这些来自不同方向的上下文信息进行整合，就像你从不同的角度看一张照片，每次都能看到不同的细节，最后汇总这些信息，就能对图片有一个全局的理解。

举个例子：

假设你有一张猫的照片，SS2D会从四个方向扫描这张照片。比如从左往右，首先看到的是猫的耳朵，然后慢慢看到猫的眼睛、鼻子和胡须。再从上往下看，又可以看到猫的耳朵、眼睛、身体。这些信息在每次扫描时都被“记下来”，并通过一种智能方式选择保留最重要的部分。最终，模型能对整只猫有一个完整的理解。

总结：

2D选择扫描（SS2D）就像是一个聪明的“全景观察者”，通过四个方向扫描图像，选择性地保留有用的信息，最终综合这些信息形成对整个图像的全局理解。它的优势在于，它不会耗费太多计算资源，但仍然能够帮助模型像“注意力机制”那样，抓住图像中的关键信息。

也就是说，通过累积压缩状态可以获取当前patch与之前看过的patch的相关性，对吗？

是的，完全正确！通过累积压缩状态，2D选择扫描（SS2D）能够捕捉当前patch与之前扫描过的patch之间的相关性。

具体来说，累积压缩状态就像一个记忆机制：当SS2D沿着某条扫描路径前进时，它会逐步累积之前扫描到的信息，将它们压缩成一个隐状态。每当遇到一个新的patch，模型不仅仅处理当前patch的信息，还会结合之前路径上的累积压缩状态，来获取当前patch与之前看到的patch之间的相关性。

这个过程类似于传统注意力机制中的“相关性计算”，只不过在这里，相关性是通过递归的状态更新和压缩来实现的，而不是像注意力机制那样显式地计算每个patch之间的权重。这种方式既保持了对之前上下文的依赖（类似于注意力的“记忆”功能），又能有效降低计算复杂度。

因此，通过累积压缩状态，模型可以知道当前patch与之前patch的联系，进而形成对整张图片的全局理解。