Vision-RWKV 是从 RWKV 改进而来的视觉模型,可以处理稀疏输入并高效执行健壮的全局处理。它能够很好地扩展到大型参数和数据集,优于其降低了空间聚合复杂性使得在高分辨率图像处理时无需窗口操作。在图像分类方面,VRWKV 在性能上优于 ViT,表现出更快的速度和更低的内存使用量,尤其是对于高分辨率输入。在密集预测任务中,它超越了基于窗口的模型,同时保持可比较的速度,因此在视觉感知任务中是一种更高效的替代方案。
Vision-RWKV 支持稀疏输入和稳定的扩展,通过类似 ViT 的块叠加图像编码器设计,包括用于注意力和特征融合的空间混合和通道混合模块。VRWKV 通过将图像转换为补丁,添加位置嵌入来形成图像标记,然后通过 L 个相同的编码器层处理图像,保持输入分辨率。
视觉版本的 RWKV 修改了原始论文的注意力机制有三个关键变化:
使用指数衰减机制将全局注意力的复杂度从二次降低到线性,大幅提高了处理高分辨率图像的计算效率。为了更好地捕获图像中的二维关系,引入了四向标记偏移操作,这个操作在空间混合和通道混合模块中都有使用。它在四个方向上 - 上、下、左、右 - 在不同通道片段中线性插值每个标记与其相邻的标记。仅略微增加了计算需求,但显著扩展了标记的感受野,增强了模型后续层中的空间关系覆盖。
模型有以下几个不同大小的变体
为了解决增加层数和递归过程中指数项积累可能导致的模型输出和训练不稳定性问题,进行了两项修改:
在 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·
本文分享自 DeepHub IMBA 微信公众号,前往查看
如有侵权,请联系 cloudcommunity@tencent.com 删除。
本文参与 腾讯云自媒体同步曝光计划 ,欢迎热爱写作的你一起参与!