VSSD：具有非因果状态空间对偶性的视觉Mamba模型

基本信息

代码地址: https://github.com/YuHengs ss/VSSD

摘要

视觉变换器显著推动了计算机视觉领域的发展，提供了强大的建模能力和全局感受野。然而，其高计算需求限制了其在处理长序列时的应用。为了解决这个问题，状态空间模型（SSMs）在视觉任务中崭露头角，因为它们提供了线性计算复杂度。最近，Mamba2中引入了状态空间对偶性（SSD），这是SSMs的一种改进变体，旨在提升模型性能和效率。然而，SSD/SSMs固有的因果性质限制了它们在非因果视觉任务中的应用。为了解决这一限制，我们引入了视觉状态空间对偶性（VSSD）模型，它具有SSD的非因果形式。具体来说，我们提出丢弃隐藏状态与令牌之间交互的幅度，同时保留它们的相对权重，这缓解了令牌贡献对先前令牌的依赖。结合多扫描策略的参与，我们证明了扫描结果可以被整合以实现非因果性，这不仅提高了SSD在视觉任务中的性能，还增强了其效率。我们在包括图像分类、检测和分割在内的各种基准测试上进行了广泛的实验，其中VSSD超越了现有的基于SSM的最先进模型。代码和权重可在 https://github.com/YuHengs ss/VSSD 获取。

1 引言

近年来，由Vision Transformer（ViT）开创的视觉变换器在计算机视觉领域取得了巨大成功。得益于注意力机制的全局感受野和强大的信息建模能力，基于视觉变换器的模型在分类、检测和分割等各种任务中取得了显著进展，超越了经典的基于CNN的模型。然而，注意力机制的二次计算复杂度使其在处理涉及长序列的任务时资源密集，这限制了其更广泛的应用。

最近，以Mamba为代表的状态空间模型（SSMs）引起了研究人员的极大关注。特别是S6块，提供了全局感受野，并展现出相对于序列长度的线性复杂度，提供了一种高效的替代方案。开创性的视觉Mamba模型，如Vim和VMamba，被开发出来以将SSMs应用于视觉任务。随后，提出了许多变体，它们使用不同的扫描路径将2D特征图展平为1D序列，用S6块对其进行建模，然后在多个扫描路径中整合结果。这些多扫描方法提高了

图1: (a) 将SSM/SSD应用于图像数据时的两个挑战。(b)和(c)是在ImageNet上的比较。我们的VSSD模型在准确性和效率方面均优于基于CNN的ConvNeXt、基于ViT的Swin Transformer和基于SSM的VMamba。所有模型的延迟均在A100 GPU上使用128的批大小和FP16精度测量。

SSMs在视觉任务中的性能，取得了与基于CNN和基于ViT的方法相竞争的结果。最近，Mamba2对S6块进行了进一步改进，提出了状态空间对偶性（SSD）的概念。Mamba2将状态空间转移矩阵AAA视为标量，并扩展了状态空间维度，从而提升了模型性能以及训练和推理效率。然而，将SSD/SSMs应用于视觉任务存在一个主要问题，即图像数据本质上是非因果的，而SSD/SSMs具有固有的因果特性。另一个问题是，将2D特征图展平为1D序列破坏了块之间固有的结构关系。我们在图1 (a)中提供了一个图示，以便更直观地理解这两个问题。在这个例子中，展平后的1D序列中的中心令牌仅限于访问先前的令牌，无法整合后续令牌的信息。此外，在2D空间中与中心令牌相邻的令牌1在1D序列中变得相距甚远，破坏了自然的结构关系。先前解决方案中的一个常见做法是在非因果视觉特征上增加扫描路径，这在某种程度上缓解了这两个问题。鉴于这些观察，一个重要的问题出现了：与多扫描方法相比，是否存在一种更有效、更高效的方式将SSD应用于非因果视觉数据？

为了解决这个问题，我们对SSD的分析表明，将矩阵AAA视为标量可以促进SSD直接转换为非因果且位置无关的形式，我们将其称为非因果SSD（NC-SSD）。具体来说，我们不使用AAA来确定要保留的隐藏状态的比例，而是使用它来指示当前令牌对隐藏状态的贡献程度。在这种情况下，每个令牌的贡献变得自引用。基于这个特性，我们证明了SSD中的因果掩码可以自然地移除，而无需特定的扫描路径。这一观察促使我们开发了一种非因果形式的SSD，其中可以推导出单个全局隐藏状态来替代先前的令牌级隐藏状态，这不仅提高了准确性，还提升了训练和推理速度。与先前主要缓解SSMs因果限制的多扫描方法不同，我们提出的NC-SSD还解决了将2D特征图展平为1D序列破坏相邻令牌连续性的问题。除了NC-SSD，还探索了其他技术，包括与标准自注意力的混合以及重叠下采样。基于这些技术，我们介绍了我们的视觉状态空间对偶性（VSSD）模型，并展示了其在相对于基于CNN、ViT和SSM的方法方面卓越的有效性和效率，如图1 (b)和©所示。具体而言，与最近提出的基于SSM的VMamba相比，我们的VSSD模型在ImageNet-1K数据集上的top-1准确率高出约1%，同时保持了相似的计算成本。此外，我们的模型在准确率-延迟曲线中也持续领先。除了在性能和效率之间更好的权衡之外，VSSD的另一个亮点在于训练速度。例如，与原始SSD或多扫描SSD（例如，具有双向扫描的Bi-SSD）相比，我们提出的模型分别将训练速度提高了近20%20\%20%和50%50\%50%。

总之，我们的贡献是双重的。首先，我们分析了状态空间对偶性，并证明它可以无缝转换为非因果模式。基于这一见解，我们引入了NC-SSD，它保留了原始SSD的全局感受野和线性复杂度优势，同时融入了固有的非因果特性，并实现了改进的训练和推理

效率。其次，利用NC-SSD作为基础组件，我们提出了VSSD模型，并进行了广泛的实验以验证其有效性。在参数和计算成本相似的情况下，我们的VSSD模型在分类、目标检测和分割的几个广泛认可的基准测试中优于其他最先进的基于SSM的模型。

2 相关工作

视觉变换器。视觉变换器的引入复兴了计算机视觉领域，该领域之前由卷积神经网络主导。然而，ViTs中自注意力机制的二次计算复杂度在处理高分辨率图像时带来了重大挑战，需要大量的计算资源。为了解决这个问题，提出了不同的解决方案，包括分层架构、窗口注意力和自注意力的变体。同时，线性注意力通过改变自注意力中查询、键和值的计算顺序，成功地将计算复杂度降低到线性尺度。尽管取得了这一进展，线性注意力的性能仍然不如二次自注意力及其变体。

状态空间模型。状态空间模型由于具有全局感受野和线性计算复杂度，越来越引起研究人员的注意。Mamba是SSMs的一个突出例子，它引入了S6块，在自然语言处理基准测试中取得了与变换器相当或更好的性能。随后的努力探索了将S6块应用于视觉任务，与CNN和基于ViT的模型相比，产生了有竞争力的结果。开发基于Mamba的视觉模型的一个核心挑战是使Mamba块固有的因果特性适应非因果图像数据。最直接的方法涉及使用不同的扫描路径将2D特征图展平为1D序列，然后用S6块建模并整合。受这些考虑的启发，已经采用了各种扫描路径并被证明是有效的，正如多项研究所证明的那样。最近，Mamba2强调了状态空间模型和结构化掩码注意力之间的显著重叠，将它们识别为彼此的对偶，并引入了状态空间对偶性的概念。在此基础上，我们证明了SSD可以通过简单的变换转变为非因果模式，而无需特定的扫描路径。

3 方法

3.1 预备知识

状态空间模型。经典状态空间模型用于描述连续系统的动态，将输入序列x(t) ∈ Rx(t)\,\in\,\mathbb{R}x(t)∈R转换为潜在空间表示h(t) ∈ RN\boldsymbol{h}(t)\,\in\,\mathbb{R}^{N}h(t)∈RN。然后利用该表示生成输出序列y(t) ∈ R\boldsymbol{y}(t)\,\in\,\mathbb{R}y(t)∈R。SSM的数学公式结构如下：

h′(t)=A∘h(t)+B∘x(t), y(t)=Ch(t),h^{\prime}(t)=\overset{\circ}{\mathbf{A}}h(t)+\overset{\circ}{\mathbf{B}}x(t),\;y(t)=\mathbf{C}h(t),h′(t)=A∘h(t)+B∘x(t),y(t)=Ch(t),

其中A∘∈RN×N,B∘∈RN×1 和 C∈R1×N\overset{\circ}{\mathbf{A}}\in\mathbb{R}^{N\times N},\overset{\circ}{\mathbf{B}}\in\mathbb{R}^{N\times1}\;\mathrm{和}\;\mathbf{C}\in\mathbb{R}^{1\times N}A∘∈RN×N,B∘∈RN×1和C∈R1×N是参数。为了将连续SSMs有效地集成到深度学习架构中，离散化是必不可少的。这个过程涉及引入一个时间尺度参数Δˉ∈\bar{\mathbf{\Delta}}\inΔˉ∈ : R并应用零阶保持技术进行离散化。通过这种方法，连续矩阵A∘\overset{\circ}{\mathbf{A}}A∘和B被转换为它们的离散对应物A\mathbf{A}A和B。因此，公式1以离散格式在公式2中重新定义，便于其在现代计算框架中的应用：

h(t)=Ah(t−1)+Bx(t), y(t)=Ch(t),其中 A=eΔA∘, B=(ΔA∘)−1(eΔA∘−I)ΔB∘≈ΔB∘,\begin{array}{r l}&{h(t)=\mathbf{A}h(t-1)+\mathbf{B}x(t),\;y(t)=\mathbf{C}h(t),}\\ &{\mathrm{其中}\;\mathbf{A}=e^{\mathbf{\Delta}\overset{\circ}{\mathbf{A}}},\;\mathbf{B}=(\mathbf{\Delta}\overset{\circ}{\mathbf{A}})^{-1}(e^{\mathbf{\Delta}\overset{\circ}{\mathbf{A}}}-\mathbf{I})\mathbf{\Delta}\overset{\circ}{\mathbf{B}}\approx\mathbf{\Delta}\overset{\circ}{\mathbf{B}},}\end{array}h(t)=Ah(t−1)+Bx(t),y(t)=Ch(t),其中A=eΔA∘,B=(ΔA∘)−1(eΔA∘−I)ΔB∘≈ΔB∘,

其中I表示单位矩阵。此外，公式2的过程可以以全局卷积方式实现为：

y=x⊙K, K=(CB,CAB,...,CAL−1B),\begin{array}{r}{\boldsymbol{y}=\boldsymbol{x}\odot\mathbf{K},\;\mathbf{K}=\left(\mathbf{C}\mathbf{B},\mathbf{C}\mathbf{A}\mathbf{B},\dots,\mathbf{C}\mathbf{A}^{L-1}\mathbf{B}\right),}\end{array}y=x⊙K,K=(CB,CAB,...,CAL−1B),

其中K ∈RL\in\mathbb{R}^{L}∈RL表示卷积核。最近，Mamba使参数B、C和△与输入相关。这种修改解决了先前SSM模型中线性时不变特性的局限性，从而增强了SSMs的适应性和性能。

3.2 非因果状态空间对偶性

最近，Mamba2引入了状态空间对偶性并将矩阵A简化为标量。这种选择性状态空间模型的特殊情况可以以线性和二次形式实现。不失一般性，选择性状态空间模型的矩阵变换形式表示如下：

y(t)=∑i=1tCtTAt:i+1Bix(i),其中 At:i=∏i=2tAi,y=SSM⁡(A,B,C)(x)=Fx,其中 Fji=CiTAj:iBi.\begin{array}{r l}&{y(t)=\sum_{i=1}^{t}\mathbf{C}{t}^{T}\mathbf{A}{t:i+1}\mathbf{B}{i}x(i),\mathrm{其中}\;\mathbf{A}{t:i}=\prod_{i=2}^{t}\mathbf{A}{i},}\\ &{y=\operatorname{SSM}(\mathbf{A},\mathbf{B},\mathbf{C})(x)=\mathbf{F}x,\mathrm{其中}\;\mathbf{F}{j i}=\mathbf{C}{i}^{T}\mathbf{A}{j:i}\mathbf{B}_{i}.}\end{array}y(t)=∑i=1tCtTAt:i+1Bix(i),其中At:i=∏i=2tAi,y=SSM(A,B,C)(x)=Fx,其中Fji=CiTAj:iBi.

当A2\mathbf{A}_{2}A2 : 被简化为标量时，公式4的二次形式可以重新表述为：

y=Fx=M⋅(CTB)x, 其中 Mij={Ai+1×⋯×Aji>j1i=j0i<j,y=\mathbf{F}x=\mathbf{M}\cdot(\mathbf{C}^{T}\mathbf{B})x,\;{\mathrm{其中}}\;\mathbf{M}{i j}={\left\{\begin{array}{l l}{A{i+1}\times\cdots\times A_{j}}&{i>j}\\ {1}&{i=j}\\ {0}&{i<j,}\end{array}\right.}y=Fx=M⋅(CTB)x,其中Mij=⎩ ⎨ ⎧Ai+1×⋯×Aj10i>ji=ji<j,

而其线性形式表示为：

h(t)=Ath(t−1)+Btx(t),y(t)=Cth(t).\begin{array}{r}{h(t)=A_{t}h(t-1)+\mathbf{B}{t}x(t),y(t)=\mathbf{C}{t}h(t).}\end{array}h(t)=Ath(t−1)+Btx(t),y(t)=Cth(t).

为了使SSMs适应图像数据，应首先将2D特征图展平为令牌的1D序列，然后进行顺序处理。由于SSMs的因果性质，每个令牌只能访问先前的令牌，信息传播本质上是单向的。这种因果特性在处理非因果图像数据时会导致次优性能，这一发现已得到先前工作的证实。此外，将2D特征图展平为1D序列破坏了它们固有的结构信息。例如，在2D映射中相邻的令牌可能在1D序列中最终相距甚远，导致视觉任务性能下降。由于SSD是SSMs的一种变体，将SSD用于视觉任务会带来与SSMs观察到的类似挑战：

挑战1：模型的因果性质限制了信息流，阻止了后续令牌影响先前的令牌。

· 挑战2：将2D特征图展平为1D序列破坏了处理过程中块之间固有的结构关系。

在将因果SSD应用于非因果图像数据的背景下，重新审视SSD的线性公式是有益的。在公式6中，标量AtA_{t}At调节了先前隐藏状态-1)和当前时间步信息的影响。换句话说，当前隐藏状态h(t)h(t)h(t)可以被视为先前隐藏状态和当前输入的线性组合，分别由AtA_{t}At 1和1加权。因此，如果我们丢弃这两项的幅度而只保留它们的相对权重，公式6可以重写为：

h(t)=h(t−1)+1AtBtx(t)=∑i=1t1AiBix(i).h(t)=h(t-1)+\frac{1}{A_{t}}\mathbf{B}{t}x(t)=\sum{i=1}^{t}\frac{1}{A_{i}}\mathbf{B}_{i}x(i).h(t)=h(t−1)+At1Btx(t)=i=1∑tAi1Bix(i).

在这种情况下，特定令牌对当前隐藏状态的贡献可以直接由其自身决定为⊥A\frac{\bot}{A}A⊥ :，而不是通过多个系数的累积乘法。随着每个令牌的贡献变得自引用，挑战2仅得到部分解决，因为由于挑战1中讨论的问题，当前令牌只能访问令牌的一个子集。

为了解决挑战1，先前的基于SSM的视觉模型经常采用多扫描路径。具体来说，在ViM的情况下，令牌序列经过正向和反向扫描，使每个令牌能够访问全局信息。尽管这些多扫描方法缓解了SSMs的因果特性，但它们没有解决挑战2，因为SSMs的长程衰减特性仍然局限于1D格式，并未扩展到2D。为了能够获取全局信息从而适应非因果图像数据，我们也从双向扫描策略开始。并且我们证明了公式7的正向和反向扫描结果

图2：SSD和NC-SSD的隐藏状态生成过程示意图。在隐藏状态更新过程中，NC-SSD使用标量 $A$ 来确定当前令牌的信息增量程度，而与使用 $A$ 决定要保留的隐藏状态比例的SSD形成对比。与生成令牌级隐藏状态的SSD不同，NC-SSD仅生成一个全局隐藏状态以适应非因果图像数据。

可以整合，以同时有效解决上述两个挑战。让H;表示双向扫描方法中第i个令牌的隐藏状态，我们可以轻松推导出：

Hi=∑j=1i1AjZj+∑j=−L−i1A−jZ−j=∑j=1L1AjZj+1AiZi,其中 Zj=Bjx(j).\mathbf{H}{i}=\sum{j=1}^{i}{\frac{1}{A_{j}}}\mathbf{Z}{j}+\sum{j=-L}^{-i}{\frac{1}{A_{-j}}}\mathbf{Z}{-j}=\sum{j=1}^{L}{\frac{1}{A_{j}}}\mathbf{Z}{j}+{\frac{1}{A{i}}}\mathbf{Z}{i},{\mathrm{其中}}\;\mathbf{Z}{j}=\mathbf{B}_{j}x(j).Hi=j=1∑iAj1Zj+j=−L∑−iA−j1Z−j=j=1∑LAj1Zj+Ai1Zi,其中Zj=Bjx(j).

如果我们将这个方程中的1AiZi\textstyle{\frac{1}{A_{i}}}\mathbf{Z}{i}Ai1Zi视为偏置并省略它，公式8可以进一步简化，导致所有令牌共享相同的隐藏状态H = ∑j=1L1AiZj\begin{array}{r}{{\bf H}\:=\:\sum{j=1}^{L}\frac{1}{A_{\it{i}}}{\bf Z}_{j}}\end{array}H=∑j=1LAi1Zj。在这种情况下，正向和反向扫描的结果可以无缝组合以建立全局上下文，有效地等同于移除因果掩码并过渡到非因果格式。因此，与因果特性相关的第一个挑战得到解决。尽管上述结果是从双向扫描方法推导出来的，但很明显，在这种非因果格式中，不同的扫描路径产生一致的结果。换句话说，设计特定的扫描路径来捕获全局信息变得不必要。此外，如公式8所示，不同令牌对当前隐藏状态的贡献不再与其空间距离相关。因此，将展平的2D特征图处理成1D序列不再损害原始的结构关系。因此，第二个挑战也得到解决。此外，由于整个计算过程可以并行进行，而不是依赖于先前SSMs必需的循环计算方法，训练和推理速度得到了提升。在修改了隐藏状态空间的迭代规则之后，我们按照Mamba2框架更新了相应的张量收缩算法或einsum符号在线性形式中的表示：

Z=contract(LD,LN→LND)(X,B)H=contract(LL,LDN→ND)(M,Z)Y=contract(LN,ND→LD)(C,H).\begin{array}{r}{\mathbf{Z}=\mathrm{contract}(\mathrm{LD},\mathrm{LN}\rightarrow\mathrm{LND})(\mathbf{X},\mathbf{B})\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\mathbf{H}=\mathrm{contract}(\mathrm{LL},\mathrm{LDN}\rightarrow\mathrm{ND})(\mathbf{M},\mathbf{Z})\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\qquad\mathbf{Y}=\mathrm{contract}(\mathrm{LN},\mathrm{ND}\rightarrow\mathrm{LD})(\mathbf{C},\mathbf{H}).}\end{array}Z=contract(LD,LN→LND)(X,B)H=contract(LL,LDN→ND)(M,Z)Y=contract(LN,ND→LD)(C,H).

该算法涉及三个步骤：第一步使用B扩展输入X，第二步展开标量SSM递归以创建全局隐藏状态H，最后一步使用C收缩隐藏状态H。为清晰起见，SSD和NC-SSD的前两个步骤在图2中描绘。与原始SSD相比，虽然第一步中的操作保持不变，但在非因果模式下，隐藏状态H中的序列长度维度被消除，因为所有令牌共享相同的隐藏状态。在最后一步中，输出Y通过C和H的矩阵乘法产生。给定Mi,j=1Ai\begin{array}{r}{\mathbf{M}{i,j}=\frac{1}{A{\it{i}}}}\end{array}Mi,j=Ai1，矩阵M可以通过消除其第一维度而简化为向量m ∈ ≡RL\equiv\mathbb{R}^{L}≡RL。在这种情况下，将m与X或B整合可以进一步简化公式9的变换为：

Y=C(BT(X⋅m)),\mathbf{Y}=\mathbf{C}(\mathbf{B}^{T}(\mathbf{X}\cdot\mathbf{m})),Y=C(BT(X⋅m)),

这可以被视为线性注意力的一种变体。然而，值得注意的是，正如A在Mamba2中扮演着重要角色一样，向量m也至关重要，正如我们的消融研究中所证明的那样。在实践中，我们直接使用学习到的A而不是，因为它们共享相同的值范围。为了更直观地理解公式10中m的作用，我们将不同头部的m平均值可视化，如图3所示。主要地，m聚焦于前景特征，使模型能够优先处理对当前任务至关重要的元素。
图4：提出的VSSD模型的整体架构。VSSD模型以一系列重叠卷积作为stem开始，随后是四个渐进处理阶段。前三个阶段配备了VSSD块，该块在图的底部详细说明，包含一个NC-SSD块和一个FFN。为简洁起见，本地感知单元在此可视化中省略。

图3：输入图像及其相应热图的可视化，这些热图是通过平均NC-SSD中不同头部的向量m得到的。

3.3 视觉状态空间对偶性模型

块设计。为了增强Mamba2中的SSD块以适应视觉应用，除了仅仅用NC-SSD替换SSD来开发我们的视觉状态空间对偶性块之外，还实施了若干修改。在构建NC-SSD块时，因果一维卷积被替换为核大小为三的深度卷积，这与先前的视觉Mamba工作一致。此外，在NC-SSD块之后集成了一个前馈网络，以促进跨通道的增强信息交换，并与经典视觉变换器的既定实践保持一致。此外，在NC-SSD块和FFN之前加入了一个局部感知单元，增强了模型的局部特征感知能力。还在不同块之间实现了跳跃连接。VSSD块的架构如图4底部所示。

与自注意力的混合。Mamba2证明了将SSD与标准多头自注意力整合可以产生额外的改进。类似地，我们的模型也融入了自注意力。然而，与Mamba2在整个网络中均匀插入自注意力不同，我们策略性地仅在最后阶段用自注意力模块替换NC-SSD块。这种修改利用了自注意力在处理高级特征方面的强大能力，正如先前在视觉任务中的工作所证明的那样。

重叠下采样层。由于分层视觉变换器和视觉状态空间模型主要采用非重叠卷积进行下采样，最近的研究表明，重叠下采样卷积可以引入有益的归纳偏置。因此，我们采用重叠卷积，遵循MLLA中使用的方式。为了保持参数数量和计算FLOPs具有可比性，我们相应地调整了模型的深度。

整体架构。我们根据上面讨论的方法开发了我们的VSSD模型，其架构如图4所示。遵循先前工作中已建立的视觉骨干网络的设计原则，我们的VSSD模型结构分为四个分层阶段。前三个阶段采用VSSD块，而最后阶段则采用MSA块。VSSD变体的详细架构如表1所示。
表1：VSSD变体的模型规格。

|------------|-----------------|-----------------------|------------------|-----|-------|
| 模型 | 块 | 通道数 | 头数 | #参数 | FLOPs |
| VSSD-Micro | [2,2,8, 4] | [48,96,192,384] | [2, 4, 8,16] | 14 | 2.3 |
| VSSD-Tiny | [2,4, 8, 4] | [64, 128, 256, 512] | [2, 4, 8, 16] | 24 | 4.5 |
| VSSD-Small | [3, 4, 18, 5] | [64, 128, 256, 512] | [2, 4, 8,16] | 40 | 7.4 |
| VSSD-Base | [3, 4, 18, 5] | [96,192,384, 768] | [3, 6, 12, 24] | 89 | 16.1 |

4 实验

4.1 分类

表2：ImageNet-1K上各种模型的准确率比较。†表示结果是使用MESA获得的。LAttn是线性注意力的缩写。

|-----------------------|-------|--------------|------|-------|
| Top-1 | 方法 | 类型 #参数 FLOPs | 微型模型 | |
| |||||
| RegNetY-1.6G[42] 卷积 | 11M | 1.6G | 78.0 | |
| EffNet-B3[50] | 卷积 | 12M | 1.8G | 81.6 |
| PVTv2-b1[56] | 注意力 | 13M | 2.1G | 78.7 |
| BiFormer [68] | 注意力 | 13M | 2.2G | 81.4 |
| NAT-M[22] | 注意力 | 20M | 2.7G | 81.8 |
| CMT-XS[19] | 注意力 | 15M | 1.5G | 81.8 |
| SMT-T[34] | 注意力 | 12M | 2.4G | 82.2 |
| Vim-T[69] | SSM | 7M | 1.5G | 76.1 |
| LVim-T[29] | SSM | 8M | 1.5G | 76.2 |
| VSSD-M | SSD | 14M | 2.3G | 82.5 |
| |||||
| RegNetY-4G [42] | 卷积 | 21M | 4.0G | 80.0 |
| ConvNeXt-T[38] | 卷积 | 29M | 4.5G | 82.1 |
| MambaOut-T[65] | 卷积 | 27M | 4.5G | 82.7 |
| EffNet-B4[50] | 卷积 | 19M | 4.2G | 82.9 |
| DeiT-S[51] | 注意力 | 22M | 4.6G | 79.8 |
| Swin-T[37] | 注意力 | 29M | 4.5G | 81.3 |
| PVTv2-B2[56] | 注意力 | 25M | 4.0G | 82.0 |
| Focal-T [62] | 注意力 | 29M | 4.9G | 82.2 |
| CSwin-T[8] | 注意力 | 23M | 4.3G | 82.7 |
| NAT-T[22] | 注意力 | 28M | 4.3G | 83.2 |
| VMambaV9-T[35] | SSM | 31M | 4.9G | 82.5 |
| LVMamba-T[29] | SSM | 26M | 5.7G | 82.7 |
| MSVMamba-T[47] | 1 | SSM | 33M | 4.6G |
| VSSD-T | SSD | 24M | 4.5G | 83.7 |
| MLLA-T[20] | LAttn | 25M | 4.2G | 83.5† |
| VSSD-T | SSD | 24M | 4.5G | 84.1† |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |

|-------------------|-------|--------------|-------|-------|
| Top-1 | 方法 | 类型 #参数 FLOPs | 1 | 小型模型 |
| |||||
| ConvNeXt-S [38] | 卷积 | 50M | 8.7G | 83.1 |
| EffNet-B5[50] | 卷积 | 30M | 9.9G | 83.6 |
| MambaOut-S[65] | 卷积 | 48M | 9.0G | 84.1 |
| Swin-S [37] | 注意力 | 50M | 8.7G | 83.0 |
| PVTv2-B3[56] | 注意力 | 45M | 6.9G | 83.2 |
| Focal-S [62] | 注意力 | 50M | 8.7G | 83.5 |
| CSwin-S [8] | 注意力 | 35M | 6.9G | 83.6 |
| NAT-S [22] | 注意力 | 51M | 7.8G | 83.7 |
| VMamba-S[35] | SSM | 44M | 11.2G | 83.5 |
| PMamba-L2[61] | SSM | 25M | 8.1G | 81.6 |
| VMambaV9-S[35] | SSM | 50M | 8.7G | 83.6 |
| LVMamba-S[29] | SSM | 50M | 11.4G | 83.7 |
| VSSD-S | SSD | 40M | 7.4G | 84.1 |
| MLLA-S[20] | LAttn | 43M | 7.3G | 84.4 |
| VSSD-S | SSD | 40M | 7.4G | 84.5† |
| ConvNeXt-B[38] | 卷积 | 89M | 15.4G | 83.8 |
| |||||
| MambaOut-B[65] | 卷积 | 85M | 15.8G | 84.2 |
| DeiT-B[51] | 注意力 | 86M | 17.5G | 81.8 |
| Swin-B[37] | 注意力 | 88M | 15.4G | 83.5 |
| CSwin-S[8] | 注意力 | 78M | 15.0G | 84.2 |
| NAT-B [22] | 注意力 | 90M | 13.7G | 84.3 |
| PMamba-L3[61] | SSM | 50M | 14.4G | 82.3 |
| VMambaV9-B[35] | SSM | 89M | 15.4G | 83.9 |
| VSSD-B | SSD | 89M | 16.1G | 84.7 |
| MLLA-B[20] | LAttn | 96M | 16.2G | 85.3† |
| VSSD-B | SSD | 89M | 16.1G | 85.4† |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |
| | | | | |

图5：我们的VSSD与基于CNN的模型、基于注意力的模型和基于SSM的VMamba之间的有效感受野比较。与基于SSM的VMamba相比，我们的VSSD有效消除了令牌间距对信息贡献的影响。

配置。我们的实验使用ImageNet-1K数据集进行，与先前研究的方法一致。每个模型都经过300个epoch的训练，其中包括20个epoch的热身阶段。使用AdamW进行优化，其中betas设置为(0.9, 0.999)，动量为0.9。余弦衰减调度器管理学习率，并结合0.05的权重衰减率。为了进一步优化模型准确性和泛化能力，我们结合了指数移动平均技术，并应用了系数为0.1的标签平滑。更多详细配置可在附录A中找到。在测试阶段，图像被中心裁剪为224×224的尺寸。

性能评估。表2展示了我们的VSSD模型与CNNs、ViTs和其他基于SSM的框架在ImageNet-1K数据集上的比较。VSSD-M模型配备了14M参数和2.3G FLOPs，获得了82.5%的top-1准确率，超过了同等价位的NAT-M 0.7%。在与归类为tiny和small的模型比较中，VSSD consistently 优于其对应模型。具体来说，具有24M参数和4.5G FLOPs的VSSD-T模型实现了83.7%的准确率，比VMambaV9-T高出1.2%。对于小型模型变体VSSD-S，它包含40M参数和7.4G FLOPs，实现了84.1%的准确率，超过了LocalVMamba-S 0.4%。在基础尺寸变体中，我们的VSSD-B具有89M参数和16.1G FLOPs，记录了84.7%的准确率，比VMambaV9-B高出0.8%。当引入MESA进行进一步优化时，我们的tiny、small和base尺寸模型的结果分别提高到84.1%、84.5%和85.4%。

除了定量比较，我们还对各种模型训练前后的有效感受野进行了比较分析，包括基于CNN的ResNet50和ConvNeXt-Tiny，基于注意力的Swin-Tiny和DeiT-Small，以及基于SSM的VMamba-Tiny，还有我们的VSSD-Tiny。中心像素的ERF使用[39]中提出的方法绘制，使用了从ImageNet-1K验证集中随机选择的50张分辨率为1024x1024的图像。为了证明所提出的NC-SSD的有效性，在分析中我们的VSSD模型没有采用第3.3节中讨论的混合自注意力和重叠下采样层等技术。值得注意的是，只有我们的VSSD和DeiT在训练前后都 consistently 表现出全局感受野。在训练后，VMamba中观察到明显的十字形衰减，而我们的方法有效消除了令牌间距对信息贡献的影响。

4.2 目标检测和实例分割

配置。我们对VSSD模型的评估在MS COCO数据集上使用Mask R-CNN框架进行与目标检测和实例分割相关的任务。所有实验均使用MMDetection库进行。与先前的研究一致，在训练阶段，图像被调整使得较短边为800像素，而较长边不超过1333像素。使用AdamW优化器进行优化，设置学习率为0.0001，批大小为16。当采用标准的"1x"训练
表3：在MS COCO数据集上使用Mask R-CNN框架的目标检测和实例分割结果。FLOPs使用1280 × 800的输入大小进行测试。

|----------------------------------------------------------|------------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------|-----------------------------------------------------------|--------------------------------|-----------------------------------------|-------------------------------------------------------------------------|
| Mask R-CNN 1x .5 39.3 35.1 42.8 36.7 $\\mathrm{AP\^{m}}$ $\\mathrm{AP_{7s}\^{b}}$ ||||||||
| 方法 PVT-T[55] EffVMamba-S[40] MSVMamba-M[47] VSSD-M | 36.7 $\\mathrm{AP\^{b}}$ | 59.2 61.8 65.8 67.5 $\\mathrm{AP_{50}\^{b}}$ | v 56.7 37.3 58.9 39.2 $\\mathrm{AP_{50}\^{m}}$ $\\mathrm{AP_{75}\^{m}}$ | #参数 33M 31M 32M 33M 48M | FLOPs 208G 197G 201G 220G | 39.3 35.1 42.8 36.7 47.7 39.9 49.8 41.3 | 37.3 39.2 42.9 44.6 |
| 39.3 43.8 45.4 | 62.9 64.5 62.2 63.3 65.5 | 49.8 46.8 48.3 50.7 50.8 | Swin-T [37] ConvNeXt-T[38] VMamba-T[35] LocalVMamba-T[29] | 42.7 44.2 46.5 46.7 46.9 | 65.2 66.6 68.5 | 39.3 40.1 42.1 42.2 | 42.2 42.8 45.3 45.5 45.4 |
| 267G 262G 286G 291G 252G | 48M 42M 45M 53M 44M 69M | 46.7 68.7 50.8 46.9 68.8 51.4 46.9 69.4 51.4 | 65.7 45.5 65.6 45.4 66.4 45.9 63.2 44.2 65.2 45.1 | MSVMamba-T[47] VSSD-T Swin-S [37] | 42.2 42.6 40.9 41.8 | 51.4 48.9 50.0 52.5 52.7 | 265G 354G 348G 400G 414G 325G |
| 44.8 66.6 45.4 67.9 48.2 69.7 48.4 69.9 | 44.2 45.1 46.4 46.5 47.1 | ConvNeXt-S [38] VMamba-S[35] LocalVMamba-S[29] | 45.4 48.2 48.4 48.4 | 70M 64M 69M 59M #参数 | 43.0 66.6 43.2 66.7 43.5 67.2 | | |
| VSSD-S 方法 PVT-T[55] | 48.4 70.1 39.8 62.2 $\\mathrm{AP_{50}\^{b}}$ $\\mathrm{AP\^{b}}$ | 70.1 53.1 43.5 Mask R-CNN3x + MS .5 62.2 43.0 37.4 64.4 45.1 38.4 $\\mathrm{AP_{50}\^{b}}$ $\\mathrm{AP_{7.}\^{b}}$ $\\mathrm{AP\^{m}}$ | 37.4 38.4 38.2 41.8 $\\mathrm{AP\^{m}}$ | AP 59.3 39.9 61.2 40.8 60.8 40.7 $\\mathrm{AP_{75}\^{m}}$ | FLOPs 208G 187G 197G 201G 220G | 33M 28M 31M 32M 33M | LightViT-T[28] EffVMamba-S[40] MSVMamba-M[47] VSSD-M Swin-T[37] |
| 45.6 38.2 50.8 41.8 52.1 42.8 50.3 41.6 | 60.8 40.7 65.1 44.9 66.5 46.0 65.1 44.9 | 46.0 68.1 46.2 67.9 48.5 69.9 48.7 70.1 48.8 70.4 | 50.3 41.6 50.8 41.7 52.9 43.2 53.0 43.4 53.4 43.6 | 65.1 44.9 65.0 44.9 66.8 46.3 67.0 46.4 67.6 46.9 | 48M 48M 42M 45M 44M | 267G 262G 286G 291G 265G | ConvNeXt-T[38] VMamba-T[35] LocalVMamba-T[29] VSSD-T |

调度时，学习率在第8和第11个epoch降低0.1倍，而扩展的"3× + MS"调度在第27和第33个epoch看到学习率降低相同的倍数。

性能评估。表3详细描述了我们的模型与成熟CNNs、ViTs和其他基于SSM的模型的比较性能。我们的VSSD模型在各种配置中表现出卓越的性能。值得注意的是，我们的VSSD-T模型展示了显著优势，在框AP和掩码AP上分别超过Swin-T +4.2和+3.3。在扩展的"3x"训练计划下，VSSD-T仍然 consistently 优于各种竞争对手。

4.3 语义分割

配置。与Swin和VMamba中描述的方法一致，我们的实验利用UperHead框架，使用ImageNet预训练骨干进行初始化。训练计划跨越160K次迭代，批大小为16，使用MMSegmentation库执行。主要实验使用512 × 512的标准输入分辨率进行。为了进一步评估我们模型的鲁棒性，实施了多尺度测试。使用AdamW优化器进行优化，学习率设定为6×10−56\times10^{-5}6×10−5。

性能评估。我们模型及其竞争对手的详细性能指标显示在表4中，包括单尺度和多尺度测试场景。具体来说，在Tiny模型类别和单尺度测试的背景下，我们的VSSD模型展示了卓越的性能，在tiny变体中分别超过Swin、ConNeXt和VMamba模型+3.5、+1.9和+0.6 mIoU。
表4：在ADE2oK数据集上使用UperNet框架的语义分割结果。所有模型的FLOPs均使用512 × 2O48的输入尺寸计算。表中"SS"代表单尺度测试，而"MS"代表多尺度测试。

4.4 消融研究

|-----------------------------------------------------|--------------------------------------|-------------------------------|---------------------|-------------------------------------------|
| 方法 EffVMamba-S[40] | mIoU SS | mIoU MS 42.1 45.4 46.0 | 29M 42M 42M 60M | 505G 875G 893G 945G |
| 41.5 45.1 45.6 | MSVMamba-M[47] VSSD-M Swin-T[37] | 45.8 46.7 48.3 | 44.4 46.0 47.3 47.9 | |
| ConvNeXt-T[38] VMamba-T[35] LocalVMamba-T[29] | 60M 55M 57M 65M 65M 53M | 939G 964G 970G 930G 942G 941G | 49.1 47.3 48.5 48.7 | EffVMamba-B[40] MSVMamba-T[47] VSSD-T |

为了验证所提出模块的有效性，我们在VSSD-Micro模型上进行了详细的消融实验。使用SSD块作为令牌混合器和遵循Swin和原始VMamaba的分块下采样器，我们建立了基线配置，详见表5的第一行。对于吞吐量测试，我们使用了A100-PCIE-40G GPU，批大小为128，FP16精度。
表5：VSSD-Micro在ImageNet-1K上的消融研究。我们的NC-SSD在准确性和效率方面 consistently 优于原始SSD和Bi-SSD。其他技术进一步提升了性能。

|-------------------------------|----|------------|--------|--------|-----|------------------|-------------|
| 操作类型下采样器层数 Top-1 #参数|FLOPs | | | 准确率(%) | | | (G)(图像/秒) (图像/秒) | 训练吞吐量推理吞吐量 |
| SSD | 分块 | 2,4, 8, 4 | 81.0 | 14.8 M | 2.1 | 1818 | 523 |
| Bi-SSD | 分块 | 2, 4,8, 4 | 81.4 | 15.2 M | 2.2 | 1741 | 399 |
| NC-SSD | 分块 | 2, 4,8, 4 | 81.6 | 14.8 M | 2.1 | 1843 | 909 |
| 混合 | 分块 | 2, 4, 8, 4 | 81.8 | 13.4M | 2.1 | 1890 | 622 |
| 混合 | 卷积 | 2,2,8,4 | 82.5 | 13.5 M | 2.3 | 1918 | 597 |

不同的SSD机制。在我们的令牌混合器消融研究中，我们探索了SSD的不同扫描路径。具体来说，引入了Bi-SSD，我们将通道分成两部分，并反转其中一部分以创建反向扫描序列。这些具有相反扫描路径的序列随后在SSD块之后连接起来。如表5所示，我们的NC-SSD模型在top-1准确率上分别比原始SSD和Bi-SSD高出0.6%和0.2%。此外，训练和推理吞吐量都得到了提升，与Bi-SSD方法相比，NC-SSD将训练吞吐量提高了近50%。

混合架构和重叠下采样器。在最后阶段加入标准注意力以及使用重叠下采样器的有效性在表5的最后两行中得到证明。具体来说，在最后阶段用标准注意力替换NC-SSD导致准确率提高0.2%，同时略微减少参数。用重叠卷积方式替换分块下采样器将准确率提高0.7%，同时将FLOPs增加0.2G。为了保持近似参数，我们将层配置从[2,4,8,4]调整为[2,2,8,4]。

m的效果。公式10将NC-SSD概念化为线性注意力的一种变体，它包含一个额外的权重向量m。图3直观地展示了m如何选择性地强调前景特征。为了定量评估m的影响，我们在NC-SSD块中进行了有和没有该组件的实验，训练100个epoch，其中5个epoch用于热身。结果如表6所示，揭示了m对模型性能的显著影响。
表6：关于NC-SSD中m效果的消融研究。符号†表示在遇到N.A之前达到的最佳准确率。

|--------|---------|---|---------------------------------|---------|
| | | | 操作尺寸 m|Top-1 #参数 FLOPs 准确率(%) | (G) |
| NC-SSD | Tiny | | 32.6†24.3M 81.8 824.3M | 4.5 4.5 |
| NC-SSD | Small x | | N.A 40.0M | 7.4 |

没有m，我们的实验表明模型训练不稳定，导致崩溃。这种不稳定性在较大的模型中尤其明显。我们报告了训练崩溃前达到的最高准确率，用†标记。对于tiny尺寸的模型，最佳准确率仅为32.6%。对于small尺寸的模型，训练在第一个epoch就崩溃了。我们假设这种不稳定性是因为，在没有线性注意力方法中通常使用的归一化技术的情况下，特征的幅度急剧上升，导致崩溃。相比之下，使用m，模型实现了81.8%的强大top-1准确率，同时保持了相同的参数数量和计算复杂度。

5 局限性

尽管提出的VSSD模型在ImageNet-1K上优于其他基于SSM的模型，但VSSD在下游任务上与基于SSM的其他模型相比，性能提升是微小的。当与最先进的视觉变换器变体进行评估时，在下游任务上的性能仍然存在显著差距。此外，本文缺乏涉及更大模型和更广泛数据集的实验，例如使用ImageNet-22K基准测试。因此，提出的VSSD模型的可扩展性仍然是一个有待进一步探索的领域。

6 结论

总之，我们的研究引入了NC-SSD，它通过修改矩阵A的作用并消除因果掩码来重新定义SSD。这些调整促进了向非因果模式的转变，显著提高了准确性和效率。广泛的实验证明了其相对于原始SSD及其基于多扫描的变体的优越性。此外，通过整合诸如混合标准注意力和重叠下采样等技术，我们的VSSD模型在几个广泛使用的基准测试中实现了与成熟CNNs、ViTs和视觉SSMs相当或更优的性能。