[論文レビュー] 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 を視覚タスクに適用するために、4方向シフトと双方向線形アテンションを用いて調整し、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 をウィンドウベースのアテンションを用いずに高解像度画像に対応させるよう適用。
• 相対的位置バイアス、レイヤースケール、追加正規化を用いて大規模トレーニングの安定性を保証する。
• 分類・検出・分割の各タスクで、低い計算コストで競争力のある性能を示す。

提案手法

• 2D画像における受容野を拡張するための四方向トークンシフトQ-Shiftを導入。
• 視覚設定で計算量を線形にするため、因果的RWKVアテンションを双方向グローバルアテンションに置換。
• 拡大時の安定性を維持するため、相対バイアスと柔軟な減衰を用いてアテンションを変更。
• 深さと解像度をまたいで出力を安定化させるためにレイヤースケールと追加レイヤー正規化を組み込む。
• 固定されたアーキテクチャブロックとパッチベースのトークン化を用いた VRWKV バリアント（Tiny〜Large）を開発。
• MAE に基づく事前学習の適合性と下流タスクでのファインチューニングを示す。

実験結果

リサーチクエスチョン

• RQ1 Vision-RWKV は線形複雑度アテンションで画像分類において ViT の性能に匹敵するか、それを上回れるか？
• RQ2 VRWKV はウィンドウ型またはグローバル ViT と比較して高入力解像度でのFLOPs・メモリ・速度の効率を維持するか？
• RQ3 VRWKV バックボーンは計算コストを抑えつつ密な予測タスク（検出・セグメンテーション）で競争力があるか？
• RQ4 MAE pre-training を活用して分類精度をさらに向上させることができるか？

主な発見

• VRWKV-T は 224^2 で ImageNet-1K の Top-1 精度 75.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) は ViT-S より明らかに低い FLOPs で競争力のある検出/分割結果を達成。
• セマンティックセグメンテーション（ADE20K）では、VRWKV バックボーンはグローバルアテンションを用いる ViT を上回りつつ、バックボーン FLOPs を削減（例：512^2 の VRWKV-S はより少ない FLOPs でより高い mIoU を達成）。
• MAE 前処理学習は追加の小さな利得をもたらす（例：ImageNet-1K で VRWKV-L が 86.0 から 86.2 へ）。