Vision-RWKV: Efficient and Scalable Visual Perception with RWKV-Like Architectures
Authors: Yuchen Duan
Weiyun Wang
Zhe Chen
Xizhou Zhu
Lewei Lu
Tong Lu
Yu Qiao
Hongsheng Li
Jifeng Dai
Wenhai Wang
Venue: ICLR 2025 Spotlight (2025)
Keywords: RWKV
Visual Perception
Linear Attention
Primary area: applications to computer vision, audio, language, and other modalities

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 that builds upon the RWKV architecture from the NLP field with key modifications tailored specifically for vision tasks. Similar to the Vision Transformer (ViT), our model demonstrates robust global processing capabilities, efficiently handles sparse inputs like masked images, and can scale up to accommodate both large-scale parameters and extensive datasets. Its distinctive advantage is its reduced spatial aggregation complexity, enabling seamless processing of high-resolution images without the need for window 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 and models are available at~\url{https://github.com/OpenGVLab/Vision-RWKV}.

Published as a conference paper at ICLR 2025
VISION-RWKV: EFFICIENT AND SCALABLE VISUAL PERCEPTION WITH RWKV-LIKE ARCHITECTURES
Yuchen Duan! 1.2 Weiyun Wang 22. Zhe Chen Xizhou Zhu 9.2.6 Lewei Lu Tong Lu', Yu Qiao², Hongsheng Li,,, Jifeng Dais,2, Wenhai Wang 1,253 The Chinese University of Hong Kong, Shhangaai AI Laboratory, FFudan University, NNajjing University, Tsinghua University, "SenseTime Research
ABSTRACT
Transformers have revolutionized computer vision and natural language process- ing, but their high computational complexity limits their application in high- resolution image processing and long-context analysis. This paper introduces Vision-RWi -WoVV (VRWKV), model that builds upon the RWKV architecture from the NLP field with key modifications tailored specifically for vision tasks. Similar to the Vision Transformer (ViT), our model demonstrates robust global processing capabilities, efficiently handles sparse inputs like masked images, and can scale up accommodate both large-scale parameters and extensive datasets. Its distinctive advantage is its reduced spatial aggregation complexity, enabling seamless pro- cessing of high-resolution images without the need for window operations. Our evaluations demonstrate that VRWKV surpasses ViT's performance in image clas- sification 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 po- tential as more efficient alternative for visual perception tasks. Code and models are available at https: : / github..... com/ htt////eeeVVaaaa /i//iiiiionn RWKV.
INTRODUCTION
Vision Transformers (ViTs) (Dosovitskiy et al., 2020; Touvron et al., 2021a; Vaswani et al., 2017; Steiner et al., 2021; He et al., 2021), renowned for their flexibility and global information process- ing capabilities, have established new benchmarks a variety of vision tasks in the past few years. However, the quadratic computational complexity associated with VITs limits their ability to effi- ciently process high-resolution images and lengthy sequences, posing a significant barrier to their broader application. As result, the exploration of visual architecture that integrates the versa- tility and comprehensive processing strengths of ViTs, while reducing computational demands, has emerged crucial area of research.
In recent developments within natural language processing (NLP), models with linear feature aggre- gation (or called "linear attention"") mechanisms like RWKV (Peng et al., 2023) and Mamba (Gu & Dao, 2023) have emerged as popular solutions for achieving heightened efficiency and processing lengthy texts. These innovative models have demonstrated attributes similar to transformers (Devlin et al., 2018; Raffel et al., 2019; Smith et al., 2022b; Liu et al., 2019; Radford et al., 2018; 8; 2019; Brown al., 2020; Lewis et al., 2019) in NLP tasks, including the ability to handle long-range dependencies and parallel processing. Furthermore, they have also proven to be scalable, perform- ing well with large-scale NLP datasets. Expanding these techniques into the visual domain shows promise in addressing the computational cost challenge encountered by ViTs.
To develop a vision model incorporating a linear attention mechanism based on the aforementioned methods, while ensuring high capacity for large-scale image data and diverse visual tasks, several critical issues need to be addressed. Firstly, the design of spatial feature aggregation operations needs to be reconsidered, taking into account the differences between image and text modalities. For instance, redesign of kernels and rewriting at the CUDA level are necessary for attention
"qqual contribution. CCorrsppondnng authors: wangwenhai @ pjlab.orggen

---

Published as a conference paper at ICLR 2025
48
2.5
12
100%
-86
47
asel
2
10
GPU Menuuy/ GB
9
Percentage Jower VRWKV-T VII-I
90%
Paroelagee Lower
45
Small
600%
64
03
Times Faster
40%
42
-VWWKV
VRWKV-T T VIT-T
41
Ting
WmT
0.5
2
3
40
0
0%
200
400
600
800
1000
224 512 640 768
1004 1280 1648 Resolution (b)
224
768 1004 1380 1000 2048
GFLOPs (a)
Resolution (c)
Figure/ 1: Performance and efficiency comparison of Vision-RWKV (VRWKV) and VIT. (a)
Bounding box average precision (AP") comparison of VRWKV and ViT (Touvron et al., 2021a) with window attention and global attention on the COCO (Lin et al., 2014) dataset. (b) Inference speed comparison of VRWKV-T and ViT-T across input resolutions ranging from 224 to 2048. (c)
GPU memory comparison of VRWKV-T and ViT-T across input resolutions from 224 to 2048.
mechanisms with causal receptive field in models like RWKV. Furthermore, the issues of gradient vanishing or exploding tend to arise gradually as the model scales up, resulting in unstable training with large parameter sizes and large-scale datasets. For example, Vision Mamba (Zhu et al., 2024) only gave appropriate results on models with less than 30M parameters. It is important to conduct an in-depth study of how linear attention models can be applied vision tasks effectively, including examining the scalability concerning data and parameters, assessing the efficiency in handling sparse visual data, and implementing necessary techniques to ensure model stability during scaling up.
Based on these points, we introduce RWVV (VRWKV). Our approach preserves the core structure and benefits of RWKV (Peng et al., 2023) while incorporating essential changes to process visual data efficiently. Specifically, we design quad-directional shift (Q-Shift) tailed for vision tasks and modify the original causal RWKV attention mechanism a bidirectional global atten- tion mechanism (Bi-WKV). The Q-Shift operation expands the semantic range of individual tokens, while the Bi-WKV enables the calculation of global attention within linear complexity in an RNN- form forward and backward. We primarily modify the exponent in the RWKV attention, releasing the limitations of the decay vector and transforming the absolute positional bias into relative bias. These changes enhance the model's ' capability while ensuring scalability and stabilliy. In this way, our VRWKV inherits the efficiency of RWKV in handling global information and sparse inputs, while also being able to model the local concept of vision tasks. Additionally, due to severe insta- bility encountered when scaling up the model, we explored a series of measures (Touvron et al., 2021b; Ba et al., 2016) to stabilize the model's outputs. These adjustments significantly improve the model's training stability when scaling up to a larger size.
Building the aforementioned design, we develop a range of VRWKV models with different model scales, spanning from the VRiKV-Tiny (6M) to the VRWKV-Large (335M). These mod- els are trained using large-scale datasets such as ImageNet-IK (Deng et al., 2009) and ImageNet- 22K (Deng et al., 2009). We train them using both common supervised classification and sparse input method MAE (He al., 2021) and evaluate their performance on visual perception tasks, including classification, detection, and segmentation. Under the same settings, VRWKV has com- parable performance to ViT in these tasks with lower computational costs while maintaining stable scalability. This achievement enables VRWKV training parallelism, high flexibility, excellent per- formance, and low inference cost simultaneously, making it a promising alternative to ViT in a wide range of vision tasks, particularly in high-resolution scenarios.
In this paper, our main contributions are:
(1) We propose VRWKV cost-effective alternative to ViT, offering comprehensive substitute with lower computational requirements. Our model retains ViT's strengths, such as capturing long- range dependencies and handling sparse inputs flexibly, while reducing complexity to linear scale. This reduction eliminates the need for window operation when processing high-resolution images, making VRWKV more efficient and scalable solution for vision tasks.
(2) To support vision tasks, we develop a bidirectional global attention mechanism combined with novel token shift method, Q-Shift, to achieve linear complexity in global attention. Additionally, we implement a set of tailored inntgiis-nnttgttiinngg relative positional bias, layer scale, and extra layer normalization--t to tackle overflow issues and ensure stable, scalable training.
2

---

Published as a conference paper at ICLR 2025
(3) Our model surpasses window-based ViTs and is comparable to global attention ViTs, demon- strating lower FLOPs and GPU memory cost with faster processing speeds as resolution increases, as shown in Figure 1. Notably, VRWKV-T achieves 75.1% top-1 accuracy trained only the ImageNet-1K (Deng et al., 2009), outperforming DeïT-T (Touvron et al., 2021a) by 2.9 points. With large-scale parameters (i.., 335M) and training data (ie.. ImageNet-22K), the top-1 accuracy of VRWKV-L is further boosted to 86.0%, which is higher than ViT-L (Dosovitskiy et al., 2020) (86.04 vs 85.15). In addition, on COCO (Lin et al., 2014), a challenging downstream benchmark, our best model VRWKV-L achieves 50.6% box mAP, 1.9 points better than ViT-L (50.6 vs 48.7).
2
RELATED WORKS
2.1 VISION ENCODER
Recent advances in vision encoders have significantly pushed the boundaries of computer vision, demonstrating remarkable performance across range of tasks. Convolutional neural networks (CNNs) served as the foundational model in computer vision. The advancement of computational resources, such as GPUs, has enabled the successful training of stacked convolutional blocks like AlexNet (Krizhevsky et al., 2012) and VGG (Simonyan & Zisserman, 2014) on large-scale image classification datasets (e.g., ImageNet (Deng et al., 2009)). This development paved the way for deeper and more sophisticated convolutional neural architectures, including GoogleNet (Szegedy et al., 2015), ResNet (He et al., 2016), and DenseNet (Huang et al., 2017).
In addition to these innovations, significant advancements have also been made with architectures like SENet (Hu et al., 2018), which introduced channel attention mechanism to enhance model sensitivity to informative features. Similarly, SKNet (Li et al., 2019) merged multiple kernel sizes to adjust the receptive field adaptively. Further extending the CNN paradigm, recent models such as RepL.NNet (Ding et al., 2022) and ConvNeXt (Liu et al., 2022) have refined the convolutional layers improve efficiency and accuracy, while InternImage (Wang al., 2023b) explored the strategies to scale up the convolution-base vision model.
Inspired by the success of self-attention layers and transformer architectures in the NLP field, the Vision Transformer (ViT) (Dosovitskiy et al., 2020) applied transformer framework on image patches, offering global receptive field and dynamic spatial aggregation. Due to the quadratically computational complexity of the vanilla attention mechanism, approaches like PVT (Wang et al., 2021; 2022) and Linformer (Wang et al., 2020) implemented global attention on down-sampled fea- ture maps, whereas other approaches like Swin (Wu et al., 2022) and HaloNet (Vaswani al., 2021; Dai et al., 2022) introduced sampling techniques to enlarge the receptive field. Mini-Intern VL (Gao et al., 2024) reduces the parameter size of VIT by employing knowledge distillation from larger ViT to a smaller one, thereby achieving efficiency in the visual encoder.
Another research direction involved replacing self-attention layers in models with linear complexity layers. Representative works include LongNet (Ding et al., 2023), RWKV (Peng et al., 2023), RetNet (Sun et al., 2023), and Mamba (Gu & Dao, 2023), though few have concentrated on visual applications. Concurrently, attempts like Vim (Zhu et al., 2024) and VMmmba (Liu et al., 2024) have sought to integrate these linear attention layers into the vision domain. However, these endeavors have only been experimented with on small-scale models (parameters < 30M for Vim and < 100M for VMamba), leaving it uncertain whether their effectiveness extends larger models.
2.2 FEATURE AGGREGATION MECHANISM
The research on feature aggregation has received significant attention. For visual data processing, convolutional operators (LeCun et al., 1995), known for their parameter sharing and local perception, enabled efficient handling of large-scale data through sliding computation. Despite their advantages, traditional CNN operators faced challenges in modeling long-range dependencies. To overcome this issue, advanced convolutional operators, such as the deformable convolution (Dai et al., 2017; Zhu et al., 2019; Xiong et al., 2024), have improved the flexibility of CNN operators, enhancing their long-range modeling capability.
As for the field of NLP, RNN-based operators (Elman, 1990; Memory, 2010; Qin et al., 2024) have historically dominated because of their effectiveness in sequence modeling. RNNs and LSTMs excel in capturing temporal dependencies, making them suitable for tasks requiring an understanding of sequence dynamics. Subsequently, a significant shift occurred. The introduction of the transformer
3

---

Published as a conference paper at ICLR 2025
Linear Head
Pred
Spatial Mix
Channel Mix
Avg Pooling
R
R
O Shift
O2
0
ASEAN B
Channel Mix
States
XL
Spatial Mix
Input Embeddings
Multagicctioon
Element- wise
Elezentt-wsse
Sigmald
Squard Hel.U
(a) Vision-RWKV Architacture
(b) Vision-RWKV Encoder Layer
Figure 2: Overall architecture of VRWKV. (a) The VRWKV architecture includes L identical VRWKV encoder layers, an average pooling layer, and linear prediction head. (b) The details of the VRWKV encoder layer. Q-Shift denotes the quad-directional shift method tailed for vision tasks. The "Bi-WKV" module served as a bidirectional RNN cell or global attention mechanism.
architecture (Vaswani et al., 2017) marked a turning point, with both NLP and computer vision fields shifting focus toward iitention-based feature aggregation. The global attention mechanism overcomes the limitations of CNNs in modeling long-range dependencies and the shortcomings of RNNs parallel computation while coming a high computational cost.
To address these issues, researchers have introduced innovations such as window attention and spa- tial reduction attention. Window attention (Liu et al., 2021; Vaswani et al., 2021; Dai et al., 2022) restricts the self- computation within local windows, drastically reducing the computational complexity while preserving the receptive field through window-level interaction. Spatial reduction attention (Wang et al., 2021; 2022), the other hand, reduces the dimensionality of the feature space before applying the attention mechanism, effectively decreasing the computational require- ments without significantly degrading the model's performance.
In addition to the efforts to optimize the global attention mechanism, various operators with linear complexity have also been explored. For instance, RWKV (Peng et al., 2023) and RetNet (Sun et al., 2023) employed exponential decay to model global information efficiently while SSMs (Gu al., 202 la;;;;;; Smith et al., 2022a; Wang al., 2023a) also exhibited linear complexity concerning sequence length and modification in Mamba (Gu & Dao, 2023) enable them to be input-dependent. Besides. XCA (Ali et al., 2021) achieved linear complexity by calculating the cross-variance be- tween input tokens. However. the low efficiency of information interaction between tokens makes the need for the assistance of additional modules necessary to complete a comprehensive feature aggregation. Despite some concurrent efforts (Liu et al., 2024; Zhu et al., 2024; Fan et al., 2023), adapting these der-vev techniques to vision tasks remains challenge in maintaining stable training across larger and more complex vision models.
VISION-RWKV
3.1 OVERALL ARCHITECTURE
In this section, we propose Viskon-RWKV (VRWKV), an efficient vision encoder with a linear complexity attention mechanism. Our principle is preserve the advantages of the original RWKV architecture (Peng et al., 2023), making necessary modifications to enable its flexible application in vision tasks, supporting sparse input, and ensuring the stability of the training process after scaling up. An overview of our VRWKV is depicted in Figure 2.
VRWKV adopts a block-stacked image encoder design like ViT, where each block consists a spatial-mix module and a channel-mix module. The spatial-mix module functions mechanism, performing linear complexity global attention computation while the channel mix mod- ule serves a feed-forward network (FFN), performing feature fusion in the channel dimension. The entire VRWKV includes a patch embedding layer and a stack of identical VRWKV encoder layers, where each layer maintains the input resolution.
attention
4

---

Published as a conference paper at ICLR 2025
Data Flow. First, we transform the H x W x 3 image into HW/p² patches, where p denotes the patch size. The patches after a linear projection add the position embedding to obtain image tokens of shape T x C, where T HW/p² denotes the total number of tokens. These tokens are then input into the VRWKV encoder with I layers.
In each layer, tokens are first fed into the spatial-mix module which plays the role of global attention mechanism. Specifically, as shown in Figure 2(b), the input tokens are first shifted and fed into three parallel linear layers to obtain the matrices Re, Ks, V. -2'x;
R = (X)
K; =
Q-Shift x (X)WK, V, = Q
(I)
Here, the Q-Shift operator a token shift function specially designed for the information exchange through nearby tokens according to the visual prior. K, and V, are then passed to calculate the global attention result, wkv e RT'xC, by a linear complexity bidirectional attention mechanism, Bi-WKV, and multiplied with a (R) which controls the output O probability:
O, = (a R2) @ wko
where wky = Bi- (K,, V.).
Operator a denotes the sigmoid function, and @ represents element-wise multiplication. The output features are then stabilized using layer normalization (Ba al., 2016) following the linear projection.
Subsequently, the tokens are passed into the channel-mix module for a channel-wise fusion. Re Ke obtained in similar manner as spatial-mix:
R. Q-Shift (XX)WR, K. = Q-Sliit (X)Wk.
(3)
In the channel mix module, V. is linear projection of Ke after the activation function and controlled by gate mechanism o R.). The output O is the linear projection of the aforementioned result:
O = (a< @ Vo)Wo,
(4)
where (
Simultaneously, residual connections (He al., 2016) established from the tokens to each nor- malization layer to ensure that training gradients do not vanish in deep networks.
3.2 LINEAR COMPLEXITY BIDIRECTIONAL ATTENTION
Different from the vanilla RWKV (Peng et al., 2023), we make the following modifications to its original attention mechanism: (1) Bidirectional attention: We extend the upper limit of original RWKV attention from (the current token) to T (the last token) to ensure that all tokens are mutually visible in the calculation of each result. Thus, the original causal attention transforms into bidirectional global attention.
(2) Relative bias: We compute the absolute value of the time difference t and divide it by the total number of tokens (denoted as T) to represent the relative bias of tokens in images of different sizes. (3) Flexible decay: We no longer restrict the learnable decay parameter w to be positive in the exponential term allowing the exponential decay attention to focus on tokens further away from the current token.
Under the collective influence of these ingenious modifications, we achieve global attention while maintaining linear complexity to the input token number T, thereby maximizing the preservation of RWKV's inherent low complexity and extending it to the visual domain.
Similar to RWKV, our bidirectional attention can also be equivalently expressed in summation form (for the sake of clarity) and an RNN form (in practical implementation).
Summation Form. The attention calculation result for the 1-th token is given by the formula:
Σ 1)/- Meg vs + ut
wkv; = Bi-WKV(K V),
(5)
ET- ([-41)-1) w guttke
Here, T represents the total number of tokens, equal HW/p², and two C-dimensional learnalle vectors that represent channel-wise spatial decay and the bonus indicating the current token, respectively. k, and u, denotes f-th feature of K and V.
The summation formula indicates that the output wkv, is a weighted sum of V along the token dimension from 0 to T 1, resulting in C-dimensional vector. It represents the result obtained
5

---

Published as a conference paper at ICLR 2025
Model
Emb Dim
Hidden Dim
Depth
Extra Norm
#Param
V-T
192
768
12
X
6.2M
V
384
1536
12
23.8M
VRWKV-B
768
3072
12
93.7M
VRWKV-L
1024
4096
24
334.9M
Default settings for Vision-RWKV of different scales. We report the embedding di- mension, hidden dimension, and model depth. "Extra Norm" means additional layer normalization layers are used stabilize the model's outputs. "#Param" denotes the number of parameters.
Table 1:
by applying the attention operation to the t-th token. The weight is determined by the spatial decay vector w, the relative bias between tokens (|t i| 1)/T. and kę collectively.
RNN Form. In the practical implementation, the above Eq 5 can be transformed into a recursive formula in the form of RNN that the result of each token can be obtained through fixed number of FLOPs. By splitting the summation term of the denominator and numerator in Eq 5 with as the boundary, we can obtain hidden states:
0t-1
[[t-d) 1/
b
Σ
([t-s]- (1/-s]
<m2p1
(6)
-11--1-1)/T-w4k,
d-11
(1I--)
1=t+1
that can be recursively computed due to its mathematical formulation. The update of hidden states only requires adding or subtracting one summation term and multiplying or dividing e ww//T Then the 1-th result can be given
wkk Gr b 1 K
0-1 1 + di + gkotu
Each update step yields attention result (i.e., wkk,)) for token, the entire wke matrix requires steps.
When the input K and V are matrices with the shape of 7 x C, the computational cost of calculating the whe matrix is given by:
(8) Here, the number 13 is approximately from the updates of (a, b, e, d), the computation of the ex- ponential, and the calculation of wkv, The above approximation shows that the complexity of the forward process is (TTC). The backward propagation of the operator can still be represented more complex RNN form, with a computational complexity of O(TC). The specific formula for forward updating and backward propagation is provided in the Appendix A.1.
WKY V)) 13 C
3.3 QUAD-DIRECTIONA. TOKEN SHIFT
We introduce quad-directional token shift (Q-Shift) flexible extension of the original token shift operation in RWKV in the first step of each spatial-mix and channel-min module. The Q-Shift operation allows all tokens shifted and linearly interpolated with their neighboring tokens as follows:
Q-Shift (-) (x)() = X
where X x* Concat w. C/1] C/ 1, C/4: C/2], X(h, w 1, C,C/ 3C/4],
3C/C//440 CJ).
Subscript (+) (R,K, V denotes interpolation of X and X controlled by the learnable vectors ((a) for the later calculation of R, K, V. respectively. h and w denote the row and column index of token X, ".* is slicing operation excluded the end index. The Q-Shift makes the attention mech- anism of different channels obtain the prior of focusing on neighboring tokens internally without introducing many additional FLOPs. also increases the receptive field each token which greatly enhances the coverage of the token in the posterior layers. X is obtained by slicing X without introducing new computations, allowing for flexible transformations during the training process for different tasks. When handling sparse inputs that do not contain original image space information, such as masked image modeling, shifting can be applied a single dimension to maximize the preservation of image priors.

---

Published as a conference paper at ICLR 2025
Method
Size
#Param
FLOPs
Top-1 Acc
ResNet-18 (He ct al., 2016)
224
11.7M
1.8G
69.9
PVT-T (Wang al., 2021)
221
13.2M
1.9G
75.1
ResNet-50 (He al., 2016)
224
25.6M
4.1G
76,6
hierarchical
Swin-T (Liu al., 2021)
221
28.3M
4.4G
81.2
PVT-M (Wang et al., 2021)
224
44.2M
6.7G
81.2
ResNet- 101 (He et al., 2016)
224
44.6M
7.9G
78.0
Swin-S (Liu al., 2021)
224
49.6M
8.7G
83.0
PVT-L (Wang ct al., 2021)
224
61.4M
9.8G
81.7
Swin-B (Liu ct al., 2021)
224
87.8M
15.1G
83.4
DeïT-T (Touvron et al., 2021
221
5.7M
1.3G
72.2
DeīT-S (Touvron et al., 2021a)
224
22.1M
4.6G
79.9
XCTT-S12 (Ali al., 2021)
2212
26.0M
4.8G
82.0
DeiT-B (Touvron al., 2021a)
224
86.6M
17.6G
81.8
non hee-hherarchicall
XCTT-L24 (Ali al., 2021)
224
189.0M
36.1G
82.9
ViT-L (Dosovitskiy et al., 2020)
384²
309.5M
191.1G
85.2
VRWKV-T
224
6.2M
1.2G
75.1
VRWKV-S
224
23.8M
4.6G
80.1
VRWKV-B
221
93.7M
18.2G
82.0
VRWKV-L
384
334.9M
189.5G
86.0
VRWKV-L
384
334.9M
189,5G
86.2
VRWKV-I."
3842
334.9M
189.5G
86.5
Table 2: Validation results on ImageNet-IK. VRWKV-T/S/B are trained on ImageNet-IK, while VRWKV-L is pre-trained Imagenet-22K and fine-tuned on Imagenet-IK. "#Param" denotes the number of parameters, and "FLOPs" represents the computational workload for the specified image resolution the "Size" column. "f" means additional MAE pre-training is applied in the pre- training process. "*" indicates Bamboo-47K is used in the pre- pre-training.
3.4 SCALE UP STABILITY
Increasing model depth and the accumulation of exponential terms during recursion can lead to instability in the training process. To mitigate this, propose simple yet effective adjustments: (1) Bounded exponential: As input resolution increases, both exponential decay and growth can quickly exceed the range of floating-point numbers. To address this, divide the exponential by the number of tokens (e.g., exp(-(It il /T w)), making the maximum decay and growth bounded. (2) Extra layer normalization: As models become deeper, we apply layer normalization (Ba et al., 2016) after the attention mechanism and the Squared ReLU operation, to prevent the model's s output from overflowing. These two adjustments promote stable scaling of input resolution and model depth, facilitating the smooth convergence of large models. Additionally, we incorporate layer scale (Touvron et al., 2021b), which further enhances model stability during scaling.
3.5 MODEL DETAILS
Following VIT, the den dimension in linear projection, and depth, are specified in Table 1. Due to the increased depth of the VRWKV-L model, additional layer normalizations discussed in Section 3.4, are incorporated at appropriate positions to ensure output stability.
for variants of VRWKV, including embedding dimension, hid-
EXPERIMENTS
We comprchensively evaluate the substitutability of our VRWKV for VIT in performance, scalabil- ity, flexibility, and efficiency. The model's effectiveness is validated in image classification, object detection, and semantic segmentation tasks.
4.1 IMAGE CLASSIFICATION
Settings. For -Tin/SSmallBBae models, conduct supervised training from scratch on ImageNet- 1K (Deng et al., 2009). Following the training strategy and data augmentation DeiT (Touvron al., 20) use batch size of 1024, AdamW (Loshchilov & Hutter, 2017) with a base learning rate of 5c-4, weight decay of 0.05, and cosine annealing schedule (Loshchilov Hutter, 2016). Images cropped the resolution of 224 x 224 for training and validation. For the -Large models,

---

Published as a conference paper at ICLR 2025
Method
Window
#Paara
FLOPs
AP
APM
VIT-T (Touvron et al., 2021a)
8.0M
95.4G
41.1
37.5
VTT-T (Touvron et al., 2021a)
8.0M
147.IG
41.6
37.9
VRWKV-T (ours)
8.4M
67.9G
41.7
38.0
VIT-S (Touvron et al., 2021a)
V
27.5M
241.2G
44.6
39.7
VIT-S (Touvron et al., 2021a)
X
27.5M
344.5G
44.9
40.1
VRWKV-S (ours)
X
29.3M
189.9G
44.8
40.2
ViT-B (Touvron et al., 2021a)
V
99.5M
686.7G
46.2
41.5
ViT-B (Touvron et al., 2021a)
X
99.SM
893.3G
46.8
41.8
VRWKV-B (ours)
X
106.6M
599.0G
46.8
41.7
VIT-L (Steiner et al., 2021)
327 OM
1799.3G
48.7
43.3
VRWKV-I. (ours)
X
351.9M
1730.6G
50.6
44.9
Table 3: Object detection and instance segmentation on COCO val2017. All models adopt the ViT-Adapter (Chen et al., 2023) to generate multi-scale features for detection heads. -T/S/B models are initialized with ImageNet-Ik weights, and all -L models use ImageNet-222 weights. "#Param" indicates the backbone parameter, and 'LLOPs represent the backbone's computational workload for a 1333 x 800 input. "Window" denotes the use of window operation in ViT layers.
first pre-train them for 90 epochs ImageNet-22K with batch size of 4096 and resolution of 192 192, and then fine-tune them for 20 epochs on ImageNet- 1 to a higher resolution 384 384.
Results. We compare the results of our VRWKV with other hierarchical and non-hierarchical back- bones on the ImageNet- 1K validation dataset. As shown in Table 2, with the same number of parameters, computational complexity, and training/testing resolutions, VRWKV achieves better sults than ViT. For example, when VRWKV-T has slightly lower FLOPs than DeiT-T(1.22 vs 1.3G), our VRWKV-T achieves 2.9 point higher than DeiT-T top-1 accuracy. When the model size scales up. VRWKV still demonstrates higher baseline performance. VRWKV-L achieves 0.8 point higher top-1 accuracy of 86.0% the resolution of 384 x 384 than ViT-L, with a slightly reduced computational cost. The superior performance from tiny to large-size models demonstrates that the VRWKV model possesses the scalability as VIT. Additionally, after using a larger dataset Bamboo- 47K (Zhang et al., 2022) in the pre-train process, the performance of VRWKV-L. can be further boosted to 86.5%, indicating that our VRWKV also possesses the ability like VIT to benefit from pre-training on large-scale datasets. The exploration of VRWKV in classification tasks demonstrates its potential to be a viable alternative to traditional ViT models.
4.2 OBJECT DETECTION
Settings. the detection tasks, we adopt Mask R-CNN (He et al., 2017) as the detection head. For the -Tin//Small/Base models, the backbones weights pre-trained ImageNet-1K for 300 epochs. For the -Large model, weights pre-trained on ImageNet-22K are used. All models use a xx training schedule (i.e., 12 epochs) with batch size of 16, and Adam W (Loshchilov & Hutter, 2017) optimizer with an initial learning rate of le-4 and weight decay of 0.05.
Results. In Table 3, report the detection results the COCO val (Lin et al., 2014) dataset using VRWKV and VIT as backbones. As the results showed in Figure ((a) and Table 3, due to the use of window attention in dense prediction tasks, VRWKV with global attention can achieve better performance than ViT with lower FLOPs. For example. VRWKV-T has approximately 30% lower backbone FLOPSs compared ViT-T using window attention, with an improvement of APb by 0.6 points. Additionally, we compare the performance of VRWKV and VIT using global attention. For instance, VRWKV-S achieves similar performance to ViT-S with 45% lower FLOPs. This demon- strates the effectiveness of VRWKV's global attention mechanism in dense prediction tasks and the advantage of lower computational complexity compared to the original attention mechanism.
4.3 SEMANTIC SEGMENTATION
Settings. In the semantic segmentation task, we use UperNet (Xiao et al., 2018) as the segmen- tation head. Specifically, all ViT models global attention in the segmentation task. For the -Tin/Smaal/Baaee models, the backbones use weights pre-trained on ImageNet- 1.k And for the Large model, weights pre-trained on ImageNet-22K used. We employ the AdamW optimizer
8

---

Published as a conference paper at ICLR 2025
Method
Window
#Param
FLOPs
mloU
ViT-T (Touvron et al., 2021a)
8.0M
20.9G
42.6
VRWKV-T (ours)
X
8.4M
16.6G
43.3
ViT-S (Touvron et al., 2021a)
27.5M .M
54.0G
46.2
VRWKV-S (ours)
X
29.3M
46.3G
47.2
ViT-B (Touvron al., 202
99.5M
157.9G
48.8
VRWKV-B (ours)
106.6
146.0G
49.2
ViT-L (Steiner et al., 2021)
X
327.0M
446.8G
53.4
VRWKV-L (ours)
X
351.9M
421.9G
53.5
Table 4: Semantic segmentation on the ADE20K val set. All models used ViT-Adapter (Chen et al., 2023) for multi-scale feature generation and are trained with UperNet as the segmentation heads. For consistency in comparison, all -T/S/B models are initialized using ImageNet- IK pre- training. whereas -L models utilize ImageNet-22K, pre-training. "#Param" refers to the number of parameters of the backbone. We report the FLOPS of backbones with the input size of 512 x 512.
with an initial learning rate of 6e-5 for the Smaal/Base/Larrg models and 12c-5 for the Tiny model, a batch size of 16, and a weight decay of 0.01. All models are trained for 160k iterations on the train- ing set of the ADE20K dataset (Zhou et al., 2017).
Results. As shown in Table 4, when using UperNet for semantic segmentation, models based on VRWKV consistently outperform those based on ViT with global attention, while also being more efficient. For example. VRWKV-S achieves point higher than ViT-S with 14% FLOPs decrease. VRWKV-L creates result of 53.5 mloU, similar to ViT-L, while the computation of the backbone has a 25G FLOPs lower. These results demonstrate that our VRWKV backbones can extract better features for semantic segmentation compared to VIT backbones while also being more efficient, benefiting from the linear complexity attention mechanism.
4.4 ABLATION STUDY
Settings. We conduct ablation studies of the tiny-size VRWKV on ImageNet- (Deng al., 2009) to validate the effectiveness of various key components like Q-Shift and Bi-WKV. The experimental settings are consistent with Section 4.1.
Token Shift. We compare three approaches: no token shift, the original shift method used in RWKV, and our proposed Q-Shift. As shown in Table 5, the variation in the shift method shows performance differences. Variant I with- out token shift leads to poor performance of 71.5, which is 3.6 points lower than model. Even with the use of our global attention, the model using the original token shift still has 0.7-point gap with our model.
Bidirectional Attention
Method
Token Shift
Top- Acc
RWKV
original
X
71.1 (-4.0)
Variant
none
71.5 (-3.6)
Variant
original
74.4 (-0.7)
Variant
Q-Shift
72.8 (-2.3)
VRWKV-T
Q-Shift
75.
Table 5: Ablation on key components of the propased VRWKV, All models are trained from scratch on ImageNet-IK: "Original" refers the token shift in RWKV (Peng et al., 2023), which mixes tokens in a single direction.
Bidirectional Attention. The bidirectional attention mechanism enables the model to achieve global attention while the original RWKV attention has causal mask internally. The result of Variant 3 shows the global attention mechanism brings 2.3 points increase in the top-1 accuracy.
Effective Receptive Field (ERF). We analyze the impact of different designs on the ERF of models based on (Ding et al., 2022) and visualize it in Figure 3(a). We visualize the ERF of the central pixel with an input size of 1024 x 1024. In Figure 3(a), "No Shift" represents the absence of the token shift method (Q-Shift), and "RWKV Attn" indicates using the original RWKV attention mechanism without our modifications for vision tasks. From the comparison in the figure, all models except the "RWKV Attn" one achieved global attention while the global capacity of the VRWKV-T model is better than that of VIT-T. Despite the assistance of Q-Shift, the central pixel in "RWKV Attn" still cannot attend to the pixels on the bottom of the image due to the large input resolution. The results

---

Published as a conference paper at ICLR 2025
A Airrtion Runtime
Attention Runtime (Bowwrd-Huiswwel
(inference)
0.6
100
Log Rotttttmoo
1.0k: fasterl
230x faster 1
VIT-T
VIWEV-TINo StiTt
10
0.4
1200
1700
2200
700
1200
2200
0.0
Resolstion
Readiution
V VWWKCUTTEW KV AUD V VRWKY-T(OUM)
Pytorch ttn
Flash
B:-WKV (Ous)
ERF of VIT and VRWKV
A ttntion Runtime
Figure 3: Comparison of effective receptive field (ERF) and attention runtime. (a) ERF for ViT and VRWKV in different settings. "No Shift" means no shift is used in spatial-mix and channel- mix modules. "RWKV Attn" means the original RWKV attention without our modifications. Our VRWKV with Q-Shift and Bi-WKV has more comprehensive ERF than other counterparts. (b) Attention runtime of inference (left) and forward + backward (right) tested on an Nvidia A100 GPU.
of "No Shift" and Q-Shift show that the Q-Shift method expands the core range of the receptive field, enhancing the inductive bias of global attention.
Efficiency Analysis. To showcase the efficiency of our linear attention mechanism, we gradually increase the input resolution from 224x x 224 to 2048 x 2048 and compare the inference and memory efficiency of VRWKV-T and ViT-T. The results were tested on an Nvidia A100 GPU, shown in Figure 1. From the curves presented in Figure 1(b), it is observed that at lower resolutions, VRWKV-T and ViT-T exhibit comparable memory usage and inference FPS. With the increase in input resolution, VRWKV-T shows a much higher speed than VIT-T. Additionally, VRWKV-T's RNN-like computational framework ensures a slow increase in GPU memory usage. By the time the resolution hits 2048 x 2048 (cquivalent to 16384 tokens), VRWKV-T's inference speed is 10 times faster than ViT-T, and its memory consumption is reduced by 80% compared to ViT-T. It is worth mentioning that the implementation of the Q-Shift operation in Py PyTorc is highly inefficient, leading to an overall decrease in model speed. However, this operation can be optimized through other means (such using CUDA extensions). Therefore, there still potential for further improvement in the speed of VRWKV with better engineering optimization.
We also compare the speed of our attention mechanism kernel, Bi- Bi-WKV, with pytorch attention and flash attention (Dao et al., 2022), reported in Figure 3(b). Our Bi-WKV is significantly faster than the attention mechanism implemented using matrix multiplication (PyTorch attention), achieving speedup of over hundred times at a resolution of 2048 x 2048 (i.e., 16384 tokens). Flash attention highly optimized for memory I/O, and its matrix multiplication aligns well with the physical archi- tecture of Nvidia GPUs, while our Bi-WKV lacks such hardware-level optimization. Nevertheless, high-resolution scenarios, our Bi- WKV still demonstrates significant speed advantage.
MAE Pre-training. ViT' can learn meaningful visual representations in masked image modeling (MIM). Yet, the rigor and effectiveness of linear attention vision models in this self-supervieed pre- training paradigm have not been validated. Our VRWKV can handle sparse inputs and leverage MIM pre-training methods like MAE (He et al., 2017) by implementing a bidirectional shift operation that removes the vertical shift in Q-Shift. The pre- trained weights can be directly fine-tuned for other tasks using a Q-Shift manner. Following the same MAE pre-training setting as ViT, and subsequent classification training similar Section 1. our VRWKV-L shows the ability acquire visual prior from masked image modeling as shown in Table 2.
CONCLUSION
We propose Vision-RWKV (VRWKV), a vision encoder with a linear computational complexity attention mechanism. We demonstrate its capability to be an alternative backbone to ViT in com- prehensive vision tasks including classification, dense predictions, and masked image modeling pre- training. With comparable performance and scalability, VRWKV exhibits lower computational com- plexity and memory consumption. Benefiting from its low complexity. VRWKV can achieve better performance in the tasks that ViT struggling to afford the high computational overhead of global attention. We hope VRWKV will be an efficient and low-cost alternative to ViT, showcasing the powerful potential of linear complexity transformers in vision fields.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This project was supported by the National Key R&D Program of China (No. 20222D0161100) 00222D01601011 the National Natural Science Foundation of China (No. 62376134). Zhe Chen is supported by the Youth PhD Student Research Project under the National Natural Science Founda- tion (No. 623B2050).
REFERENCES
Alaaeldin Ali, Hugo Touvron, Mathilde Caron, Piotr Bojanowski, Matthijs Douze, Armand Joulin, Ivan Laptev, Natalia Neverova, Gabriel Synnaeve, Jakob Verbeek, et al. Xcit: Cross-covariance image transformers. NeurIPS, 34, 2021. 7.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 1607.06470 2016. 2, 5,
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. NeurIPS, 33:1877-1901, 2020. 1
Zhe Chen, Yuchen Duan, Wenhai Wang, Junjun He, Tong Lu, Jifeng Dai, and Yu Qiao. Vision transformer adapter for dense predictions. 2023. 8,9
Jifeng Dai, Haozhi Qi, Yuwen Xiong, Yi Li, Guodong Zhang, Han Hu, and Yichen Wei. Deformable convolutional networks, In ICCV, pp. 764-773, 2017. 3
Jifeng Dai, Min Shi, Weiyun Wang, Sitong Wu, Linjie Xing, Wenhai Wang, Xizhou Zhu, Lewei Lu, Jie Zhou, Xiaogang Wang, al. Demystify transformers & convolutions in modern image deep networks. arXiv preprint arXi::2211.0573
2022. 3,4
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. NeurIPS, 35:16344-16335 2022. 10
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, pp. 248-255, 2009. 2, 3, 7, 9, 16, 17
Jacob Devlin, Ming-Wes Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.048005 2018.
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to 1,0 1,000,000,000 tokens. arXiv preprint arXiv:2307.024880 2023.
Xiaohan Ding, Xiangyu Zhang, Jungong Han, and Guiguang Ding. Scaling up your kernels to 31x31: Revisiting large kernel design in cnns. In CVPR, pp. 11963-11975, 2022. 3,
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, al. An image is worth 16x 16 words: Transformers for image recognition at scale. In ICLR, 2020. 1, 3,
Jeffrey L Elman. Finding structure in time. Cognitive Science, 14(2):179-211 1990.
Qihang Fan, Huaibo Huang, Mingrui Chen, Hongmin Liu, and Ran He. Rmt: Retentive networks meet vision transformers. arXiv preprini arXiv:2309. 11523, 2023. 4
Zhangwei Gao, Zhe Chen, Erfei Cui, Yiming Ren, Weiyun Wang, Jinguo Zhu, Hao Tian, Shenglong Ye, Junjun He, Xizhou Zhu, et al. Mini-internvl: flexible-trance pocket multi-modal model with 5% parameters and 90% performance. Visual Intelligence, 2(1):1-17, 2024.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint /2200223/2.00752, 2023. 3,
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00399) 2021a. 4

---

Published as a conference paper at ICLR 2025
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-timg models with linear state space layers. NeurlPS, 34:572-585, 2021b. 4
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In CVPR, pp. 770-778, 2016. 5,
Kaiming He, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask r-cnn. In ICCV, pp. 2961- 2969, 2017. 8, 10
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dolkár, and Ross Girshick. Masked autoencoders are scalable vision learners. arXiv preprint arXiv:2111.06377, 2021. 1.2
Jie Hu, Li Shen, and Gang Sun. Squeere-and-exciiation networks. In CVPR, pp. 7132-7141, 2018.
Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinherger. Densely connected convolutional networks. In CVPR, pp. 4700-4708, 2017. 3
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. NeurIPS, 25, 2012.
Yann LeCun, Yoshua Bengio, et al. Convolutional networks for images, speech, and time series. The Handbook of Brain Theory and Neural Networks, 33)) 1995. 3
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov, and Luke Zettlemoyer. Bart: Denoising sequence-to-sequene pre- training for natural language generation, translation, and comprehension. arXiv preprint 11X1v:1910.13461, 2019.
Xiang Li, Wenhai Wang, Xiaolin Hu, and Jian Yang. Selective kernel networks. In CVPR, pp- 510-519, 2019. 3
Tsung-V Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In ECCV, pp. 740-755, 2014. 2, 3,8
Yinhan Liu, Myle Out, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 1rXh::1977.11622 2019. 1
Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang. Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model. arXiv preprint vrXiv:241110016 2024. 4
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, pp. 10012- 10022, 2021. 7,7
Zhuang Liu, Hanzi Mao, Chao- Yuan Wu, Christoph Feichtenhofer. Trevor Darrell, and Saining Xie. A convnet for the 2020s. In CVPR, pp. 11976-11986, 2022. 3
Nya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint arXiv:1608.03983 2016. 7
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint
arXiv:1711.05101, 2017. 7,8
Long Memory. Long short-term memory. Neural Computation, 98):1735-1788, 2010. 3
Maxim Milakov and Natalia Gimelshein, Online normalizer calculation for softmax. arXiv preprint arXiv: 180:. 02867, 2018. 16
12

---

Published as a conference paper at ICLR 2025
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, et al. Rwkv: Reinventing mns for the transformer era. arXiv preprint arXiv:2305.13048, 2023. 2, 4,
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for quence modeling. Advances in Neural Information Processing Systems, 36, 2024. 3
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training. 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019. 1
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: /910. 10683, 2019. 1
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint 1nXiv:1109.1555 2014. 3
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.049:3 04933, 2022a.
Shaden Smith, Mostofa Patwary, Brandon Norick, Patrick LeGresley, Samyam Rajbhandari,, Jared Casper, Zhun Liu, Shrimai Prabhumoye, George Zerveas, Vijay Korthikanti, et al. Using deep- speed and megatron to train aegatronnturingg nlg 530b, a large-scale generative language model. arXiv preprint ar arXiv:2201.11999 2022b.
Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer. How to train your vit? data, augmentation, and regularization in vision transformers. arXiv preprint ar XiXv:2106.110700 2021. 8,
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuging Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint wwxiiv23077088211. 2023. 3.4
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Du- mitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In CVPR, pp. 1 1-9, 2015. 3
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In ICML, pp. 10347-10357, 2021a. 2, 3, 7, 8, 9,
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In ICCV, pp. 32-42, 2021b. 2,7
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NeurlPS, 30, 2017. 1,4, 17
Ashish Vaswani, Prajit Ramachandran, Aravind Srinivas, Niki Parmar, Blake Hechtman, and Jonathon Shlens. Scaling local self-attention for parameter efficient visual backbones. In CVPR, pp. 12894-12904, 2021. 3.
Jue Wang, Wentao Zhu, Pichao Wang, Xiang Yu, Linda Liu, Mohamed Omar, and Raffay Hamid. Selective structured state-spaces for long-form video understanding. In CVPR, pp. 6387-6397, 2023a.
Sinong Wang, Belinda 7. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768. 2020. 3
Wenhai Wang. Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song. Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without
convolutions. In ICCV, pp. 568-578, 2021. 3, 4,
13

---

Published as a conference paper at ICLR 2025
Wenhai Wang. Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pvtv2: Improved baselines with pyramid vision transformer. CVMJ, pp. 1-10, 2022. 3,4
Wenhai Wang, Jifeng Dai, Zhe Chen, Zhenhang Huang, Zhiqi Li, Xizhou Zhu, Xiaowei Hu, Tong Lu, Lewei Lu, Hongsheng Li, et al. Internimage: Exploring large-scale vision foundation models with deformable convolutions. In CVPR, pp. #4408-14419 2023b. 3
Sitong Wu, Tianyi Wu, Haoru Tan, and Guodong Guo. Pale transformer: general vision trans-
former backbone with pale-shaped attention. In AAAI, volume 36, pp. 2731-2739, 2022. 3
Tete Xiao, Yingcheng Liu, Bolei Zhou, Yuning Jiang, and Jian Sun. Unified perceptual parsing for scene understanding. In ECCV, pp. 434, 2018. 8
Yuwen Xiong, Zhiqi Li, Yuntao Chen, Feng Wang, Xizhou Zhu, Jiapeng Luo, Wenhai Wang, Tong Lu, Hongsheng Li, Yu Qiao, et al. Efficient deformable convnets: Rethinking dynamic and sparse operator for vision applications. arXiv preprint 10..:::2011069977 2024.
Yuanhan Zhang, Qinghong Sun, Yichun Zhou, Zexin He, Zhenfei Yin, Kun Wang. Lu Sheng, Yu Qiao, Jing Shao, and Ziwei Liu. Bamboo: Building mega-scale vision dataset continually with human-machine synergy. arXiv preprint arXiv:2203.07844 2022. 8
Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In CVPR, pp. 633-641, 2017. 9
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vi- sion mamba: Efficient visual representation learning with bidirectional state space model. arXiv preprint -=XXi:2401100411,, 2024. 2, 3, 4, 16, 17
Xizhou Zhu Han Hu, Stephen Lin, and Jifeng Dai. Deformable convnets v2: More deformable, better results. CVPR, pp. 9308-9316, 2019.
14

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1
RNN FORM FORWARD AND BACKWARD
The attention mechanism in the spatial-mix module uses an RNN form forward and backward to achieve linear complexity of the input token number T. The following sections give more details of the operation.
States
Recurrence Relation
Initial Value
W d + 2
0
by W (bi-1 ke+ e Ve41)
b-1
>is e -1t-11utkke
Ct -ct-1 t1
C-1
d
de = w (dh-1 - ce
d-1 Σ e -1-1)ottt
day
w (1r -0(-1)
dm
dul
dw
db
db,
w dbi-1 + b.-1 1- k (e41) v(41)
db
Lio1 - 1)ee (1- (i-kkkek
de
des
w ( -(--1
de
dw
du
de
dde
w dd + di-1 c411
dd_1
ET- -(i 1)e (f- 1jwtts
dw
des
dw
go
g+; w (gor-1 gli- Vi
jiden)
goa Σ/ 8t Vi
iden 1i- 13w
gb
gb; = w gbb-1 + 81/b-1 Iden
gbo = 0
ge
gO: w (ger-1 gy/+ y/
iden ys)
goo
gyss iden yi (i-1))
gd
gd: gd:-1 + -
ye
gdo
Table 6: RNN states of Bi-WKV in forward and backward process. The update in recurrence relations has a fixed number of FL.OPs. w and w+ denotes the grow or decay vector e//T and ww//T The calculation of initial values O(TC) which does not affect the final complexity.
A.1.1 BACKWARD EQUATION
The backward process acquires the gradient of output matrix wkv (denotes y) passed from the previous layer (denotes as gy) calculate the gradient of each input. The saved inputs are vectors w, u E IRC key and value matrices K, V E TTxC' (the batch dimension is omitted). The new input is the gradient gy E T' provided by the backpropagation. The outputs include the gradients gw, gu E CCC matrices gK, gV E I''xC' corresponding to the inputs, respectively. From the RNN form of the forward process, the backward can be represented in an RNN form with a linear complexity related to the token number T. Some intermediate variables are listed as follows:
pum y
0t-1 b e e ddde 1/(a-1 d + e ettke Ve
Vi mum iden
1/
(10)
The outputs of backward propagation are listed as follows:
gw
St
iden
da
dw
dbc-1
dw
ua( do dw
dde-1 )),
dw
(11)
T-1
gu
gi/t Mt
iden
(-Hi (
c
atke
vt).
(12)
t=0
gkt =gb:
A gd: et gvi We iden (e tetu ye cctn)
(13)
+ ga: k Vi ge: et
gu: =gb tt + ga: tt + gy: ye
idon butu
(14)
The RNN states and their initial values and recurrence relations are provided in Table 6. From the recurrence relations, all updates have complexity of OCC), which means the number of FLOPs for each update is fixed. Therefore, the final backward complexity is OssTC' where s denotes the sum of the FL.OPs for all equations.
15

---

Published as a conference paper at ICLR 2025
80
80
80
Top-1 Aocuraay
70
Top-1 A
Top-1 Aocuracy
700
60
60
50
50
50
40
deit-t
40
deit-s
40
deit-b
30
vIwkv-t
30
vIWkV-s
30
vIWkv-b
20
20
224
384
768
1024
224
384
448
512
768
1024
224
384
768
1024
resolution
resolution
resolution
Figure 4: Performance of VRWKV and DeiT (Touvron et al., 2021a) on ImageNet-11 1K (Deng et al., 2009). All models are trained on fixed resolution of 224 x 224 and evaluated on different resolutions. Our VRWKV shows an obvious robustness different resolutions.
Method
#Param
Top-1 Acc
Vim-T
7M
76.1
VRWKV-T
75.1
Vim-S
26M
80.5
VRWKV-S
24M
80.1
Vim-B
98M
81.9
VRWKV-B
94M
82.0
Vim-L
NA
NA
VRWKV-L
335M
86.0
Table 7: Comparison with Vision Mamba (Zhu et al., 2024)(Vim)
ImageNet- (Deng et al.,
2009). "NA" denotes not available.
A.1.2 IMPLEMENTATION DETAILS
numerical trick compute safe exponential (Peng et al., 2023; Milakov Gimelshein, 2018) used to avoid overflow in the exponential terms the recurrence in the forward and backward process. An example of the update of state a is shown as follows:
q = max(pt-1 w/T,kt), a'
(15)
PC- q) a' exp(k: g) Ue.
(16)
4.
(17)
The exponential terms in the new state a' are forced to be smaller than by subtracting the max value. The subtracted part stored in is divided automatically when calculating wke.
A.2 ROBUSTNESS ON IMAGE RESOLUTION
Settings. In this experiment, we aim to explore whether the proposed VRWKV exhibits distinct properties compared to ViT. To this end, evaluated the performance of different variants of DeîT (Touvron et al., 2021a) and VRWKV at different resolutions on the ImageNet- 1K (Deng et al., resolution of 224 224, evaluated the models across range of resolutions, from 224 x 224 to 1024 x 1024.
2009) classification task. While the training was standardized
Results. As shown in Figure 4, our VRWKV demonstrates stronger robustness when evaluated on higher resolution. In contrast to DeiT' (Touvron et al., 2021a), VRWKV performs better as the resolution slightly increases. For example, VRWKV-B achieved top-1 accuracy 82.5% at 384 x 384 resolution, marking an improvement of 0.5 points over its accuracy the training resolution. When the test resolution scales up to 1024 1024, VRWKV-B still maintains an accuracy of 67.2%, while DeiT-B only achieves accuracy of 57.5%. This indicates that VRWKV has stronger potential and robustness in high-resolution scenarios and a good alternative to VIT for high-resolution tasks.
16

---

Published as a conference paper at ICLR 2025
A.3 COMPARISON TO VISION MAMBA
In the field of visual linear attention mechanisms, Vision Mamba (Zhu et al., 2024)(Vim) stands significant implementation that has garnered considerable interest. In this study, we compare the performance and efficiency of two models in visual tasks.
3
Vanilla A
infermnce tine(ns)
2.5
Vision Mamba
2
Bi-WKV
1.5
Classification Performance. We compare the classifica- tion accuracy on ImageNet-IK IK (Deng et al., 2009). As reported in Table 7, Vim has higher Top-1 Ace in tiny and small sizes, while the base size models achieve compa- rable performance. Benefitting from our careful stability design, VRWKV can scale up larger models while Vim faces instability issues in the training process.
0.5
224
736 input resolution
Inference Efficiency. We compare the inference speed of three attention mechanisms: Vanilla Attn (Vaswani et al., 2017), Bi- Bi-WKV, and Vision Mamba, shown in Figure 5. As the input resolution increases, the inference cost of vanilla attention quickly surpasses that of Bi-WKV and Vision Mamba. With our optimizations and design on CUDA, our Bi-WKV demonstrates faster speeds than Vision Mamba at the same input resolution.
Figure 5: Inference time of aften- tion mechanisms. Input resolutions are scanned from 224 to 1024. All experi- ments are run on Nvidia A100.
17


Reviews:
[{"review_id": "IPjRrz7Qcr", "reviewer": "Reviewer_QwNu", "summary": "The paper introduces a new network architecture, VISION-RWKV, a vision-adapted version of the RWKV network from the NLP community, which employs an RNN-based linear attention mechanism. The authors have made necessary adaptations to the RWKV to better suit visual tasks, recognizing the differences between vision and language processing tasks.", "questions": "see weaknesses", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "1.The paper is well-written and easy to follow;\n\n2.Extensive empirical evidence supports the effectiveness of the model, indicating its practical application potential;\n\n3.The adaptation of RWKV for visual tasks goes beyond mere transfer, incorporating modifications that enhance its suitability for image processing.", "weaknesses": "Major：Although the paper discusses RWKV and mamba as prominent RNN-based linear attention models transitioning from NLP to computer vision, it lacks a direct comparison with the mamba (vision) model. Such a comparison would be valuable for assessing the respective strengths and weaknesses of each model in the field of computer vision.\n\nMinor：From an innovation perspective, the approach of adapting other domains' mature architectures to computer vision, similar to what was done post-transformer, appears somewhat incremental.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "624oHBu3Sm", "reviewer": "Reviewer_cUXV", "summary": "Authors propose an approach that adapts NLP models for RWKV to vision tasks, termed as V-RWKV. Their proposed method could be looked upon as a cost effective solution than ViTs. They propose Quad-directional Shift and change the causal attention to bidirectional one towards learning locl concepts that are more relevant in vision rather than simple text. They evaluate their approach on the image classification, object detection and semantic segmentation tasks.", "questions": "Clarifying the practical benefit from their proposed approach and why not simply use vanilla ViTs considering the current gains are not that considerable when looking at Tables 2 and 4. Hence, the reason I am leaning towards a marginal reject but since the method seems to provide interesting direction and their results for detection seems sufficiently good I am not going lower.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- Comprehensive results across three tasks.\n- Interesting direction to replace ViTs with other efficient techniques that provide on-par or better performance.", "weaknesses": "- The gains in efficiency seem to be relatively minor, i.e., it is not an order of magnitude which still brings the question whether it is worth exploring these type of models to replace ViTs to begin with or not.\n\n For example, in Table 2 ViT-L vs. VRWKV-L 191.1 vs 189.5G FLOPS and parameters at 309.5M vs. 334.9M respectively. I think the reduction in FLOPS when scaling to Large variant seems to be around 3G. Similarly in Table 4, FLOPS 446.8 vs. 421.9G with the expense of an increase in the number of parameters. I am not expert in such type of methods focused on improving efficiency but I do not see the results are impressive enough to show the benefit from the V-RWKV design, especially that it is increasing the parameters.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ku5Nr8lhqM", "reviewer": "Reviewer_weiY", "summary": "The authors describe a Vision-RWKV architecture that employs novel techniques like bi-directional RWKV, quad token shifting method Q-shift etc. These techniques helps Vision-RWKV  architecture surpass window-based ViTs and comparable to global attention ViTs along with lower FLOPs, lower GPU memory cost and faster processing as shown in Figure 1. Although the MAE finetuning is not as straightforward as typical finetuning on downstream task, overall this paper has a good contribution for the vision research community.", "questions": "- Can you please clarify what you meant in line 527-528 by task fine-tuning using Q-shift manner?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Novel contribution for quad-directional token shifting called Q-shift. This essentially increases the models range of semantic understanding\n- Authors expanded causal RWKV to a bidirectional global RWKV. They modified the exponent in the RWKV attention that leads to transforming the absolution positional bias to relative bias.\n- Well written paper with through experimentation including MAE pretraining", "weaknesses": "- The authors implemented a bidirectional shift operation that removed the vertical shift in Q-shift, thereby enabling for MAE pretraining. IMO this is a source of complexity. As a result, MAE finetuning needs to be done in Q-shift manner", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)