Vision-RWKV: Efficient and Scalable Visual Perception with RWKV-Like Architectures

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.

한 줄 요약

Vision-RWKV(VRWKV)는 고해상도 이미지 처리에서 ViT 대비 10배 빠르고 80% 적은 메모리를 사용하는 시각 인식 모델이다.

핵심 기여도

VRWKV는 RWKV 기반의 **선형 복잡도 주의 메커니즘**을 도입하여 고해상도 이미지 처리 시 ViT 대비 10배 빠른 처리 속도와 80% 낮은 메모리 소비를 달성함.
**Q-Shift**와 **Bi-WKV** 모듈을 통해 ViT의 윈도우 연산 없이 글로벌 어텐션을 실현함.
**ImageNet-1K**에서 VRWKV-T는 75.1% 정확도를 달성해 DeiT-T 대비 2.9% 개선됨.
**ImageNet-22K**에서 VRWKV-L은 86.0% 정확도를 달성함.

핵심 아이디어

Vision-RWKV는 NLP 분야의 RWKV 모델을 시각 분야로 이식하여, **선형 복잡도 주의 메커니즘**을 도입함으로써 ViT의 제곱 복잡도 문제를 해결하려는 시도이다. ViT는 고해상도 이미지 처리 시 메모리와 계산 비용이 급증하지만, VRWKV는 **Bi-WKV**라는 양방향 글로벌 어텐션 메커니즘을 통해 이 문제를 완화한다. 또한, **Q-Shift**라는 새로운 토큰 쉬프트 메커니즘을 통해 토큰 간의 상호작용 범위를 확장하여, **Effective Receptive Field(ERF)**를 향상시킴. 이는 고해상도 이미지에서의 글로벌 정보 처리 능력을 높이는 데 기여한다.

기술적 접근법

**Bi-WKV**: 원래 RWKV의 인과적 어텐션(Causal Attention)을 **양방향 글로벌 어텐션**으로 수정하여, 고해상도 이미지에서도 효율적으로 글로벌 정보를 처리함.
**Q-Shift**: 토큰의 시맨틱 범위를 확장하는 새로운 토큰 쉬프트 메커니즘으로, ERF를 향상시킴.
**Relative Positional Bias**: 절대 위치 편향을 상대적 편향으로 변경하여, 모델의 확장성과 안정성을 향상.
**Layer Scale & Extra Layer Norm**: 모델 확장 시 발생하는 **overflow 문제**를 해결하기 위해 추가적인 정규화와 스케일링 기법 적용.
**모델 스케일**: VRWKV-T(6M), VRWKV-L(335M) 등 다양한 규모의 모델을 구축하고, **ImageNet-1K**, **ImageNet-22K**에서 훈련.

주요 결과

**ImageNet-1K**에서 VRWKV-T는 75.1% 정확도를 달성해 DeiT-T 대비 2.9% 개선됨.
**ImageNet-22K**에서 VRWKV-L은 86.0% 정확도를 달성함.
**2048×2048 해상도**에서 VRWKV-T는 ViT-T 대비 **10배 빠른 처리 속도**와 **80% 낮은 메모리 소비**를 기록함.
**Bi-WKV**는 2048×2048 해상도에서 PyTorch 어텐션 대비 **100배 이상 빠른 처리 속도**를 보임.

의의 및 한계

VRWKV는 고해상도 이미지 처리에서 ViT의 계산 복잡도 문제를 해결하며, **선형 복잡도 어텐션**의 실용성을 입증함. 특히, **Q-Shift**와 **Bi-WKV**를 통해 윈도우 연산 없이 글로벌 어텐션을 유지하는 점에서 혁신적이다. 또한, **MAE**를 활용한 사전 학습이 가능해, 비지도 학습 환경에서도 유용함. 그러나 **Q-Shift의 PyTorch 구현**은 효율성이 낮아, **CUDA 확장** 등 엔지니어링 최적화가 필요하다는 한계가 있음. 또한, 현재까지는 **이미지 분류**와 **밀집 예측**에서 평가되었으며, 다른 시각 작업(예: 3D 비전, 비디오 분석)에 대한 평가가 부재함.

실용적 활용

VRWKV는 고해상도 이미지 처리가 필요한 **의료 영상 분석**, **위성 이미지 해석**, **로봇 비전**, **실시간 영상 처리** 등에서 유용하게 활용될 수 있다. 특히, **GPU 메모리가 제한된 장치**나 **대규모 데이터셋**을 다루는 연구 환경에서 효율적인 대안으로 사용 가능하다.