PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Yuchen Duan, Weiyun Wang, Zhe Chen, Xizhou Zhu, Lewei Lu, Tong Lu, Yu Qiao, Hongsheng Li, Jifeng Dai, Wenhai Wang

arXiv:2403.02308v3cs.CV

TL;DR

High computational complexity limits transformer-based vision models on high-resolution images and long sequences. The paper introduces VRWKV, adapting RWKV with vision-specific token shifting and bidirectional linear-complexity attention. VRWKV delivers comparable or better visual-task performance with lower computational and memory costs, especially as resolution increases.

  • Problem

    ViTs’ quadratic computational complexity limits efficient processing of high-resolution images and lengthy sequences, while vision adaptations of linear-attention models must address spatial aggregation and scaling stability.

  • Method

    VRWKV adapts RWKV for vision using Q-Shift, bidirectional Bi-WKV global attention, relative positional bias, and stability strategies for scalable training.

  • Results

    VRWKV matches or exceeds ViT across classification and dense prediction while using lower computational and memory costs and becoming faster at higher resolutions.

  • Takeaways & Limitations

    VRWKV is presented as an efficient, low-cost alternative backbone to ViT for classification, dense prediction, and masked image modeling.

Abstract

from arXiv · show

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.

1 INTRODUCTION

ViTs provide global visual processing but incur quadratic costs that hinder high-resolution images and long sequences. VRWKV adapts RWKV with vision-specific mechanisms and scalable training strategies, targeting efficient global processing across visual tasks.

  • 1 INTRODUCTION: ViTs’ quadratic computational complexity limits efficient processing of high-resolution images and lengthy sequences.This motivates architectures that retain global processing while reducing computational demands.
  • 1 INTRODUCTION: VRWKV combines Q-Shift with bidirectional Bi-WKV to model local visual concepts and global information with linear complexity.Q-Shift expands token receptive fields, while Bi-WKV computes global attention in an RNN-like forward and backward process.
  • 1 INTRODUCTION: VRWKV retains efficient handling of global information and sparse inputs while supporting local visual modeling and scalable, stable training.The model incorporates relative positional bias, layer scale, and extra normalization among its stability strategies.
  • 1 INTRODUCTION: VRWKV models span 6M to 335M parameters and are trained on ImageNet-1K or ImageNet-22K for classification, detection, and segmentation.The models use supervised classification and masked image modeling training settings.
  • 1 INTRODUCTION: 75.1% top-1 accuracy is achieved by VRWKV-T, outperforming DeiT-T by 2.9 points; VRWKV-L reaches 86.0% on ImageNet-22K and 50.6% box mAP on COCO.The reported comparisons are 75.1% versus DeiT-T and 50.6 versus 48.7 box mAP against ViT-L.

2 RELATED WORKS

Vision research has progressed from convolutional and recurrent operators toward attention-based architectures, but global attention remains computationally expensive. Linear-complexity operators address efficiency, while visual applications require adapting their sequence-oriented mechanisms to image structure.

  • 2 RELATED WORKS: ViT introduced global receptive fields and dynamic spatial aggregation, while vanilla attention incurs quadratic computational complexity.Subsequent methods reduce cost through down-sampled feature maps or localized attention.
  • 2 RELATED WORKS: RWKV, RetNet, Mamba, and related operators provide linear-complexity sequence processing, but relatively few works have focused on visual applications.Vision-focused efforts include Vim and VMamba.
  • 2 RELATED WORKS: CNNs efficiently process local features but face challenges modeling long-range dependencies, motivating feature-aggregation alternatives.Convolutional operators use parameter sharing and sliding computation for large-scale visual data.
  • 2 RELATED WORKS: RNNs capture temporal dependencies efficiently, whereas transformers improve parallel computation and long-range modeling at higher computational cost.This trade-off motivates efficient alternatives that combine broad dependency modeling with lower complexity.
  • 2 RELATED WORKS: Window attention restricts computation to local windows, reducing complexity while preserving receptive fields through interactions among windows.Spatial reduction attention instead reduces feature dimensionality before attention computation.

3 VISION-RWKV

Vision-RWKV adapts RWKV for vision with bidirectional linear-complexity global attention, Q-Shift token mixing, and stabilization mechanisms for scalable training. Its design preserves global processing while maintaining linear complexity in the token count.

  • Overall architecture: VRWKV combines a ViT-like encoder with spatial-mix global attention and channel-mix feature fusion.The encoder uses patch embedding followed by stacked VRWKV layers, each maintaining input resolution.
  • Quad-directional token shift: Q-Shift exchanges information with neighboring tokens before computing R, K, and V, expanding token receptive fields without many additional FLOPs.For sparse inputs such as masked image modeling, the shift can operate in one dimension to preserve image priors.
  • Bidirectional attention: Bi-WKV makes all tokens mutually visible by replacing causal RWKV attention with bidirectional global attention.Relative spatial bias is incorporated through the token-distance term in the attention formulation.
  • Linear complexity: 13 × T × C FLOPs are required for Bi-WKV, giving forward and backward computational complexity O(TC).The practical implementation expresses the attention calculation in an RNN form with T update steps.
  • Stable scaling: VRWKV stabilizes scaling with bounded exponential terms, residual connections, normalization, relative positional bias, and extra normalization layers.These mechanisms address overflow and gradient stability as depth, resolution, and model size increase.

4 EXPERIMENTS

VRWKV is evaluated across classification, dense prediction, ablations, receptive fields, efficiency, and masked pretraining. Across these experiments, it generally matches or exceeds ViT performance while reducing computational or memory costs, especially at high resolution.

  • Experimental scope: VRWKV is evaluated for classification, object detection, semantic segmentation, and masked image modeling, including comparisons with ViT backbones.The experiments assess substitutability across multiple visual perception tasks.
  • Image classification: 2.9 points higher top-1 accuracy than DeiT-T is achieved by VRWKV-T with slightly lower FLOPs, while VRWKV-L reaches 86.0% top-1 accuracy versus ViT-L’s 85.15%.The classification comparison uses matched model scales, computational complexity, and resolutions.
  • Object detection: 30% lower backbone FLOPs and a 0.6-point APb improvement are reported for VRWKV-T versus window-attention ViT-T on COCO detection.VRWKV-S also has similar performance to ViT-S with 45% lower FLOPs under global attention.
  • Semantic segmentation: VRWKV-S achieves 1 point higher performance than ViT-S with 14% fewer FLOPs, while VRWKV-L reaches 53.5 mIoU with 25G fewer backbone FLOPs than ViT-L.These comparisons use UperNet on ADE20K with global attention for ViT models.
  • Ablation study: Q-Shift improves receptive-field coverage, while bidirectional attention increases top-1 accuracy by 2.3 points compared with the causal attention setting.Without token shift, performance is 3.6 points below the proposed model; the original shift remains 0.7 points behind it.
  • Efficiency analysis: At 2048×2048 resolution, VRWKV-T is 10 times faster and uses 80% less memory than ViT-T, although PyTorch’s Q-Shift implementation reduces overall speed.The comparison corresponds to 16,384 tokens on an Nvidia A100 GPU.

5 CONCLUSION

The paper presents VRWKV as a linear-complexity vision encoder and evaluates it as an alternative ViT backbone. It reports comparable performance and scalability with lower computational complexity and memory consumption, particularly for tasks where global attention is costly.

  • Conclusion: VRWKV is proposed as a vision encoder with a linear-complexity attention mechanism.The model is positioned as an alternative backbone to ViT across classification, dense prediction, and masked image modeling.
  • Conclusion: VRWKV combines comparable performance and scalability with lower computational complexity and memory consumption than ViT.The conclusion emphasizes its potential as an efficient, low-cost alternative for visual perception.

A.1 RNN FORM FORWARD AND BACKWARD

Bi-WKV expresses both forward and backward computation in RNN form, keeping updates linear in the token count while supporting bidirectional global attention. The section also reports resolution robustness and comparison coverage for VRWKV.

  • A.1 RNN FORM FORWARD AND BACKWARD: The spatial-mix attention mechanism uses RNN-form forward and backward computation to achieve linear complexity in the token number T.
  • A.1 RNN FORM FORWARD AND BACKWARD: Table 6 lists the Bi-WKV RNN states, initial values, and recurrence relations, with fixed FLOPs per update.
  • A.1.1 BACKWARD EQUATION: Backward propagation receives gy and computes gradients for w, u, K, and V using an RNN form with linear complexity in T.
  • A.1.1 BACKWARD EQUATION: The displayed backward equations include the output-gradient relation for gvt.
  • A.1.1 BACKWARD EQUATION: All recurrence updates have O(C) cost, giving final backward complexity O(sTC), where s sums the FLOPs of the equations.
  • A.1.1 BACKWARD EQUATION: Figure 4 evaluates VRWKV and DeiT trained at 224 × 224 across different evaluation resolutions to assess robustness.

A.1.2 IMPLEMENTATION DETAILS

The implementation stabilizes exponential recurrence updates by subtracting a maximum value, keeping exponential terms bounded during forward and backward computation.

  • A.1.2 IMPLEMENTATION DETAILS: A numerical trick computes safe exponentials in the forward and backward recurrences to avoid overflow.
  • A.1.2 IMPLEMENTATION DETAILS: The state a′ update uses exponentials involving w, p, q, k_t, and v_t.
  • A.1.2 IMPLEMENTATION DETAILS: Subtracting the maximum value forces new-state exponential terms below 1, with the subtracted quantity stored in p and divided during wkv computation.

A.2 ROBUSTNESS ON IMAGE RESOLUTION

The resolution-robustness experiment trains models at 224 × 224 and evaluates them from 224 × 224 through 1024 × 1024. VRWKV retains stronger accuracy as test resolution increases.

  • A.2 ROBUSTNESS ON IMAGE RESOLUTION: Models are trained at 224 × 224 and evaluated across resolutions from 224 × 224 to 1024 × 1024 on ImageNet-1K.
  • A.2 ROBUSTNESS ON IMAGE RESOLUTION: 82.5% top-1 accuracy at 384 × 384 gives VRWKV-B a 0.5-point improvement over its training-resolution accuracy.
  • A.2 ROBUSTNESS ON IMAGE RESOLUTION: At 1024×1024, VRWKV-B maintains 67.2% accuracy, whereas DeiT-B achieves 57.5%.
  • A.2 ROBUSTNESS ON IMAGE RESOLUTION: VRWKV performs better as resolution slightly increases, while DeiT does not show the same behavior.

A.3 COMPARISON TO VISION MAMBA

The comparison with Vision Mamba covers ImageNet-1K classification and inference efficiency across attention mechanisms. VRWKV scales to larger models, while Bi-WKV becomes faster than competing mechanisms as resolution grows.

  • Classification Performance: On ImageNet-1K, Vision Mamba has higher Top-1 Acc in tiny and small sizes, while base-size models perform comparably.
  • Classification Performance: VRWKV scales to larger models, whereas Vision Mamba encounters training instability during scaling.
  • Inference Efficiency: Figure 5 scans inference time from 224 to 1024 input resolution for Vanilla Attn, Bi-WKV, and Vision Mamba on Nvidia A100.
  • Inference Efficiency: As resolution increases, vanilla attention’s inference cost quickly surpasses Bi-WKV and Vision Mamba.
  • Inference Efficiency: With CUDA optimizations, Bi-WKV runs faster than Vision Mamba at the same input resolution.
Loading 2403.02308v3…