[论文解读] Vision-RWKV: Efficient and Scalable Visual Perception with RWKV-Like Architectures

Yuchen Duan, Wei‐Yun Wang+8 more|arXiv (Cornell University)|Mar 4, 2024

Advanced Image and Video Retrieval Techniques被引用 10

一句话总结

Vision-RWKV 将 RWKV 适配到视觉任务，采用四向平移和双向线性注意力，在分类和密集预测任务上保持强劲性能的同时，比 ViT 具有更低的计算/内存需求。

ABSTRACT

Transformers have revolutionized computer vision and natural language processing, but their high computational complexity limits their application in high-resolution image processing and long-context analysis. This paper introduces Vision-RWKV (VRWKV), a model adapted from the RWKV model used in the NLP field with necessary modifications for vision tasks. Similar to the Vision Transformer (ViT), our model is designed to efficiently handle sparse inputs and demonstrate robust global processing capabilities, while also scaling up effectively, accommodating both large-scale parameters and extensive datasets. Its distinctive advantage lies in its reduced spatial aggregation complexity, which renders it exceptionally adept at processing high-resolution images seamlessly, eliminating the necessity for windowing operations. Our evaluations demonstrate that VRWKV surpasses ViT's performance in image classification and has significantly faster speeds and lower memory usage processing high-resolution inputs. In dense prediction tasks, it outperforms window-based models, maintaining comparable speeds. These results highlight VRWKV's potential as a more efficient alternative for visual perception tasks. Code is released at https://github.com/OpenGVLab/Vision-RWKV.

研究动机与目标

• 在保留 RWKV 在长距离/全局信息处理方面的高效性的同时，激发一个适合视觉任务的编码器，降低视觉任务特定的复杂性。
• 调整 RWKV 以在不使用窗口化注意力的情况下处理高分辨率图像。
• 通过相对位置偏置、层缩放和额外的归一化来确保大规模训练的稳定性。
• 在分类、检测和分割等任务上展示在较低计算成本下的竞争性能。

提出的方法

• 引入 Q-Shift，实现四向令牌平移，以扩展二维图像的感受野。
• 用双向全局注意力替代因果 RWKV 注意力，在视觉场景中实现线性复杂度。
• 通过相对偏置和灵活衰减来调整注意力，以在扩展时保持稳定性。
• 引入层缩放和额外层归一化，以在不同深度和分辨率下稳定输出。
• 开发 VRWKV 的变体（Tiny 到 Large），具有固定的结构块和基于补丁的标记化。
• 展示基于 MAE 的预训练兼容性并在下游任务上进行微调。

实验结果

研究问题

• RQ1 在线性复杂度注意力下，Vision-RWKV 是否能达到甚至超越 ViT 在图像分类上的性能？
• RQ2 在高输入分辨率下，VRWKV 相较于窗口化或全局 ViT 是否能保持高效（FLOPs、内存、速度）？
• RQ3 在密集预测任务（检测/分割）中，VRWKV 主干是否以更低的计算成本保持竞争力？
• RQ4 VRWKV 是否能够利用 MAE 预训练在分类准确度上获取进一步提升？

主要发现

• VRWKV-T 在 224^2 上的 ImageNet-1K 达到 75.1% 的 top-1 准确率，参数为 6.2M，FLOPs 为 1.2G，较 DeiT-T 提升 2.9 点。
• VRWKV-L（384^2）在 ImageNet-1K 上达到 86.0% 的 top-1（在 ImageNet-22K 上预训练），参数 334.9M，FLOPs 189.5G，准确度超过 ViT-L。
• 在 COCO 目标检测中，VRWKV-L 获得 50.6 AP^b 和 44.9 AP^m（优于 ViT-L），且骨干FLOPs 更低。
• VRWKV-S（224^2）在检测/分割方面取得具有竞争力的结果，FLOPs 明显低于 ViT-S。
• 在语义分割（ADE20K）中，VRWKV 主干在全局注意力下的表现优于 ViT，同时降低了主干 FLOPs（例如 VRWKV-S 在 512^2 下以更少的 FLOPs 实现更高的 mIoU）。
• 使用双向平移的 MAE 预训练带来额外的小幅增益（例如在 ImageNet-1K 上，VRWKV-L 从 86.0 提升到 86.2）。