Vision-RWKV: Efficient and Scalable Visual Perception with RWKV-Like Architectures

Vision-RWKV: Efficient and Scalable Visual Perception with RWKV-Like Architectures

相关链接：arxiv github
关键字：Vision-RWKV、Visual Perception、Linear Attention、RWKV、Transformer

摘要

Transformers 在计算机视觉和自然语言处理领域引起了革命，但它们的高计算复杂度限制了它们在高分辨率图像处理和长上下文分析中的应用。本文介绍了 Vision-RWKV (VRWKV)，这是一个从NLP领域的RWKV模型改编而来，对视觉任务做了必要修改的模型。与ViT(Vision Transformer)类似，我们的模型旨在高效处理稀疏输入，并展示出强大的全局处理能力，同时也能有效地扩展，适应大规模参数和大量数据集。它的独特优势在于它降低了空间聚合的复杂性，使其能够无缝处理高分辨率图像，无需窗口操作。我们在图像分类评估中证明了VRWKV在显著提高速度和降低内存使用的同时，能达到与ViT相当的分类性能。在密集预测任务中，它也优于基于窗口的模型，并保持了相近的速度。这些结果突显了VRWKK作为视觉感知任务更高效替代方案的潜力。

核心方法

• Quad-directional Shift (Q-Shift): 对视觉任务量身定制的数据转换方法，通过四向移位和线性插值扩展了单个代币的语义范围。
• Bidirectional Global Attention Mechanism: 转换了原始的RWKV注意力机制为双向全局注意力机制，以线性计算复杂度在RNN形式的前向和后向中计算全局注意力。
• 在RWKV注意力机制中进行调整，去除了衰变向量的限制，将绝对位置偏差转化为相对偏差，增强了模型处理能力的同时确保了可扩展性和稳定性。
• 为了稳定模型的输出，在不断扩大的网络中引入了额外的层归一化（Layer Normalization）。

实验说明

以下是部分实验结果的展示：

Method	#Param	FLOPs	Top-1 Acc
DeiT-T [48]	5.7M	1.3G	72.2
DeiT-S [48]	22.1M	4.6G	79.9
DeiT-B [48]	86.6M	17.6G	81.8
ViT-L [13]	309.5M	191.1G	85.2
VRWKV-T	6.2M	1.2G	75.1
VRWKV-S	23.8M	4.6G	80.1
VRWKV-B	93.7M	18.2G	82.0
VRWKV-L	334.9M	189.5G	85.3

详细说明：

• 参数量（#Param）和计算量（FLOPs）表示了模型的大小和处理一个224×224分辨率图像时的工作量。
• Top-1 准确率是指在ImageNet-1K数据集上进行图像分类任务的准确率。
• 可以看出，VRWKV在各个尺寸的模型中与ViT表现相当，但在计算资源使用上更为高效。

结论

我们提出的Vision-RWKV (VRWKV) 是一个具有线性计算复杂度注意力机制的高效视觉编码器。我们证明它在包括分类、密集预测和掩模图像建模预训练等综合视觉任务中可以作为ViT的一个替代后骨架。与ViT相比，VRWKV展示了类似的性能和可扩展性，并且具有更低的计算复杂性和内存消耗。得益于它的低复杂性，VRWKV在一些ViT难以负担全局注意力高计算开销的任务中，可以实现更好的性能。我们希望VRWKV能成为ViT的一个高效和低成本的替代方案，在视觉领域展示出线性复杂度Transformers的强大潜力。