Vision-RWKV:基于RWKV架构的高效可扩展视觉感知模型
Vision-RWKV:基于RWKV架构的高效可扩展视觉感知模型
Vision-RWKV 是从 RWKV 改进而来的视觉模型，可以处理稀疏输入并高效执行健壮的全局处理。它能够很好地扩展到大型参数和数据集，优于其降低了空间聚合复杂性使得在高分辨率图像处理时无需窗口操作。在图像分类方面，VRWKV 在性能上优于 ViT，表现出更快的速度和更低的内存使用量，尤其是对于高分辨率输入。在密集预测任务中，它超越了基于窗口的模型，同时保持可比较的速度，因此在视觉感知任务中是一种更高效的替代方案。
架构
Vision-RWKV 支持稀疏输入和稳定的扩展，通过类似 ViT 的块叠加图像编码器设计，包括用于注意力和特征融合的空间混合和通道混合模块。VRWKV 通过将图像转换为补丁，添加位置嵌入来形成图像标记，然后通过 L 个相同的编码器层处理图像，保持输入分辨率。
视觉版本的 RWKV 修改了原始论文的注意力机制有三个关键变化：
- 引入双向注意力以确保所有标记彼此可见，通过将注意力范围扩展到整个标记集来将因果注意力转换为双向全局注意力；
- 基于标记之间的时间差异添加一个相对偏差，通过标记的总数进行归一化，以解决不同大小图像中的相对位置问题；
- 允许在指数项中有一个灵活的衰减参数，使模型能够关注距当前标记较远的标记，跨越不同通道；
使用指数衰减机制将全局注意力的复杂度从二次降低到线性，大幅提高了处理高分辨率图像的计算效率。为了更好地捕获图像中的二维关系，引入了四向标记偏移操作，这个操作在空间混合和通道混合模块中都有使用。它在四个方向上 - 上、下、左、右 - 在不同通道片段中线性插值每个标记与其相邻的标记。仅略微增加了计算需求，但显著扩展了标记的感受野，增强了模型后续层中的空间关系覆盖。
模型有以下几个不同大小的变体
为了解决增加层数和递归过程中指数项积累可能导致的模型输出和训练不稳定性问题，进行了两项修改：
- 有界指数 — 指数项除以标记数目，以保持衰减和增长在可控范围内；
- 额外的层归一化，添加在注意力机制和 Squared ReLU 操作之后，以防止更深层模型的输出溢出。
实验结果
在 ImageNet-1K 验证数据集上，VRWKV 模型在各种模型尺寸上展现出卓越性能，优于诸如 ViT 等分层和非分层的骨干模型，具有更高的 top-1 精度和较低或可比的计算复杂度。并且VRWKV-T 在准确性上超过了 DeiT-T，VRWKV-L 在减少计算成本的同时实现了比 ViT-L 更高的准确性。在大规模数据集上进行预训练进一步提升了 VRWKV 的性能，表明其可扩展性和作为传统 ViT 模型替代方案的潜力。
在 COCO 数据集上的目标检测中，VRWKV 显示出比 ViT 更好的性能，并且 FLOPs 明显较低，受益于其在密集预测任务中的全局注意力机制。
语义分割任务上它的效率和准确性上优于 ViT 模型，展示了其线性复杂度注意力机制的有效性。
其他贡献
在 ImageNet-1K 数据集上对 winy 大小的 VRWKV 模型进行消融研究，验证了 Q-Shift 和双向注意力等关键组件，显示出显著的性能改进。没有标记偏移，性能显著下降，而原始偏移方法仍落后于新引入的 Q-Shift。实施双向注意力使 top-1 精度提高了 2.3 个点。
有效感受野的分析突显出除使用原始 RWKV 注意力外的所有模型均实现了全局注意力，其中 VRWKV-T 在全局容量方面优于 ViT-T。Q-Shift 显著增强了感受野的核心范围，增强了全局注意力的归纳偏差。
效率分析表明，在更高的分辨率（高达 2048x2048）下，VRWKV-T 在 Nvidia A100 GPU 上的推理速度和内存使用率明显优于 ViT-T，这要归功于其线性注意力机制和类似 RNN 的计算框架。
在 ImageNet-1K 验证上，MAE 预训练进一步提升了 VRWKV 的性能，显示了其从稀疏输入和掩膜图像建模中受益的能力，增加了 top-1 精度。
论文地址：
https://arxiv.org/abs/2403.02308
代码
https://github.com/OpenGVLab/Vision-RWKV
作者：Andrew Lukyanenko·