北大 & 中科大提出 PlainMamba | 即插即用，保持高性能的同时大大减少计算量 !

AIGC 先锋科技

发布于 2024-07-08 12:57:40

1820

发布于 2024-07-08 12:57:40

文章被收录于专栏：AIGC 先锋科技

作者介绍了PlainMamba：一种简单的非层次化状态空间模型（SSM），专为通用视觉识别而设计。最近的Mamba模型已经展示了SSM如何在序列数据上与其他架构高度竞争，并已初步尝试将其应用于图像。在本文中，作者进一步将 Mamba 的选择性扫描过程适应到视觉领域，通过（i）一种连续的2D扫描_过程，通过确保扫描序列中标记的相邻性来提高空间连续性，以及（ii）方向感知更新，使模型能够通过编码方向信息来辨别标记的空间关系。作者的架构设计简单易用且易于扩展，由堆叠相同的 PlainMamba块组成，使得模型在所有层次中宽度保持恒定。通过移除特殊标记的需求，进一步简化了架构。作者在包括图像分类、语义分割、目标检测和实例分割在内的各种视觉识别任务上评估了PlainMamba。作者的方法在性能上超过了之前的非层次化模型，并且与层次化替代方案具有竞争力。特别是在需要高分辨率输入的任务中，PlainMamba在保持高性能的同时，计算量大大减少。代码和模型可在 https://github.com/ChenhongyiYang/PlainMamba 获取。

1 Introduction

开发高性能的视觉编码器一直是计算机视觉领域最重要的目标之一。借助高质量的视觉特征，一系列下游任务，如语义分割、物体识别和检测，都可以相对容易地解决。

早期提取视觉表示的方法依赖于手工设计的特征，如SIFT [61]和SURF [4]。随后，一个重大突破是采用了卷积神经网络（CNNs），它们处理具有局部上下文的图像并强制空间等方差。最近，视觉 Transformer （ViTs）[24]摒弃了这种强制的归纳偏置，转而采用对图像块进行操作的可学习上下文。然而，尽管 Transformer 及其自注意力机制取得了巨大成功，但注意力的二次成本被认为是进一步扩展这些模型的障碍。

这激发了人们对状态空间模型（SSMs）的兴趣。由于与线性循环网络紧密相关，SSMs在潜在无限上下文长度的同时，保持了输入序列长度的线性复杂度[31]，与注意力机制相比，这提供了显著的加速。然而，要让SSMs在学习竞争力表示方面取得有效进展，需要几项重要的改进，包括强制状态空间变量成为正交基投影[34]。最近的Mamba[31]架构进一步将基于SSM的模型与现代化 Transformer 对齐，例如使状态空间变量依赖于输入 -- 类似于自注意力中的 Query 、键和值。当用于NLP时，这些设计导致了可以扩展到现代化基于 Transformer 的LLM[8, 81]的大小和性能的状态空间模型，同时提高了长距离上下文和推理速度。

目前，人们对于将Mamba架构适配到视觉领域产生了浓厚的兴趣。然而，在作者开始这样做之前，作者需要思考按照什么指导原则来设计作者新的模型。正如作者在图1中展示的，通过审视最近提出的视觉编码器的发展，作者发现向诸如DeiT这样的普通模型中添加更多的归纳偏置，例如层次结构，确实可以提高模型在ImageNet这类开源基准测试上的性能。

然而，作者不应忽视这样一个事实：普通的ViT 被多个流行的视觉基础模型所采用，这表明架构设计的简洁性对于多种原因来说都是关键。首先，保持恒定的模型宽度（即非层次化）使得集成来自多个 Level 的特征变得更加容易，这在诸如语义分割 [46] 这类密集预测任务中很常见。它还使得在不同模态之间结合特征变得更加容易，例如在CLIP 或LLaVa 中，或者作为越来越复杂的AI驱动系统的组成部分。此外，更简单的组件可以更容易地为硬件加速进行优化。另外，也观察到过度设计的模型可能导致在常用基准测试和下游任务之间的性能存在显著差距。这意味着基准性能可能不再反映现实世界的实用性，因为过度工程化往往会增加模型的复杂性，从而使得其他人更难以复用。

受上述发现的启发，作者提出了PlainMamba：一种用于视觉识别的简单Mamba架构。这个模型融合了来自CNNs、Transformers和基于新型SSM模型的思想，旨在为视觉模态提供易于使用的模型。与先前的视觉状态空间模型[108, 58]相比，作者通过在网络的所有层中堆叠相同的块以保持恒定的模型宽度，并移除了对CLS Token 的需求，简化了架构。这使得模型易于扩展和重用，同时实现了具有竞争力的性能。

作者的贡献如下：

作者提出了一种新的视觉状态空间模型，作者称之为PlainMamba。这种架构改进并简化了之前将Mamba架构扩展到视觉模态的尝试。

作者通过两种方式改进了SSM块，即适应选择性扫描以更好地处理二维空间输入：(i) 作者的连续二维扫描方法确保扫描序列在空间上是连续的，以增强语义的连续性。(ii) 受位置编码启发的方向感知更新，允许模型编码每个扫描顺序的方向性，从而进一步改善空间上下文。

作者用三种不同规模（7M、26M和50M）对PlainMamba架构进行了测试，并展示了它在从ImageNet1K分类到语义分割和目标检测等一系列任务上的竞争力表现。具体来说，作者展示了PlainMamba在性能上超过了其非分层对应架构，包括SSMs和Transformers，并且与分层竞争对手持平。

2 Related Work

2.0.1 Visual Feature Extractors

如何有效地从图像中提取视觉特征一直是计算机视觉领域的一个长期挑战。在深度学习的早期，CNNs[39, 47, 73, 95]在模型架构领域占据主导地位。它们通过使用卷积滤波器引起的空间先验，利用了视觉特征的局部性。此外，堆叠多层增加了它们的感受野。多年来，已经提出了许多不同的CNN Backbone 架构[47, 44, 13]，引入了利用空间信息[73, 95]的新方法，构建更深层模型[39, 77]，提高效率[78, 66, 71]，添加多尺度连接[69]，扩展架构[91]，并引入注意力机制[5, 10, 42, 75, 83, 85]。近年来，ViTs已成为图像建模的强大工具[24]。与CNNs相比，它们对数据的假设更少（特征局部性[98]，平移和尺度不变性）。

通过用自注意力模块替换卷积层， Transformer 可以捕获全局关系，并在许多常见的图像基准测试中取得了最先进的结果[20, 54, 107]。为了将原始的 Transformer 架构[82]适应于视觉任务，图像被分割成块并转换为标记，然后送入 Transformer 编码器。在这个框架内，许多研究工作都集中在提高性能上（例如，LeViT，[30]结合 Transformer 编码层和卷积），或降低自注意力的昂贵二次复杂度[16, 18]。

对ViT架构的另一种流行扩展是添加分层结构[27, 59, 86, 92, 97]，类似于CNNs中使用的多尺度特征金字塔。例如，Swin Transformer[59]使用移位的窗口来跨尺度共享特征信息。这些多尺度特征随后用于广泛的下游任务。最近的研究探讨了如何在ViTs本身内使用这些分层特征[9, 23, 25, 36, 37, 48, 51, 68, 72]。一些研究[59]探讨了将多分辨率特征作为注意力键和值来学习多尺度信息的用法。然而，这些扩展增加了模型的复杂性，使得在后期阶段难以有效地使用其特征，从而阻碍了其广泛应用。实际上，最近的研究[50, 100]回到了原始的ViT架构，因为其非分层性质极大地简化了其特征的使用。特别是，普通的ViT为在不同任务上进行预训练和微调提供了更大的灵活性。

2.2.2 State Space Models

状态空间模型（SSMs）因其能随序列长度线性扩展而成为 Transformer （transformers）和卷积神经网络（CNNs）的有效替代品[33, 76]。SSMs通过转换状态空间，有效捕捉扩展序列中的依赖关系。为了减轻这类模型的初始计算成本，S4 [32] 对状态矩阵实施了低秩约束，而S5 [74] 引入了并行扫描以进一步提高效率。此外，H3 [29] 通过改进硬件利用率，在常见基准测试中取得了具有竞争力的结果。最后，Mamba [31] 将SSM矩阵参数化为输入的函数，从而使其能够作为一种可学习的选择机制，并提供更大的灵活性。后续工作通过使用分层结构[58]和双向块[108]，将选择性的SSMs扩展到了图像[3, 57, 70, 88, 89]和视频[62]上，而Mamba-ND [49] 引入了针对多维数据的架构。MambaIR [35] 解决图像恢复问题，而Pam-Mamba [40] 专注于全色锐化。DiS [28] 通过用SSM主干网替换U-Net，将SSMs引入到扩散模型中。虽然从上述工作中汲取灵感，但PlainMamba通过添加更宽的深度卷积来改进Mamba [31] 的选择性SSM块。与跨扫描模块（CSM）[58]和Mamba-ND [49]相比，PlainMamba尊重图像块的空间序列特性（见图2）。与[108]不同，作者不使用CLS标记。

2.2.2 Simplifying Visual Feature Extractors

简化统一现有方法与提高性能同样重要。简单的非层次化架构具有鲁棒性，概念简单，并且易于扩展。此外，它们也为未来的研究奠定了基础。ViTs [24] 通过将图像转换为块状 Token 来移除CNN的金字塔结构。这种方式，它们轻松地适应了 Transformer 架构以处理视觉任务。源自序列建模的另一个技巧是使用CLS Token 进行预测，这已被证明对于视觉任务是不必要的[103]。FlexiVit [7] 统一了单一架构中具有不同输入分辨率的图像，而GPViT [100] 通过非层次化的 Transformer 改善了特征分辨率。

同样，ConvNext [60] 引入了一个简单的CNN模型，与最先进的 Transformer 方法相竞争。其他研究，如MLP-Mixer [79] 和后续工作 [41]，仅使用了多层感知器来引入简单的架构。简单的非层次化ViT [24] 已成为许多不同任务的简单构建块。SAM [46] 使用预训练的ViT作为图像编码器，仅做最小改动，以实现大规模的图像分割。DinoV2 [63, 19] 使用ViT通过自我监督在精选数据集上预训练模型，以学习通用视觉特征。同样，CLIP [65] 模型的图像编码器由一个带有少量修改的基本ViT组成，允许通过对比目标学习图像-文本表示。DALLE-2 [1] 采用ViT图像编码器来提取用于条件文本图像生成的视觉特征。LlaVA [55, 56] 将视觉编码器（来自CLIP的预训练ViT）和LLM结合用于视觉-语言任务。受到ViT的通用性和灵活性的启发，这项工作旨在通过引入PlainMamba：一个简单的Mamba架构，继续简化特征提取器的同时保持强大的性能。

3 Method

在本节中，作者将详细介绍PlainMamba。所提出模型的概览在图2（a）中展示。在3.1节中，作者回顾了状态空间建模和Mamba的背景。接着，在3.2节中作者介绍了作者的模型架构，而在3.3节中则介绍了作者的块架构。最后，在3.4节中，作者展示了PlainMamba的不同变体。

Preliminaries

3.1.1 State Space Models.

状态空间模型（SSMs）通常用于模拟连续线性时不变（LTI）系统[90]，其中输入信号

x(t)\in\mathbb{R}

通过状态变量

h(t)\in\mathbb{R}^{m}

映射到输出信号

y(t)\in\mathbb{R}

，遵循以下规则：

h^{\prime}(t)=\mathbf{A}h(t)+\mathbf{B}x(t),\;\;\;y(t)=\mathbf{C}h^{\prime}(t) +\mathbf{D}x(t) \tag{1}

其中

\mathbf{A}\in\mathbb{R}^{m\times m}

\mathbf{B}\in\mathbb{R}^{m\times 1}

\mathbf{C}\in\mathbb{R}^{1\times m}

和

\mathbf{D}\in\mathbb{R}^{1\times 1}

是参数。

为了使上述系统适用于一个离散系统，例如，一个序列到序列的任务，使用时间尺度参数

\mathbf{\Delta}

来将参数

\mathbf{A}

和

\mathbf{B}

转换为它们的离散化对应项

\bar{\mathbf{A}}

和

\bar{\mathbf{B}}

。在Mamba [31] 及其后续工作 [58, 108] 中，这是通过以下零阶保持（ZOH）规则实现的：

\bar{\mathbf{A}}=\exp{(\mathbf{\Delta}\mathbf{A})},\;\;\;\bar{\mathbf{ B}}=(\mathbf{\Delta}\mathbf{A})^{-1}(\exp{(\mathbf{\Delta}\mathbf{A})}-\mathbf{I}) \cdot\mathbf{\Delta}\mathbf{B} \tag{2}

随后，输入序列

\{x_{i}\}

（对于

i=1,2,...

）可以以类似的方式映射到其输出序列

\{y_{i}\}

。

h^{\prime}_{i}=\bar{\mathbf{A}}h_{i-1}+\bar{\mathbf{B}}x_{i},\; \;\;y_{i}=\mathbf{C}h^{\prime}_{i}+\mathbf{D}x_{i} \tag{3}

Mamba. 由于SSMs通常用于模拟LTI系统，它们的模型参数被所有时间步

共享。然而，如Mamba [31] 中发现的，这种时间不变特性严重限制了模型的代表性。为了缓解这个问题，Mamba 提高了时间不变约束，并使参数

\mathbf{B}

、

\mathbf{C}

和

\mathbf{\Delta}

依赖于输入序列

\{x_{i}\}

，他们称这个过程为 _选择性扫描_，从而得到依赖于标记的

\{\mathbf{B}_{i}\}

、

\{\mathbf{C}_{i}\}

和

\{\mathbf{\Delta}_{i}\}

。此外，将SSM与门控MLP [43] 结合起来以获得更好的表示能力。具体来说，输出序列

\{y_{i}\}

是从

\{x_{i}\}

中计算出来的，如下所示：

x^{\prime}_{i}=\sigma\big{(}\text{DWConv}\big{(}\text{Linear}(x_ {i})\big{)}\big{)},\;\;\;z_{i}=\sigma\big{(}\text{Linear}(x_{i})\big{)} \tag{4}

\mathbf{B}_{i},\mathbf{C}_{i},\mathbf{\Delta}_{i}=\text{Linear}(x^{ \prime}_{i}),\;\;\;\bar{\mathbf{A}}_{i},\bar{\mathbf{B}}_{i}=\text{ZOH}(\mathbf{ \Delta},\mathbf{B}_{i},\mathbf{\Delta}_{i})

(5)

h^{\prime}_{i}=\bar{\mathbf{A}}_{i}h_{i-1}+\bar{\mathbf{B}}_{i} x^{\prime}_{i},\;\;\;y^{\prime}_{i}=\mathbf{C}_{i}h^{\prime}_{i}+\mathbf{D}x^{ \prime}_{i},\;\;\;y_{i}=y^{\prime}_{i}\odot z_{i} \tag{6}

其中

\sigma

表示SiLU激活函数，而

\odot

表示逐元素相乘。

Overall architecture of PlainMamba

在图2中，作者展示了PlainMamba的模型架构。作者的模型分为三个主要部分：(1)一个卷积标记器，将输入的2D图像转换成一系列视觉标记；(2)主网络，包含一系列

个相同的PlainMamba块，用于学习高质量的视觉表示；以及(3)针对不同类型下游任务的特定任务 Head 。

更详细地说，分词器将输入图像

I\in\mathbb{R}^{H_{I}\times W_{I}\times 3}

下采样成一个视觉标记列表

x\in\mathbb{R}^{H\times W\times C}

，其中

是通道数。作者遵循ViT [24]的做法，将默认的下采样因子设置为16。在将初始视觉标记与位置嵌入[82]结合以保留空间信息之后，标记通过一系列

PlainMamba 块进行转换，这些块的设计旨在通过保持输入-输出形状的一致性来简化使用。架构的最后阶段包含一个特定于任务的 Head ，这取决于特定的下游应用。例如，在图像分类任务中，图像标记被全局池化成一个向量，然后该向量被送入一个线性分类 Head 以产生最终输出。

PlainMamba在几个关键方面与现有的视觉 Transformer [24, 57]和同时期的视觉Mamba[58, 108]架构区分开来。首先，它不使用任何特殊标记，例如常用的CLS标记。其次，与采用分层结构来管理特征分辨率的方法[59, 87, 51]相比，PlainMamba相反，在所有区块中保持恒定的特征分辨率。这种设计选择考虑了在各种各样的视觉基础模型[63, 65, 46]中取得的近期进展，在这些模型中使用了普通的非分层ViT，而不是其分层对应物。

PlainMamba Block

PlainMamba的整体架构由若干个相同的PlainMamba块组成，这些块构成了学习高质量视觉特征的基础。作者在图2中展示了PlainMamba块的结构。作者从原始的Mamba块中获得灵感，其中SSM与门控MLP在一个统一模块中结合。然而，作者对选择性扫描过程进行了几处关键调整，以充分利用图像输入的二维特性。这种适应对于有效地从语言模型的固有一维处理范式过渡到图像的二维领域至关重要。为此，作者引入了两种新技术：(1) 连续二维扫描 和 (2) _方向感知更新_。第一种技术确保每个视觉标记总是与先前扫描的标记相邻。这样做可以减轻位置偏差，并鼓励对图像空间进行更均匀的理解，从而提高模型从视觉输入中学习的能力。第二种技术将2D相对位置信息显式嵌入到选择性扫描过程中，这使得模型能够更好地解释扁平化视觉标记的位置上下文。

4.1.3 Continuous 2D Scanning

选择扫描机制本质上是为顺序数据（如文本）设计的。将该机制适应于2D图像数据需要将2D视觉标记展平成1D序列，以应用状态空间模型（SSM）的更新规则。先前的研究，例如VisionMamba [108]和VMamba [58]，已经证明了使用多种扫描顺序来提升模型性能的有效性——如在多个方向上进行行扫描和列扫描。然而，如图3（a）和（b）所示，在这些方法中，每种扫描顺序只能覆盖一种类型的2D方向，例如从左到右，当移动到新的一行（或列）时会导致空间上的不连续。此外，由于方程3中的参数

\mathbf{A}

起到衰减作用，这样的空间不连续还可能导致相邻标记衰减到不同程度，加剧了语义上的不连续，并可能导致性能下降。

作者的“连续二维扫描”通过确保每次扫描的视觉标记始终与上一次扫描的标记在二维空间上相邻，来解决这一挑战。如图3（c）所示，在作者的方法中，视觉标记也是按照四种不同的顺序进行扫描的。当扫描到一行（或一列）的末尾时，下一个要扫描的标记将是下一列（或行）中与之相邻的标记，而_不是_对面的标记。然后，扫描以相反的方向继续，直到到达图像的最后一个视觉标记。因此，作者的方法保持了空间和语义的连续性，并且在扫描非相邻标记时避免了潜在的信息丢失。此外，在实际应用中，模型通常接收相同大小的输入图像，这意味着作者的方法可以通过预计算排列索引来轻松实现并高效运行。

3.3.3 Direction-Aware Updating

如公式3所示，在选择性扫描中， Token

x_{i}

对隐藏状态

h_{i}

的贡献由参数

\bar{\mathbf{B}}_{i}

决定，这个参数是从

x_{i}

本身导出的。在语言模型中，序列顺序自然指示了 Token 之间的位置关系，使模型能够“记住”它们的相对位置。然而，在作者的连续二维扫描中，当前 Token 相对于前一个 Token 可以是四种可能方向中的一种。这对模型基于仅

\mathbf{B}_{i}

来辨别连续 Token 之间精确空间关系的能力提出了挑战。

作者的 方向感知更新 因此被提出以应对这一挑战。从视觉 Transformer 中的相对位置编码机制[24]获得灵感，作者采用一组可学习参数

\{\mathbf{\Theta}_{k}\in\mathbb{R}^{m\times 1}\}

（对于

k=1,2,...,5

），代表四个基本方向加上一个特殊的BEGIN方向用于初始标记。这些参数与数据相关的

\mathbf{B}_{i}

相加，用方向信息丰富选择扫描过程。具体来说，根据方程3中的

x_{i}

和

z_{i}

，作者的 方向感知更新 表述如下：

h^{\prime}_{k,i}=\bar{\mathbf{A}}_{i}h_{k,i-1}+(\bar{\mathbf{B}}_{i}+\bar{ \mathbf{\Theta}}_{k,i})x_{i} \tag{7}

y^{\prime}_{i}=\sum_{k=1}^{4}\bigl{(}\mathbf{C}_{i}h^{\prime}_{k,i}+\mathbf{D} x_{i}\bigr{)},\ \ y_{i}=y^{\prime}_{i}\odot z_{i} \tag{8}

在文中，

跨越了作者 Continuous 2D Scanning 引入的四种不同的扫描方向。另外，对于每次扫描的初始标记，作者改为添加最终的

\bar{\mathbf{\Theta}}_{k=5}

向量。术语

\bar{\mathbf{\Theta}}_{k,i}

表示使用

\mathbf{\Delta}_{i}

对

\mathbf{\Theta}_{k,i}

进行离散化。

Model Variants of PlainMamba

作者提出了三种不同变体的PlainMamba模型。作者在表1中报告了它们的架构细节。具体来说，从PlainMamba-L1到PlainMamba-L2，作者扩大了模型的宽度，即特征通道数，并将模型的深度保持在24层。从PlainMamba-L2到PlainMamba-L3，作者同时扩大了模型的宽度和深度。FLOPs是使用224

\times

224的输入进行测量的，作者遵循官方Mamba代码库来计算选择性扫描过程的FLOPs。

4 Experiments

ImageNet-1K Classification

设置： 作者按照[100]构建了作者的代码库。具体来说，对于ImageNet-1k的实验，作者使用AdamW优化器训练所有PlainMamba模型300个周期。遵循[59]，作者将批量大小设置为2048，权重衰减设置为0.05，峰值学习率设置为0.002。使用了余弦学习率调度。对于数据增强，作者使用了常见的配方[22, 23, 59, 80]，其中包括Mixup [104]，Cutmix [102]，随机擦除[106]和Rand增强[15]。

结果：在表2中，作者报告了ImageNet-1K实验的结果。作者将PlainMamba与三种不同的视觉特征提取器进行了比较：CNNs、视觉 Transformer 以及SSMs。此外，比较还包括了分层和非分层模型。首先，与SSMs相比较，作者的模型表现优于最近提出的Vision Mamba [108] 和 Mamba-ND [49]。例如，尽管PlainMamba-L2与Mamba-ND-T具有相似的模型大小，但前者的准确率却高出2.4%。这些结果验证了PlainMamba作为一个非分层SSM的有效性。其次，与CNNs和 Transformer 相比较，作者的模型在非分层对应模型中取得了更好的性能。例如，PlainMamba-L2比DeiT-Small的准确率高出1.7%。此外，与分层模型相比较，PlainMamba也取得了相似的性能。例如，当模型大小约为25M时，作者的模型比Swin-Tiny的准确率高出0.3%，验证了PlainMamba作为一个通用特征提取器的能力。另一方面，分层的VMamba [58]和其他分层 Transformer 确实比作者的准确率更高。正如作者在第1部分所解释的，分层模型在视觉识别方面的表现往往优于非分层模型。由于作者的主要动机是开发一个简单的Mamba架构，因此稍低的ImageNet准确率是可以接受的。

COCO object detection and instance segmentation

设置： 遵循[100]，作者测试了PlainMamba在COCO目标检测和实例分割上的能力，使用了两阶段的Mask R-CNN [38]和单阶段的RetinaNet [53]。对于这两个模型，作者都报告了1

\times

调度下的结果。遵循[100]，作者使用ViTAdapter [11]来计算多尺度特征以适应FPN网络结构。为了保持公平的比较，作者采用了[59]中提出的常用训练设置。

结果：作者在表3中报告了Mask R-CNN目标检测和实例分割的结果，在表4中报告了RetinaNet目标检测的结果。对于Mask R-CNN，PlainMamba与最先进架构持平。例如，在使用1

\times

训练计划时，具有相似FLOPs和更少参数的PlainMamba-L1达到了44.1 AP

{}^{bb}

和39.1 AP

{}^{mk}

，而Swin-Small则达到了44.8 AP

{}^{bb}

和40.9 AP

{}^{mk}

。作者还观察到分层模型通常比非分层模型表现得更好。尽管作者的模型性能低于某些分层模型，例如同时期的VAMaba [58]，但PlainMamba在其非分层对应模型中取得了最佳性能。例如，在使用1

\times

训练计划时，与配备了ViTAdapter [11]的DeiT-T相比，PlainMamba的AP

{}^{bb}

高出3.1，AP

{}^{mk}

高出1.6。这些结果证明了PlainMamba能够提取好的局部特征，这对于实例分割等目标 Level 任务非常重要。同样，PlainMamba在与单阶段RetinaNet目标检测器配合使用时也表现出色。例如，仅有一半模型大小和相似FLOPs的PlainMamba-L1，比Swin-Tiny的AP高出0.2。另一方面，作者也承认PlainMamba的性能不如分层VMMaba [58]。作者将这种劣势归因于基于FPN [52]的Mask R-CNN的多分辨率架构，这种架构更自然地适合分层设计。

ADE20K Semantic Segmentation

设置： 作者遵循常见做法[23; 59; 100]并以UperNet [94]作为分割网络来衡量PlainMamba在语义分割方面的能力。与XCiT [2]不同，作者没有将常量分辨率特征图显式地调整到多尺度。按照[59]，作者以批量大小16对所有模型进行160次迭代训练，并将默认的训练图像大小设置为512×512。

结果：作者在表5中报告了作者的模型在ADE20K语义分割任务上的表现。与ImageNet-1k和COCO实验类似，这里的竞争模型包括了三种视觉特征提取器中的分层和非分层 Backbone 网络。结果显示，PlainMamba在非分层模型中取得了最佳性能。例如，在相似的参数量下，PlainMamba-L2以更低的计算成本优于高分辨率（ Patch 大小为8）的XCiT-S12/8模型[2]。此外，PlainMamba-L2也比分层的Swin-Transformer-Tiny[59]表现更好，在模型尺寸和FLOPs更低的情况下获得了更高的mIoU。同时，PlainMamba也比同期的Vision Mamba[108]表现得更好。例如，PlainMamba-L2比ViM-S的mIoU高出1.9。这个结果验证了作者的模型在提取细粒度视觉特征方面的有效性，这对于像素级的语义分割任务至关重要。

Ablation Studies and Discussions

设置： 在这里，作者进行了消融研究，以测试作者的模型设计，并更深入地理解所提出的方法。在大多数实验中，作者使用了作者的L1模型，该模型拥有不到1000万的参数。这些模型都是在ImageNet-1K上进行预训练的，遵循第4.1节中描述的相同训练设置。

深度与宽度 在为给定的参数数量设计神经架构时，通常需要在网络的深度（即层数）和宽度（即特征维度）之间找到一个良好的平衡。尽管这个问题在现有架构中已经有所研究[93; 99]，但之前的结论是否适用于视觉SSMs仍然不清楚。在表6中，作者研究了所提的PlainMamba的深度与宽度之间的权衡。首先，结果表明，更深的模型往往比浅层模型表现得更好。例如，当参数数量约为7.4M时，12层模型比6层模型ImageNet top-1准确率高2.2%，而24层模型比12层模型再高1.1%。然而，当作者进一步增加深度到36层，同时相应减少宽度时，top-1准确率保持相似。另一方面，作者还注意到，更深的模型不如浅层但更宽的模型高效。例如，24层模型的FLOPs比12层模型高出0.3G。这些结果说明在网络深度和宽度之间保持良好平衡的必要性。

PlainMamba模块设计。在这里，作者通过将PlainMamba模块与Vision Mamba [108]和V Mamba [58]中的模块设计进行比较，来测试不同PlainMamba模块设计的有效性。为了公平比较，作者为所有设计使用了相同的模型深度和宽度设置，并采用相同的训练方案训练所有模型。作者还从Vision Mamba [108]模块中移除了CLS标记，并使用全局平均池化作为替代。作者在表7中报告了结果。可以看出，作者的设计取得了最佳效果。具体来说，Vision Mamba仅达到74.4%的ImageNet准确率，比作者的低3.5%。作者还注意到，带有Vision Mamba模块的模型劣于原始的Vision Mamba模型，这是由于移除了CLS标记所导致的。这些结果表明，即使在没有CLS标记的情况下，作者的模型仍然保持了其性能。此外，作者的设计与V Mamba模块[58]相比，准确率高出0.8%，这表明改进来自于作者提出的连续2D扫描和方向感知更新，这验证了作者在将SSM适应于2D图像方面所提出技术的有效性。

与ViT的效率对比 SSMs（例如Mamba）的一个特别优势是它们在保持效率的同时捕获全局信息的能力。在图4中，作者比较了PlainMamba与视觉 Transformer 的效率。具体来说，为了确保公平比较，作者创建了一个通道数为224的DeiT模型，得到一个拥有740万参数的模型，用于与PlainMamba-L1进行比较。具体而言，作者比较了模型在不同输入大小下的FLOPs和峰值推理内存。

结果显示，当输入大小扩展到高分辨率（例如4096×4096）时，作者的模型能够保持较低的运算成本。然而，DeiT在使用如此高分辨率输入时，其FLOPs和内存消耗会迅速增加。另一方面，作者还注意到，在使用低分辨率图像（例如128×128）时，作者模型的效率不如同等大小的DeiT。为了进一步研究它们在效率上的这种差异，作者将FLOPs分解为三部分[79]：1）_标记混合_，2）_通道混合_，和3）其他。具体来说，标记混合指的是DeiT中的多头注意力部分和PlainMamba中的选择扫描部分，通道混合指的是DeiT中的前馈网络和PlainMamba中的输入与输出投影。作者在表8中报告了这些结果。这些结果表明，在低和高分辨率下，PlainMamba的FLOPs在这三个部分中均匀分布。相反，当使用128×128输入时，DeiT的FLOPs主要由通道混合和其他两部分（可以忽略不计）组成。然而，由于自注意力操作的二次复杂性，当使用4096×4096输入时，DeiT在标记混合的FLOPs增长到23T，比PlainMamba贵23倍。这些结果验证了在使用高分辨率输入时PlainMamba的高效率。

5 Conclusion

在本文中，作者介绍了PlainMamba，一个基于SSM的简单视觉识别模型。作者的模型在概念上很简单，因为它不使用任何特殊标记，也没有分层结构，这使得它成为广泛使用的普通视觉 Transformer （vision transformer）的完美对应物。为了更好地将SSM适应于建模2D视觉特征，作者提出了两种技术，即（i）连续2D扫描和（ii）方向感知更新。作者在广泛的视觉识别任务上对PlainMamba进行了基准测试，包括图像分类、目标检测、实例分割和语义分割。结果显示，PlainMamba在性能上优于先前的非分层模型，包括并行的基于SSM的模型，并且可以与高性能的分层模型相媲美。此外，当用于高分辨率输入时，与视觉 Transformer 相比，PlainMamba显著减少了计算需求。作者希望作者的模型可以作为未来在这一领域研究的 Baseline 方法。