集智书童 | 从 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成为一个高质量和高效的分割任何模型。

作者在各种数据集和基准上评估了作者的方