Vision-RWKV: Efficient and Scalable VisualPerception with RWKV-Like Architectures

摘要: Transformer已经彻底改变了计算机视觉和自然语言处理，但它们的高计算复杂性限制了它们在高分辨率图像处理和长上下文分析中的应用。本文介绍了视觉-RWKV (VRWKV)模型，它是在自然语言处理领域使用的RWKV模型的基础上，对视觉任务进行了必要的修改。与视觉转换器(ViT)类似，我们的模型旨在有效地处理稀疏输入并展示强大的全局处理能力，同时也有效地扩展，适应大规模参数和广泛的数据集。它的独特优势在于降低了空间聚合的复杂性，这使得它非常擅长无缝地处理高分辨率图像，消除了窗口操作的必要性。我们的评估表明，VRWKV在图像分类方面的性能超过了ViT，并且在处理高分辨率输入方面具有更快的速度和更低的内存使用。在密集预测任务中，它优于基于窗口的模型，保持相当的速度。这些结果突出了VRWKV作为视觉感知任务更有效的替代方案的潜力,代码发布在https://github.com/OpenGVLab/Vision-RWKV

主要贡献：

(1)我们提出VRWKV作为ViT的低成本替代品，以更低的计算成本实现全面替代。我们的模型不仅保留了ViT的优点，包括捕获远程依赖关系的能力和处理稀疏输入的灵活性，而且还将复杂性降低到线性水平。这种显著的减少消除了在处理高分辨率图像时需要基于窗口的注意力，使VRWKV成为视觉任务更有效和可扩展的解决方案。

(2)为了适应视觉任务，我们引入了双向全局注意和一种新的标记转移方法Q-Shift，实现了全局注意的线性复杂性。为了保证稳定的可扩展性，我们做了一些努力，包括在注意机制中使用相对位置偏差来避免溢出，在我们的模型中采用层尺度，以及在关键矩阵的计算中增加额外的层规范化。

(3)我们的模型超越了基于窗口的vit，与全局关注vit相当，随着分辨率的增加，显示出更低的FLOPs和更快的处理速度。值得注意的是，仅在ImageNet-1K上训练的VRWKV-T达到了75.1%的top-1准确率[8]，比DeiT-T[49]高出2.9分。在大规模参数(即335M)和训练数据(即ImageNet-22K)下，VRWKV-L的top-1准确率进一步提升至86.0%，高于ViT-L [12] (86.04 vs. 85.15)。此外，在具有挑战性的下游基准COCO[27]上，我们的最佳模型VRWKV-L达到50.6%的箱形mAP，比viti - l高出1.9个点(50.6比48.7)。

网络架构：

(a) VRWKV架构包括L个相同的VRWKV编码器层、一个平均池化层和一个线性预测头。(b) VRWKV编码器层的详细信息。Q-Shift表示用于视觉任务的四向移位方法。“Bi-WKV”模块作为一个双向RNN单元或一个全局注意机制。

如(b)所示，与RWKV不同的是，为了适应视觉任务，作者提出了Q-Shift模块和Bi-WKV模块，我们首先解释一下Q-Shift：

先看论文里的公式：

下标(∗)∈{R, K, V}表示由可学习向量µ(∗)控制的X和X†的3个插值，分别用于后面的R, K, V的计算。h和w表示令牌X的行和列索引，“:”是排除结束索引的切片操作。Q-Shift使不同通道的注意机制在不引入许多额外flop的情况下获得内部关注相邻令牌的优先权。Q-Shift操作还增加了每个令牌的接受域，从而大大增强了令牌在后一层的覆盖范围。

详细解释一下：X是一个形状为[H, W, C]的三维张量，代表一个具有高度H、宽度W和通道数C的图像或特征图。在这个公式中，X†[h, w]代表了在位置(h, w)处的输出。Concat表示连接操作，它将四个张量按照指定的轴连接起来。

具体地说，X†[h, w]是由四个部分组成的，这四个部分分别来自于输入张量X的不同位置：

1. X[h - 1, w, 0 : C/4]：这个部分来自于输入张量X的位置(h - 1, w)，它包含了从通道0到C/4的所有通道的数据。
2. X[h + 1, w, C/4 : C/2]：这个部分来自于输入张量X的位置(h + 1, w)，它包含了从通道C/4到C/2的所有通道的数据。
3. X[h, w - 1, C/2 : 3C/4]：这个部分来自于输入张量X的位置(h, w - 1)，它包含了从通道C/2到3C/4的所有通道的数据。
4. X[h, w + 1, 3C/4 : C]：这个部分来自于输入张量X的位置(h, w + 1)，它包含了从通道3C/4到C的所有通道的数据。    这样，通过连接这四个部分，就得到了X†[h, w]处的输出。

接下来我们解释一下Bi-WKV模块：

这里，T表示tokens的总数，等于HW/p2, w和u是两个c维可学习向量，分别表示通道方向的空间衰减和指示当前标记的奖励。kt和vt表示K和V的第t个特征。

求和公式表明，输出wkvt是V沿着标记维数从0到T−1的加权和，得到c维向量。

它表示通过对第t个token应用注意力操作获得的结果。权重由空间衰减向量w、标记(|t−i|−1)/ t和ki之间的相对偏差共同决定。

我们也可以将这个模块看作RNN的形式：

将式(5)中分母和分子的求和项以t为边界进行拆分，可以得到4种隐藏状态:

隐藏状态的更新只需要加或减一个求和项，乘或除e - w/T，第t个输出结果为：

实验：

我们全面评估了VRWKV方法在性能、可扩展性、灵活性和效率方面对ViT的可替代性。我们在广泛使用的图像分类数据集ImageNet上验证了我们模型的有效性[8]。对于下游密集预测任务，我们选择COCO[27]数据集上的检测任务和ADE20K[62]数据集上的语义分割任务。

1、图像分类任务

2、目标检测

3、语义分割

消融实验：

结论：

我们提出了视觉- rwkv (VRWKV)，一个具有线性计算复杂度注意机制的视觉编码器。我们证明了它在综合视觉任务中作为ViT的替代主干的能力，包括分类、密集预测和掩膜图像建模预训练。VRWKV具有相当的性能和可扩展性，具有较低的计算复杂度和内存消耗。得益于其低复杂度，VRWKV可以在ViT难以承担全局关注的高计算开销的任务中实现更好的性能。我们希望VRWKV能够成为一种高效、低成本的ViT替代品，展示线性复杂变压器在视觉领域的强大潜力。

结语：

RWKV的线性复杂度是这个模型的最大亮点，大模型部署在普通台式机和笔记本成为了可能，甚至是在手机端部署。RWKV将推动大模型进行一次架构迁移！