[논문 리뷰] Vision-RWKV: Efficient and Scalable Visual Perception with RWKV-Like Architectures

Yuchen Duan, Wei‐Yun Wang외 8명|arXiv (Cornell University)|2024. 03. 04.

Advanced Image and Video Retrieval Techniques인용 수 10

한 줄 요약

Vision-RWKV는 RWKV를 시각 분야에 맞게 4방향 시프트와 양방향 선형 주의(attention)를 도입하여 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를 윈도 기반 주의 없이 고해상도 이미지를 처리하도록 적응시킨다.
• 상대 위치 편향(relative positional bias), 계층 스케일(layer scale), 추가 정규화를 통해 대규모 교육의 안정성을 보장한다.
• 더 낮은 계산 비용으로 분류, 탐지 및 분할에서 경쟁력 있는 성능을 입증한다.

제안 방법

• 2D 이미지에서 수용 영역을 확장하기 위한 사방향 토큰 이동(Q-Shift)을 도입한다.
• 비전 설정에서 선형 복잡도를 달성하기 위해 인과적 RWKV 주의를 양방향 글로벌 주의로 대체한다.
• 확대 시 안정성을 유지하기 위해 상대 편향(relative bias)과 유연한 감쇠를 갖춘 주의를 수정한다.
• 깊이와 해상도에 걸쳐 출력을 안정시키기 위해 계층 스케일과 추가 계층 정규화를 도입한다.
• 고정된 아키텍처 블록과 패치 기반 토큰화로 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)는 ViT-S보다 현저히 낮은 FLOPs로 탐지/세분화에서 경쟁력 있는 결과를 달성한다.
• ADE20K의 시맨틱 세분화에서 VRWKV 백본은 전역 주의를 사용하는 ViT를 능가하면서 백본 FLOPs를 줄인다(예: VRWKV-S 512^2에서 더 적은 FLOPs로 더 높은 mIoU를 달성).
• MAE 사전 학습은 bidirectional shift를 활용하여 추가적인 작은 이점을 가져다주며(예: ImageNet-1K에서 VRWKV-L이 86.0에서 86.2로 증가).