集智书童 | 从 Mamba 到 RWKV-SAM | 速度提高2倍，分类和语义分割都优于 Mamba 模型 ！

本文来源公众号“集智书童”，仅用于学术分享，侵权删，干货满满。

原文链接：从 Mamba 到 RWKV-SAM | 速度提高2倍，分类和语义分割都优于 Mamba 模型 ！

基于Transformer的分割方法在处理高分辨率图像时面临着有效推理的挑战。最近，一些线性注意力架构，如Mamba和RWKV，因其能够高效处理长序列而受到广泛关注。在这项工作中，作者通过探索这些不同的架构来设计一个高效的“分割任何物体”模型。

具体来说，作者设计了一个混合骨架，包含卷积和RWKV操作，它在准确性和效率上都取得了最佳效果。

此外，作者还设计了一个有效的解码器，利用多尺度标记来获得高质量的 Mask 。

作者将作者的方法称为RWKV-SAM，这是一个简单、有效、快速的SAM-like模型的 Baseline 。

此外，作者构建了一个包含各种高质量分割数据集的基准，并使用这个基准共同训练一个既高效又高质量的分割模型。

根据基准结果，与Transformer和其他线性注意力模型相比，作者的RWKV-SAM在效率和分割质量上都表现出色。

例如，与同规模的Transformer模型相比，RWKV-SAM的速度提高了2倍以上，并且在各种数据集上能够实现更好的分割性能。此外，RWKV-SAM在分类和语义分割结果上优于最近的视觉Mamba模型。

1 Introduction

在大型分割数据集上训练的Segment Anything Model（SAM）[23]最近因其跨多种分割任务的显著通用性和有效性而受到广泛关注。通过接收人类或其他模型提供的视觉提示，如点和框作为输入，SAM可以在各种场景中生成 Mask ，从而启用了诸如图像编辑[9]、遥感[14]、医学图像分割[38]等众多下游应用。尽管SAM具有强大的泛化能力，但它仍存在一些可能阻碍在某些场景中实际应用的缺点。首先，SAM的计算成本极高。其次，SAM的分割质量在某些情况下仍然不足；例如，SAM总是生成过于平滑的边缘，这在很多情况下并不适用。以上两个缺点限制了SAM在实时场景和高质量分割结果需求领域的应用。

现有工作通常只关注解决第一个问题或第二个问题。例如，一些研究[80; 71; 63]，如EdgeSAM[80]和Efficient SAM[63]，旨在探索SAM的高效架构。然而，分割质量仍然有限。另一方面，有几项工作[22; 58]探索了高分辨率和高质量的SAM。它们给SAM带来了额外的计算成本，从而减慢了推理速度。因此，应该在高质量和高效率之间寻求平衡，以便更好地在现实世界中部署SAM。

最近，一系列始于自然语言处理领域（如RWKV[47]，Mamba[10]）并延续到计算机视觉领域（如VMamba[34]，Vision-RWKV[8]）的工作开始关注设计能够在线性时间内处理长距离依赖的方法（线性注意力模型）。与计算复杂度随序列长度增加而呈平方增长的 Transformer 相比，线性注意力模型重新构建了注意力机制，使其与序列长度成线性比例增长，从而在序列非常长时显著降低计算成本。因此，线性注意力模型可以处理非常长的序列，同时保持其全局感知能力。然而，以前没有研究探索将这些架构应用于类似SAM的可提示分割任务。

在这项工作中，作者尝试同时解决这些问题，利用最近的线性注意力模型构建一个高效且高质量的SAM。特别是，作者提出了RWKV-SAM来处理SAM的计算成本和分割质量问题。SAM的高计算成本可以归因于两个原因：1)庞大的参数数量，2)由于输入特征尺寸增长， Transformer 层中注意力设计带来的二次时间复杂度。尽管先前的工作通过减少模型大小（例如，EfficientSAM[63]）来处理SAM的效率问题，但这些解决方案仍然面临二次时间复杂度，这意味着它们无法在高分辨率输入（例如高分辨率1024x1024输入）上实现良好的效率。作者提出了一个利用RWKV[47]的高效分割主干，以提高在高分辨率下的效率同时保持全局感知。作者的高效分割主干包含三个阶段，解码器可以利用这些阶段来细化生成的 Mask 。此外，作者探索了不同的解码器设计来融合不同尺度的特征，并在结合高质量数据集上训练模型，使作者的RWKV-SAM成为一个高质量和高效的分割任何模型。

作者在各种数据集和基准上评估了作者的方法。如图1所示，作者的RWKV-SAM在效率和质量上都超过了以前的方法。尽管它只需要SAM大约1/16的推理时间，但作者的方法实现了更准确、高质量的分割结果。尽管模型大小与EfficientSAM相当，但作者的RWKV-SAM运行速度超过了2倍。由于来自主干的低层次局部特征的信息，作者的方法在细节上优于HQ-SAM。与先前的线性模型（如Mamba）相比，作者的RWKV-SAM在运行速度和多种基准上的表现都更胜一筹。当使用极高分辨率输入时，作者的模型运行得更快。

作者对这项工作做出了以下贡献：

(1)作者提出了RWKV-SAM，它包含一个高效的分割主干，产生不同分辨率的特征图，并利用RWKV操作降低时间复杂度。

(2)作者探索了不同的设计来利用解码器中的多尺度特征图，并在高质量分割数据集上训练RWKV-SAM，以实现高质量的分割能力。

(3)作者在几个基准上证明了RWKV-SAM的有效性，超过了以前的方法同时保持了效率。

(4)作者对各种线性注意力模型进行了详细的比较研究，包括Vision Mamba和VRWKV。

据作者所知，这是首次以公平比较的方式探索这些模型的工作。

2 Related Work

高效分割。现有的方法主要集中在对闭合集和特定领域[27]的高效分割研究。大部分关于高效分割的研究致力于驾驶场景。此外，还有多项研究针对高效全景分割和快速视频实例分割。最近，各种研究开发了高效分割技术，以促进移动设备上的模型执行针对“分割任何模型”。Mobile SAM [71]介绍了一种简化的编码器蒸馏方法。Fast SAM [78]采用单阶段实例分割框架，直接解码类别无关的 Mask 。Edge SAM [80]使用新的提示引导蒸馏，在现实世界移动设备上部署SAM模型。Efficient SAM [63]在这项工作中，除了实时性约束外，作者还追求高质量的分割。

高效 Backbone 网络。这一方向主要集中发展高效的CNNs， Transformer 和混合架构，以学习视觉表示。最近，一些研究探索了线性注意力模型，包括视觉中的RWKV 和Mamba。然而，所有这些工作都试图替换 Transformer 进行表示学习，忽略了在不同尺度上生成特征。作者探索了一种用于快速、高分辨率分割的高效 Backbone 网络，其中作者采用了CNN和RWKV混合架构。根据实验，作者提出的 Backbone 网络在类似的参数和延迟约束下实现了更好的表示。

高质量分割。之前针对高质量分割的工作通过设计特定模块，提出细粒度数据集，关注以目标为中心的设置，以及添加细化器，针对特定任务。为了允许更开放的设置，一些研究[22; 33; 58]探索了将SAM作为基础模型来提高分割质量。然而，它们无法实时运行。特别是，HQ-SAM[22]与原始SAM相比带来了额外的成本。与这些工作相比，作者有两个目标。一个是设计一个新的模型，以实时分割高质量的目标 Mask 。另一个是建立一个完整的训练流程，包括数据集，以使高效模型能够进行高质量分割。

线性注意力模型。当标记数量增加时， Transformer 存在计算成本问题，这正是高质量分割面临的挑战。最近，一些研究表明了替代 Transformer 架构的巨大潜力。特别是，状态空间模型[10; 11]已被证明能够建模长距离依赖。此外，RWKV[47; 48]是另一种具有更快推理速度的方法。作者旨在探索这些架构，在“分割任何模型”元架构下进行高效、高质量分割。特别是，作者发现在高效分割设置下，针对高分辨率图像输入，RWKV比Mamba运行得更快。因此，作者旨在将RWKV架构作为作者的 Backbone 网络进行探索。

3 Method

概述。作者旨在构建一个高效、高质量的“分割任何物体”模型。这要求模型具备以下特性：

首先，模型应该有一个即使在高清分辨率下也高效的 Backbone 网络；

其次，模型应能够利用现有的SAM知识，避免在完整的SA-1B [23]上进行训练；

第三，模型应能够利用来自 Backbone 网络的特征金字塔，并使用高质量数据进行训练以生成高质量的 Mask 。

为了构建一个满足这三个特性的模型，作者设计了一个基于RWKV的 Backbone 网络（第3.1节），该网络具有特征金字塔，在高清分辨率下高效，并且与其他 Transformer 或线性注意力模型相比具有较好的性能（请参考表3和图3）。在第3.2节中，作者介绍了作者的训练流程，以利用来自SAM模型和高质量数据集的知识。作者还展示了用于融合来自不同分辨率的特征的解码器。

Efficient Segmentation Backbone

原始的Segment Anything模型[23]采用了一种基于 Transformer 的架构。尽管它取得了强大的性能，但计算开销巨大。尽管EfficientSAM[63]极大地减少了模型参数的数量，但它并未改变视觉 Transformer 的架构，在处理高分辨率时仍然需要较长的推理时间。其主要原因在于视觉 Transformer 架构的内在属性。随着分辨率的提高， Patch 的数量呈二次增长，导致计算需求增加。

为了在较高分辨率下构建一个高效的分割 Backbone 网络，作者遵循NLP社区[10; 47]中的线性时间序列建模的精神，提出了基于RWKV的高效视觉 Backbone 网络。通常，作者的 Backbone 网络具有3阶段设计，并包含两种类型的块：移动卷积块（MBConv）[55]和视觉-RWKV块（VRWKV）[8]。

表1：不同 Backbone 网络变体的设置。

宏观层面设计。图2展示了作者高效分割 Backbone 网络的概览。作者 Backbone 网络的宏观层面设计是由ViTamin[1]启发而来。在前两个阶段，作者采用基于卷积的块，即MBConv，生成高分辨率特征图。作者在每个阶段之前将特征图下采样2。高分辨率特征图可用于解码器中的 Mask 细化。在第三个阶段之前，特征图的下采样因子为16，这意味着特征图中的每个像素可以被视为一个“标记”。在第三个阶段，作者堆叠了一系列VRWKV块，以标记作为输入。与普通的视觉 Transformer 或Vision-RWKV[8]相比，作者的 Backbone 网络具有不同尺度的特征图，而不是单一的固定分辨率。作者在表1中呈现了不同变体的 Backbone 网络设置。这些多尺度特征图允许作者的模型自适应地关注各种空间细节，增强了处理不同目标大小的复杂场景的能力。微观层面设计。MBConv块使用了“倒置瓶颈”设计[55]。它包含一个1 x 1卷积来扩展通道尺寸，一个3 x 3深度卷积进行空间混合，以及另一个1 x 1卷积将通道投影回原始通道尺寸。遵循ViTamin[1]，作者为了简化使用LayerNorm而不是BatchNorm。作者在MBConv块中将扩展比例设置为4。

RWKV-SAM: Data, Model, and Training Pipeline

重新审视SAM。 原始的Segment Anything Model（SAM）包含一个重的ViT-H [7] 主干网络，一个接收框或点作为视觉提示的提示编码器，以及一个轻量级的解码器，包含在16倍下采样步幅上的两个Transformer层。轻量级解码器接收主干网络的输出和提示编码器，并生成相应的 Mask 。SAM在包含1100万张图片的大型自动标注SA-1B数据集上进行训练，这需要256个A100 GPU连续68小时。为了构建一个高效、高质量的分割任何物体的模型，作者的RWKV-SAM涉及到训练数据、模型结构和训练流程的设计。

表2：训练数据集。

训练数据。 SAM使用的SA-1B [23] 标注是自动生成的。尽管它有助于扩大训练数据规模，但标注不包含细节。为了缩小这一差距，作者引入了三个异构数据集进行联合训练。第一个数据集是COCONut-B [6]，包含242K张图片，包括COCO [32] 标注和未标注的图片。COCONut-B的标注是通过辅助的手动标注流程生成的，这产生了超出原始标注的高质量标注。第二个数据集是EntitySeg [49]，它包含30k高分辨率（2000px至8000px）由人类标注员标注的图片。第三个数据集是DIS5K [51] 数据集。DIS5K数据集提供了卓越的单个目标、高度准确的标注。

4 Experiments

Experimental Setup

数据集。 作者的训练包括如第3.2节所述的两个阶段。在第一阶段，作者使用1%的SA-1B数据来提炼高效的分割 Backbone 网络。作者对1%的SA-1B进行24个周期的训练（相当于在COCO上进行24个周期）。在第二阶段，作者使用COCONut-B [6]，EntitySeg [49]和DIS5K [51]数据集的组合进行训练。由于不同数据集之间的数据样本有所不同，作者将EntitySeg和DIS5K数据集重复，以使三个数据集按2:1:1的比例平衡，每个周期产生482k个样本。作者对组合数据集进行6个周期的训练（相当于在COCO上进行24个周期）。

评估协议。 由于作者的方法是在包括单目标和高复杂性场景的各种高质量数据集上训练的，作者考虑不同场景选择了三种类型的基准进行评估。第一个基准是在COCO数据集上。作者使用强大的检测器ViTDet-H [30]生成边界框作为视觉提示，作为Segment Anything模型中的提示输入。第二个基准是在DIS [51]数据集（验证集）上。使用由 Mask 标注生成的边界框作为提示输入，作者在单目标高质量数据集上测试生成 Mask 的性能。第三个基准也包括单目标数据集，但包含了COIFT [31]和HR-SOD [70]数据集，以测试在单目标高质量数据集上的零样本性能。作者在COCO上报告 Mask 平均精度（mAP）和 Mask 边界平均精度（mBAP），并在单目标分割数据集上报告mIoU和边界mIoU（mBIoU）。

Main Results

图像网预训练. 作者提出的有效分割 Backbone 网络首先在ImageNet-1K [54] 数据集上进行预训练。作者按照 VRWKV [8] 的方法，使用 Swin-Transformer [35] 的训练策略，以224x224的分辨率训练 Backbone 网络120个周期。为了验证 Backbone 网络的性能，作者在验证集上测试了ImageNet分类性能。如表4所示，在参数数量较少的情况下，作者的方法与基于Mamba或RWKV的先前方法相比，性能更好或相当。例如，当与小版本的Vim [83]，VRWKV [8]和RWKV-SAM进行比较，在相同的嵌入维度384下，尽管作者的模型尺寸更小，但作者的方法仍然优于先前的方法。作者认为这种改进可能得益于作者的宏观级设计，它在前两个块中使用卷积层来获取不同尺度的特征，而不是直接通过将图像转换为图像块进行下采样。

语义分割. 为了测试 Backbone 网络在语义分割方面的性能，作者使用预训练的 Backbone 网络作为特征提取器，并集成了 UperNet [62] 作为解码器。作者在 ADE20K 数据集 [79] 上以相同的设置训练模型160k次迭代。对于 RWKV-SAM，作者对 RWKV-SAM  Backbone 网络进行了轻微修改，在第三阶段的末尾加入了一个 MBConv 块以生成最小尺度（1/32）的特征。这个特征与有效分割 Backbone 网络前三个块的输出相结合，形成了一个特征金字塔。

如表5所示，作者报告了与Vim [83]的比较结果。结果表明，即使参数更少，作者的方法性能也更好。

分割任何模型. 作者将配备了有效分割 Backbone 网络的 RWKV-SAM 与第4.1节提到的基准测试中的先前工作进行对比。如表3所示，尽管 SAM [23] 和 EfficientSAM [63] 在 COCO 数据集上表现出良好的性能，但与作者的方法和 HQ-SAM [22] 在高质量数据集上的表现相比，它们略显不足。在 COIFT [31] 和 HR-SOD [70] 数据集上，尽管作者的训练数据集不包含同一领域的数据样本，但作者的方法与更大的 SAM [23] 相比，仍然展示了强大的零样本性能。作者还使用分割任何解码器训练了先前的线性注意力模型。VRWKV-S [8] 也使用了 VRWKV 块，但它只有一个特征尺度。因此，它不能使用高分辨率特征来优化分割结果。根据结果，它在高质量数据集上的性能相对较差。使用双向mamba层 [10] 的 Vim [83] 在 COCO 数据集的边界度量上也有相对较差的性能。

Ablation Study and Analysis

效率分析。在图3中，作者报告了不同输入图像尺寸下的RWKV-SAM-Small和ViT-Small（由EfficientSAM [63]使用）的延迟。这些结果是在单个NVIDIA A100 GPU上测试的。这两个 Backbone 网络都采用了384的嵌入维度。在相对较低的输入分辨率下，两个模型的FPS相似（例如，以512x512图像作为输入，ViT-Small为86.8 FPS，而RWKV-SAM-Small为73.4 FPS）。然而，当输入图像的分辨率增加时，ViT模型的延迟呈二次增长。相比之下，RWKV-SAM-Small的延迟线性增长，这意味着在输入图像尺寸较大时具有优势。如图3所示，当使用高分辨率图像作为输入时，RWKV-SAM的延迟仍然比ViT有显著优势。在segment anything模型中，典型的输入尺寸是1024x1024。尽管使用相同的输入尺寸，由于更有效的计算机制，RWKV-SAM-Small实现的FPS（40.3）显著高于ViT-Small（17.8）。因此，使用类似规模的模型（ViT-Small：22.4M，RWKV-SAM-S：19.7M），RWKV-SAM作为segment anything模型的 Backbone 网络具有更高的效率。

解码器中融合模块的影响。在解码器中，作者探索了不同的设计来融合不同尺度的特征。第一个设计为每个尺度使用两个卷积层以下采样低级特征并调整通道，然后沿通道维度融合三个特征后的两个卷积层。第二个设计用RWKV块替换卷积层，这使得特征具有全局感知能力。第三个设计逐渐从低分辨率特征融合到高分辨率，并使用DCN进行融合，遵循FaPN [19]的设计。这种方法可能使模型有更多机会捕捉相对较远距离的信息。如表6（左）所示，第一个设计在实现与其他设计相当性能的同时具有最佳的效率。第二个设计可能会损害性能，表明在解码器中使用RWKV块不如使用卷积等局部操作有效。作者怀疑这可能是因为RWKV块破坏了图像特征的连续性。因此，作者使用第3.2节提到的第一个设计。

编码器设计的消融研究。如第3.2节所述，作者使用了3阶段设计。在前两个阶段，作者使用MBConv块学习低级表示。作者探索了向第一阶段添加一些RWKV块以评估宏观级设计的影响。在表6（右）中，作者报告了在1024x1024输入尺寸下的FPS和替代设计的准确度。RWKV-SAM Backbone 网络总共有14个RWKV块，默认情况下它们位于第三阶段。在表中，（x-y）表示将x和y块放在前两个阶段，其余块在第三阶段。根据结果，尽管模型尺寸有所减小（前两个阶段的隐藏嵌入尺寸较小），但在1024x1024输入尺寸下的推理速度较慢。这种设计也导致模型性能下降。因此，作者将RWKV块放在第三阶段以获得更好的推理速度。

可视化。在图4中，作者将作者的方法与SAM [23]和HQ-SAM [22]进行了比较。根据可视化结果，作者观察到即使是在非常复杂的场景中，作者的方法也实现了更优的分割质量。例如，在第一个例子中，SAM无法区分围栏门，而HQ-SAM失去了很多细节。相比之下，作者的RWKV-SAM在细节方面取得了最佳结果。

5 Conclusion

在本文中，作者开发了RWKV-SAM，它包括一个高效的分割 Backbone 网络和一个完整的训练流程，以使分割任何模型具备高质量分割能力。得益于线性复杂度，作者的方法在保持强大性能的同时，在高分辨率上实现了卓越的效率。

在所提出的基准上进行训练后，作者的RWKV-SAM展示了卓越的高质量分割性能。

作者还对最近提出的线性注意力模型进行了基准测试，包括Mamba和RWKV，显示RWKV-SAM在它们之中表现出色。作者的RWKV-SAM能够以高质量和高效率分割任何物体，使其成为一个通用的分割工具，可以促进下游应用。作者希望作者的研究能激发使用线性注意力模型的新架构设计，用于密集预测任务。

6 参考

[1].Mamba or RWKV: Exploring High-Quality and High-Efficiency Segment Anything Model.

THE END !

文章结束，感谢阅读。您的点赞，收藏，评论是我继续更新的动力。大家有推荐的公众号可以评论区留言，共同学习，一起进步。

原文地址：https://blog.csdn.net/csdn_xmj/article/details/142444746

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！