RWKV: Reinventing RNNs for the Transformer Era (original) (raw)

\setitemize

noitemsep \setenumeratenoitemsep

Bo Peng1,2 Eric Alcaide2,3,411footnotemark: 1 Quentin Anthony2,511footnotemark: 1
Alon Albalak2,6 Samuel Arcadinho2,7 Stella Biderman2,8 Huanqi Cao9 Xin Cheng10
Michael Chung11 Xingjian Du1 Matteo Grella12 Kranthi Kiran GV2,13 Xuzheng He2
Haowen Hou14 Jiaju Lin1 Przemysław Kazienko15 Jan Kocoń15 Jiaming Kong16
Bartłomiej Koptyra15 Hayden Lau2 Krishna Sri Ipsit Mantri17 Ferdinand Mom18,19
Atsushi Saito2,20 Guangyu Song21 Xiangru Tang22 Bolun Wang23 Johan S. Wind24
Stanisław Woźniak15 Ruichong Zhang9 Zhenyuan Zhang2 Qihang Zhao25,26
Peng Zhou23 Qinghua Zhou5 Jian Zhu27 Rui-Jie Zhu28,29
1Generative AI Commons2EleutherAI3U. of Barcelona4Charm Therapeutics5Ohio State U.6U. of C., Santa Barbara
7Zendesk8Booz Allen Hamilton9Tsinghua University10Peking University11Storyteller.io12Crisis2413New York U.
14National U. of Singapore15Wroclaw U. of Science and Technology16Databaker Technology17Purdue U.18Criteo AI Lab
19Epita20Nextremer21Moves22Yale U.23RuoxinTech24U. of Oslo25U. of Science and Technology of China
26Kuaishou Technology27U. of British Columbia28U. of C., Santa Cruz29U. of Electronic Science and Technology of China

Abstract

Transformers have revolutionized almost all natural language processing (NLP) tasks but suffer from memory and computational complexity that scales quadratically with sequence length. In contrast, recurrent neural networks (RNNs) exhibit linear scaling in memory and computational requirements but struggle to match the same performance as Transformers due to limitations in parallelization and scalability. We propose a novel model architecture, Receptance Weighted Key Value (RWKV), that combines the efficient parallelizable training of transformers with the efficient inference of RNNs.

Our approach leverages a linear attention mechanism and allows us to formulate the model as either a Transformer or an RNN, thus parallelizing computations during training and maintains constant computational and memory complexity during inference. We scale our models as large as 14 billion parameters, by far the largest dense RNN ever trained, and find RWKV performs on par with similarly sized Transformers, suggesting future work can leverage this architecture to create more efficient models. This work presents a significant step towards reconciling trade-offs between computational efficiency and model performance in sequence processing tasks.111Code at: https://github.com/BlinkDL/RWKV-LM

1 Introduction

Deep learning has greatly advanced artificial intelligence, impacting a range of scientific and industrial uses. These often involve complex sequential data processing tasks such as natural language understanding, conversational AI, time-series analysis, and indirectly sequential formats like images and graphs Brown et al. (2020); Ismail Fawaz et al. (2019); Wu et al. (2020); Albalak et al. (2022). Predominant among these techniques include RNNs and Transformers Vaswani et al. (2017), each with specific benefits and drawbacks. RNNs require less memory, particularly for handling long sequences. However, they suffer from the vanishing gradient problem and non-parallelizability in the time dimension during training, limiting their scalability Hochreiter (1998); Le and Zuidema (2016).

Refer to caption

Figure 1: Average performance of RWKV models compared to transformers across twelve NLP tasks. For further details, see section 5.

Transformers emerged as a powerful alternative, adept at managing local and long-range dependencies and supporting parallelized training Tay et al. (2022). Models such as GPT-3 Brown et al. (2020), ChatGPT OpenAI (2022); Kocoń et al. (2023), LLaMA Touvron et al. (2023), and Chinchilla Hoffmann et al. (2022) showcase the potential of Transformers in NLP. However, the self-attention mechanism’s quadratic complexity makes it computationally and memory intensive for tasks involving long sequences and constrained resources. This has stimulated research to enhance Transformers’ scalability, sometimes sacrificing some of their effectiveness Wang et al. (2020); Zaheer et al. (2020); Dao et al. (2022a).

Table 1: Inference complexity comparison with different Transformers. Here T𝑇Titalic_T denotes the sequence length, d𝑑ditalic_d the feature dimension, c𝑐citalic_c is MEGA’s chunk size of quadratic attention, and s𝑠sitalic_s is the size of a local window for AFT.

To tackle these challenges, we introduce the Receptance Weighted Key Value (RWKV) model, combining the strengths of RNNs and Transformers while circumventing key drawbacks. RWKV alleviates memory bottleneck and quadratic scaling associated with Transformers Katharopoulos et al. (2020) with efficient linear scaling, while maintaining the expressive properties of the Transformer, such as parallelized training and robust scalability. RWKV reformulates the attention mechanism with a variant of linear attention, replacing traditional dot-product token interaction with more effective channel-directed attention. This implementation, without approximation, offers the lowest computational and memory complexity; see Table 1.

The motivation behind RWKV is to balance computational efficiency with expressive capacity in neural networks. It offers a solution for handling large-scale models with billions of parameters, exhibiting competitive performance at a reduced computational cost. Experiments suggest RWKV addresses scaling and deployment challenges in AI, especially for sequential data processing, pointing towards more sustainable and efficient AI models.

Our contributions in this paper are as follows:

2 Background

Here we briefly review the fundamentals of RNNs and Transformers.

2.1 Recurrent Neural Networks (RNNs)

Popular RNN architectures such as LSTM Hochreiter and Schmidhuber (1997) and GRU Chung et al. (2014) are characterized by the following formulation (shown for LSTM, others can be reasoned similarly):

ftsubscript𝑓𝑡\displaystyle f_{t}italic_f start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =σg⁢(Wf⁢xt+Uf⁢ht−1+bf),absentsubscript𝜎𝑔subscript𝑊𝑓subscript𝑥𝑡subscript𝑈𝑓subscriptℎ𝑡1subscript𝑏𝑓\displaystyle=\sigma_{g}(W_{f}x_{t}+U_{f}h_{t-1}+b_{f}),= italic_σ start_POSTSUBSCRIPT italic_g end_POSTSUBSCRIPT ( italic_W start_POSTSUBSCRIPT italic_f end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + italic_U start_POSTSUBSCRIPT italic_f end_POSTSUBSCRIPT italic_h start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_b start_POSTSUBSCRIPT italic_f end_POSTSUBSCRIPT ) , (1)
itsubscript𝑖𝑡\displaystyle i_{t}italic_i start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =σg⁢(Wi⁢xt+Ui⁢ht−1+bi),absentsubscript𝜎𝑔subscript𝑊𝑖subscript𝑥𝑡subscript𝑈𝑖subscriptℎ𝑡1subscript𝑏𝑖\displaystyle=\sigma_{g}(W_{i}x_{t}+U_{i}h_{t-1}+b_{i}),= italic_σ start_POSTSUBSCRIPT italic_g end_POSTSUBSCRIPT ( italic_W start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + italic_U start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT italic_h start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_b start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ) , (2)
otsubscript𝑜𝑡\displaystyle o_{t}italic_o start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =σg⁢(Wo⁢xt+Uo⁢ht−1+bo),absentsubscript𝜎𝑔subscript𝑊𝑜subscript𝑥𝑡subscript𝑈𝑜subscriptℎ𝑡1subscript𝑏𝑜\displaystyle=\sigma_{g}(W_{o}x_{t}+U_{o}h_{t-1}+b_{o}),= italic_σ start_POSTSUBSCRIPT italic_g end_POSTSUBSCRIPT ( italic_W start_POSTSUBSCRIPT italic_o end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + italic_U start_POSTSUBSCRIPT italic_o end_POSTSUBSCRIPT italic_h start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_b start_POSTSUBSCRIPT italic_o end_POSTSUBSCRIPT ) , (3)
c~tsubscript~𝑐𝑡\displaystyle\tilde{c}_{t}over~ start_ARG italic_c end_ARG start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =σc⁢(Wc⁢xt+Uc⁢ht−1+bc),absentsubscript𝜎𝑐subscript𝑊𝑐subscript𝑥𝑡subscript𝑈𝑐subscriptℎ𝑡1subscript𝑏𝑐\displaystyle=\sigma_{c}(W_{c}x_{t}+U_{c}h_{t-1}+b_{c}),= italic_σ start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT ( italic_W start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + italic_U start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT italic_h start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_b start_POSTSUBSCRIPT italic_c end_POSTSUBSCRIPT ) , (4)
ctsubscript𝑐𝑡\displaystyle c_{t}italic_c start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =ft⊙ct−1+it⊙c~t,absentdirect-productsubscript𝑓𝑡subscript𝑐𝑡1direct-productsubscript𝑖𝑡subscript~𝑐𝑡\displaystyle=f_{t}\odot c_{t-1}+i_{t}\odot\tilde{c}_{t},= italic_f start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ⊙ italic_c start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_i start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ⊙ over~ start_ARG italic_c end_ARG start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , (5)
htsubscriptℎ𝑡\displaystyle h_{t}italic_h start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =ot⊙σh⁢(ct).absentdirect-productsubscript𝑜𝑡subscript𝜎ℎsubscript𝑐𝑡\displaystyle=o_{t}\odot\sigma_{h}(c_{t}).= italic_o start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ⊙ italic_σ start_POSTSUBSCRIPT italic_h end_POSTSUBSCRIPT ( italic_c start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) . (6)

Although RNNs can be factored into two linear blocks (W𝑊Witalic_W and U𝑈Uitalic_U) and an RNN-specific block (1)–(6), as noted by Bradbury et al. (2017), the data dependency relying on previous time steps prohibits parallelizing these typical RNNs.

2.2 Transformers and AFT

Introduced by Vaswani et al. (2017), Transformers are a class of neural networks that have become the dominant architecture for several NLP tasks. Instead of operating on sequences step-by-step like RNNs, Transformers rely on attention mechanisms to capture relationships between all input and all output tokens:

Attn⁢(Q,K,V)=softmax⁢(Q⁢K⊤)⁢V,Attn𝑄𝐾𝑉softmax𝑄superscript𝐾top𝑉\displaystyle\mathrm{Attn}(Q,K,V)=\mathrm{softmax}(QK^{\top})V,roman_Attn ( italic_Q , italic_K , italic_V ) = roman_softmax ( italic_Q italic_K start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT ) italic_V , (7)

where the multi-headness and scaling factor 1dk1subscript𝑑𝑘\frac{1}{\sqrt{d_{k}}}divide start_ARG 1 end_ARG start_ARG square-root start_ARG italic_d start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT end_ARG end_ARG is omitted for convenience. The core Q⁢K⊤𝑄superscript𝐾topQK^{\top}italic_Q italic_K start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT multiplication is an ensemble of pairwise attention scores between each token in a sequence, which can be decomposed as vector operations:

Attn⁢(Q,K,V)tAttnsubscript𝑄𝐾𝑉𝑡\displaystyle\mathrm{Attn}(Q,K,V)_{t}roman_Attn ( italic_Q , italic_K , italic_V ) start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =∑i=1Teqt⊤⁢ki⊙vi∑i=1Teqt⊤⁢ki.absentsuperscriptsubscript𝑖1𝑇direct-productsuperscript𝑒subscriptsuperscript𝑞top𝑡subscript𝑘𝑖subscript𝑣𝑖superscriptsubscript𝑖1𝑇superscript𝑒subscriptsuperscript𝑞top𝑡subscript𝑘𝑖\displaystyle=\frac{\sum_{i=1}^{T}e^{q^{\top}_{t}k_{i}}\odot v_{i}}{\sum_{i=1}% ^{T}e^{q^{\top}_{t}k_{i}}}.= divide start_ARG ∑ start_POSTSUBSCRIPT italic_i = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_T end_POSTSUPERSCRIPT italic_e start_POSTSUPERSCRIPT italic_q start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT italic_k start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_POSTSUPERSCRIPT ⊙ italic_v start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG start_ARG ∑ start_POSTSUBSCRIPT italic_i = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_T end_POSTSUPERSCRIPT italic_e start_POSTSUPERSCRIPT italic_q start_POSTSUPERSCRIPT ⊤ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT italic_k start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_POSTSUPERSCRIPT end_ARG . (8)

AFT Zhai et al. (2021), alternately formulates

Attn+⁢(W,K,V)t=∑i=1tewt,i+ki⊙vi∑i=1tewt,i+ki,superscriptAttnsubscript𝑊𝐾𝑉𝑡superscriptsubscript𝑖1𝑡direct-productsuperscript𝑒subscript𝑤𝑡𝑖subscript𝑘𝑖subscript𝑣𝑖superscriptsubscript𝑖1𝑡superscript𝑒subscript𝑤𝑡𝑖subscript𝑘𝑖\displaystyle\mathrm{Attn}^{+}(W,K,V)_{t}=\frac{\sum_{i=1}^{t}e^{w_{t,i}+k_{i}% }\odot v_{i}}{\sum_{i=1}^{t}e^{w_{t,i}+k_{i}}},roman_Attn start_POSTSUPERSCRIPT + end_POSTSUPERSCRIPT ( italic_W , italic_K , italic_V ) start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT = divide start_ARG ∑ start_POSTSUBSCRIPT italic_i = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_t end_POSTSUPERSCRIPT italic_e start_POSTSUPERSCRIPT italic_w start_POSTSUBSCRIPT italic_t , italic_i end_POSTSUBSCRIPT + italic_k start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_POSTSUPERSCRIPT ⊙ italic_v start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG start_ARG ∑ start_POSTSUBSCRIPT italic_i = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_t end_POSTSUPERSCRIPT italic_e start_POSTSUPERSCRIPT italic_w start_POSTSUBSCRIPT italic_t , italic_i end_POSTSUBSCRIPT + italic_k start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_POSTSUPERSCRIPT end_ARG , (9)

where {wt,i}∈RT×Tsubscript𝑤𝑡𝑖superscript𝑅𝑇𝑇\{w_{t,i}\}\in R^{T\times T}{ italic_w start_POSTSUBSCRIPT italic_t , italic_i end_POSTSUBSCRIPT } ∈ italic_R start_POSTSUPERSCRIPT italic_T × italic_T end_POSTSUPERSCRIPT is the learned pair-wise position biases, and each wt,isubscript𝑤𝑡𝑖w_{t,i}italic_w start_POSTSUBSCRIPT italic_t , italic_i end_POSTSUBSCRIPT is a scalar.

Inspired by AFT, RWKV takes a similar approach. However, for simplicity, it modifies the interaction weights so that it can be transformed into an RNN. Each wt,isubscript𝑤𝑡𝑖w_{t,i}italic_w start_POSTSUBSCRIPT italic_t , italic_i end_POSTSUBSCRIPT in RWKV is a channel-wise time decay vector multiplied by the relative position and traced backward from current time as it decays:

wt,i=−(t−i)⁢w,subscript𝑤𝑡𝑖𝑡𝑖𝑤\displaystyle w_{t,i}=-(t-i)w,italic_w start_POSTSUBSCRIPT italic_t , italic_i end_POSTSUBSCRIPT = - ( italic_t - italic_i ) italic_w , (10)

where w∈(R≥0)d𝑤superscriptsubscript𝑅absent0𝑑w\in(R_{\geq 0})^{d}italic_w ∈ ( italic_R start_POSTSUBSCRIPT ≥ 0 end_POSTSUBSCRIPT ) start_POSTSUPERSCRIPT italic_d end_POSTSUPERSCRIPT, with d𝑑ditalic_d the number of channels. We require w𝑤witalic_w to be non-negative to ensure that ewt,i≤1superscript𝑒subscript𝑤𝑡𝑖1e^{w_{t,i}}\leq 1italic_e start_POSTSUPERSCRIPT italic_w start_POSTSUBSCRIPT italic_t , italic_i end_POSTSUBSCRIPT end_POSTSUPERSCRIPT ≤ 1 and the per-channel weights decay backwards in time.

3 RWKV

The RWKV model architecture is defined by four fundamental elements that are intrinsic to the time-mixing and channel-mixing blocks:

These core elements interact multiplicatively at each timestep, as depicted in Figure 2.

Refer to caption Refer to caption

Figure 2: Elements within an RWKV block (left) and the complete RWKV residual block, equipped with a final head for language modeling (right).

3.1 Architecture

The RWKV model is composed of stacked residual blocks. Each block consists of a time-mixing and a channel-mixing sub-block, embodying recurrent structures to leverage past information.

This model uses a unique attention-like score update process, which includes a time-dependent softmax operation improving numerical stability and mitigating vanishing gradients (for rigorous proof, see Appendix H). It ensures that the gradient is propagated along the most relevant path. Additionally, layer normalization Ba et al. (2016) incorporated within the architecture aids in stabilizing the gradients, effectively addressing both vanishing and exploding gradient issues. These design elements not only enhance the training dynamics of deep neural networks but also facilitate the stacking of multiple layers, leading to superior performance over conventional RNN models by capturing complex patterns across different levels of abstraction (see also Appendix I).

Refer to caption

Figure 3: RWKV architecture for language modeling.

3.1.1 Token Shift

In this architecture, all linear projection vectors (R𝑅Ritalic_R, K𝐾Kitalic_K, V𝑉Vitalic_V in time-mixing, and R′superscript𝑅′R^{\prime}italic_R start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT, K′superscript𝐾′K^{\prime}italic_K start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT in channel-mixing) involved in computations are produced by linear interpolation between current and previous timestep inputs, facilitating a token shift.

The vectors for time-mixing computation are linear projections of linear combinations of the current and previous inputs of the block:

rtsubscript𝑟𝑡\displaystyle r_{t}italic_r start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =Wr⋅(μr⊙xt+(1−μr)⊙xt−1),absent⋅subscript𝑊𝑟direct-productsubscript𝜇𝑟subscript𝑥𝑡direct-product1subscript𝜇𝑟subscript𝑥𝑡1\displaystyle=W_{r}\cdot(\mu_{r}\odot x_{t}+(1-\mu_{r})\odot x_{t-1}),= italic_W start_POSTSUBSCRIPT italic_r end_POSTSUBSCRIPT ⋅ ( italic_μ start_POSTSUBSCRIPT italic_r end_POSTSUBSCRIPT ⊙ italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + ( 1 - italic_μ start_POSTSUBSCRIPT italic_r end_POSTSUBSCRIPT ) ⊙ italic_x start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT ) , (11)
ktsubscript𝑘𝑡\displaystyle k_{t}italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =Wk⋅(μk⊙xt+(1−μk)⊙xt−1),absent⋅subscript𝑊𝑘direct-productsubscript𝜇𝑘subscript𝑥𝑡direct-product1subscript𝜇𝑘subscript𝑥𝑡1\displaystyle=W_{k}\cdot(\mu_{k}\odot x_{t}+(1-\mu_{k})\odot x_{t-1}),= italic_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ⋅ ( italic_μ start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ⊙ italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + ( 1 - italic_μ start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) ⊙ italic_x start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT ) , (12)
vtsubscript𝑣𝑡\displaystyle v_{t}italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =Wv⋅(μv⊙xt+(1−μv)⊙xt−1),absent⋅subscript𝑊𝑣direct-productsubscript𝜇𝑣subscript𝑥𝑡direct-product1subscript𝜇𝑣subscript𝑥𝑡1\displaystyle=W_{v}\cdot(\mu_{v}\odot x_{t}+(1-\mu_{v})\odot x_{t-1}),= italic_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ⋅ ( italic_μ start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ⊙ italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + ( 1 - italic_μ start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ) ⊙ italic_x start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT ) , (13)

as are the channel-mixing inputs:

rt′subscriptsuperscript𝑟′𝑡\displaystyle r^{\prime}_{t}italic_r start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =Wr′⋅(μr′⊙xt+(1−μr′)⊙xt−1),absent⋅subscriptsuperscript𝑊′𝑟direct-productsubscriptsuperscript𝜇′𝑟subscript𝑥𝑡direct-product1subscriptsuperscript𝜇′𝑟subscript𝑥𝑡1\displaystyle=W^{\prime}_{r}\cdot(\mu^{\prime}_{r}\odot x_{t}+(1-\mu^{\prime}_% {r})\odot x_{t-1}),= italic_W start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_r end_POSTSUBSCRIPT ⋅ ( italic_μ start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_r end_POSTSUBSCRIPT ⊙ italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + ( 1 - italic_μ start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_r end_POSTSUBSCRIPT ) ⊙ italic_x start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT ) , (14)
kt′subscriptsuperscript𝑘′𝑡\displaystyle k^{\prime}_{t}italic_k start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =Wk′⋅(μk′⊙xt+(1−μk′)⊙xt−1).absent⋅subscriptsuperscript𝑊′𝑘direct-productsubscriptsuperscript𝜇′𝑘subscript𝑥𝑡direct-product1subscriptsuperscript𝜇′𝑘subscript𝑥𝑡1\displaystyle=W^{\prime}_{k}\cdot(\mu^{\prime}_{k}\odot x_{t}+(1-\mu^{\prime}_% {k})\odot x_{t-1}).= italic_W start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ⋅ ( italic_μ start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ⊙ italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + ( 1 - italic_μ start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) ⊙ italic_x start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT ) . (15)

The token shift is implemented as a simple offset in the temporal dimension at each block using the PyTorch (Paszke et al., 2019) library as nn.ZeroPad2d((0,0,1,-1)).

3.1.2 WKV Operator

The computation of the W⁢K⁢V𝑊𝐾𝑉WKVitalic_W italic_K italic_V operator in our model parallels the method used in Attention Free Transformer (AFT) Zhai et al. (2021). However, unlike AFT where W𝑊Witalic_W is a pairwise matrix, our model treats W𝑊Witalic_W as a channel-wise vector that is modified by relative position. In our model, this recurrent behavior is defined by the time-dependent update of the W⁢K⁢V𝑊𝐾𝑉WKVitalic_W italic_K italic_V vectors, formalized in the following equation:

w⁢k⁢vt𝑤𝑘subscript𝑣𝑡\displaystyle wkv_{t}italic_w italic_k italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =∑i=1t−1e−(t−1−i)⁢w+ki⊙vi+eu+kt⊙vt∑i=1t−1e−(t−1−i)⁢w+ki+eu+kt.absentsuperscriptsubscript𝑖1𝑡1direct-productsuperscript𝑒𝑡1𝑖𝑤subscript𝑘𝑖subscript𝑣𝑖direct-productsuperscript𝑒𝑢subscript𝑘𝑡subscript𝑣𝑡superscriptsubscript𝑖1𝑡1superscript𝑒𝑡1𝑖𝑤subscript𝑘𝑖superscript𝑒𝑢subscript𝑘𝑡\displaystyle=\frac{\sum_{i=1}^{t-1}e^{-(t-1-i)w+k_{i}}\odot v_{i}+e^{u+k_{t}}% \odot v_{t}}{\sum_{i=1}^{t-1}e^{-(t-1-i)w+k_{i}}+e^{u+k_{t}}}.= divide start_ARG ∑ start_POSTSUBSCRIPT italic_i = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_t - 1 end_POSTSUPERSCRIPT italic_e start_POSTSUPERSCRIPT - ( italic_t - 1 - italic_i ) italic_w + italic_k start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_POSTSUPERSCRIPT ⊙ italic_v start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_u + italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_POSTSUPERSCRIPT ⊙ italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_ARG start_ARG ∑ start_POSTSUBSCRIPT italic_i = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_t - 1 end_POSTSUPERSCRIPT italic_e start_POSTSUPERSCRIPT - ( italic_t - 1 - italic_i ) italic_w + italic_k start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_POSTSUPERSCRIPT + italic_e start_POSTSUPERSCRIPT italic_u + italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_POSTSUPERSCRIPT end_ARG . (16)

To circumvent any potential degradation of W𝑊Witalic_W, we introduce a vector U𝑈Uitalic_U that separately attends to the current token. More information about this can be found in Appendix I.

3.1.3 Output Gating

Output gating is implemented in both time-mixing and channel-mixing blocks using the sigmoid of the receptance, σ⁢(r)𝜎𝑟\sigma(r)italic_σ ( italic_r ). The output vector otsubscript𝑜𝑡o_{t}italic_o start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT post the W⁢K⁢V𝑊𝐾𝑉WKVitalic_W italic_K italic_V operator is given by:

otsubscript𝑜𝑡\displaystyle o_{t}italic_o start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =Wo⋅(σ⁢(rt)⊙w⁢k⁢vt).absent⋅subscript𝑊𝑜direct-product𝜎subscript𝑟𝑡𝑤𝑘subscript𝑣𝑡\displaystyle=W_{o}\cdot(\sigma(r_{t})\odot wkv_{t}).= italic_W start_POSTSUBSCRIPT italic_o end_POSTSUBSCRIPT ⋅ ( italic_σ ( italic_r start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) ⊙ italic_w italic_k italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) . (17)

In the channel-mixing block, a similar operation is performed:

ot′subscriptsuperscript𝑜′𝑡\displaystyle o^{\prime}_{t}italic_o start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =σ(rt′)⊙(Wv′⋅max(kt′,0)2),\displaystyle=\sigma(r^{\prime}_{t})\odot(W^{\prime}_{v}\cdot\max({k^{\prime}_% {t}},0)^{2}),= italic_σ ( italic_r start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) ⊙ ( italic_W start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ⋅ roman_max ( italic_k start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , 0 ) start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ) , (18)

where we adopt the squared ReLU activation function So et al. (2021).

3.2 Transformer-like Training

RWKV can be efficiently parallelized using a technique called time-parallel mode, reminiscent of Transformers. The time complexity of processing a batch of sequences in a single layer is O⁢(B⁢T⁢d2)𝑂𝐵𝑇superscript𝑑2O(BTd^{2})italic_O ( italic_B italic_T italic_d start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ), primarily consisting of matrix multiplications Wλsubscript𝑊𝜆W_{\lambda}italic_W start_POSTSUBSCRIPT italic_λ end_POSTSUBSCRIPT, where λ∈{r,k,v,o}𝜆𝑟𝑘𝑣𝑜\lambda\in\{r,k,v,o\}italic_λ ∈ { italic_r , italic_k , italic_v , italic_o } (assuming B𝐵Bitalic_B sequences, T𝑇Titalic_T maximum tokens, and d𝑑ditalic_d channels). In contrast, updating attention scores w⁢k⁢vt𝑤𝑘subscript𝑣𝑡wkv_{t}italic_w italic_k italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT involves a serial scan (see Appendix D for more detail) and has complexity O⁢(B⁢T⁢d)𝑂𝐵𝑇𝑑O(BTd)italic_O ( italic_B italic_T italic_d ).

The matrix multiplications can be parallelized similarly to Wλsubscript𝑊𝜆W_{\lambda}italic_W start_POSTSUBSCRIPT italic_λ end_POSTSUBSCRIPT, where λ∈{Q,K,V,O}𝜆𝑄𝐾𝑉𝑂\lambda\in\{Q,K,V,O\}italic_λ ∈ { italic_Q , italic_K , italic_V , italic_O } in conventional Transformers. The element-wise W⁢K⁢V𝑊𝐾𝑉WKVitalic_W italic_K italic_V computation is time-dependent but can be readily parallelized along the other two dimensions Lei et al. (2018)333For extremely long sequences, more sophisticated methods such as Martin and Cundy (2017) that parallelize over sequence length could be used..

3.3 RNN-like Inference

Recurrent networks commonly utilize the output at state t𝑡titalic_t as input at state t+1𝑡1t+1italic_t + 1. This usage is also observed in the autoregressive decoding inference of language models, where each token must be computed before being passed to the next step. RWKV takes advantage of this RNN-like structure, known as time-sequential mode. In this context, RWKV can be conveniently formulated recursively for decoding during inference, as demonstrated in Appendix D.

3.4 Additional Optimizations

Custom Kernels

To address inefficiencies in the W⁢K⁢V𝑊𝐾𝑉WKVitalic_W italic_K italic_V computation arising from the sequential nature of the task when using standard deep learning frameworks, we have developed a custom CUDA kernel. This kernel enables the execution of a single compute kernel on training accelerators, while all other parts of the model, such as matrix multiplications and point-wise operations, are already inherently parallelizable and efficient.

Small Init Embedding

During the initial stage of training a transformer model Vaswani et al. (2017), we observe that the embedding matrix undergoes slow changes, presenting a challenge for the model to move away from its initial noisy embedding state. To address this issue, we propose an approach that involves initializing the embedding matrix with small values and subsequently applying an additional LayerNorm operation. This accelerates and stabilizes the training process, allowing for the training of deep architectures with post-LN components. The effectiveness of this approach is demonstrated in Figure 9, illustrating improved convergence by enabling the model to quickly transition away from the initially small embedding. This is achieved through small changes occurring in a single step, which subsequently lead to substantial alterations in directions and further notable changes after the LayerNorm operation.

Custom Initialization

Building on principles from previous works (He et al., 2016; Jumper et al., 2021), we adopt an initialization strategy where parameters are set to values resembling an identity mapping while breaking symmetry to establish a clear information flow. The majority of weights are initialized to zero, and linear layers do not employ biases. Detailed formulas are given in Appendix E. We observe that the choice of initialization plays a crucial role in both the speed and quality of convergence (refer to Appendix F for further details).

3.5 Implementation

RWKV is implemented using the PyTorch Deep Learning Library Paszke et al. (2019). We integrate additional optimization strategies inspired by DeepSpeed Rasley et al. (2020) into the system, improving its efficiency and scalability.

The model begins with an embedding layer, as detailed in Section 3.4. Following this are several identical residual blocks arranged sequentially. These are depicted in Figures 2 and 3 and adheres to the principles outlined in Section 3.1.1. After the last block, a simple output projection head, consisting of a LayerNorm Ba et al. (2016) and a linear projection, is employed for logits generation for next-token prediction and computation of the cross-entropy loss during training.

4 Trained Models and Computing Costs

To demonstrate the scalability of RWKV, we train six models ranging from 169 million to 14 billion parameters as shown in Table 2. All models are trained for one epoch (330 billion tokens) on the Pile (Gao et al., 2020; Biderman et al., 2022).

Table 2: RWKV model architectures and FLOP counts. Further details of these hyperparameters are elaborated upon in Appendix G.

The number of parameters for each model is computed using the formula: # parameters=2⁢V⁢D+13⁢D2⁢L+D⁢(11⁢L+4)# parameters2𝑉𝐷13superscript𝐷2𝐿𝐷11𝐿4\text{\# parameters}=2VD+13D^{2}L+D(11L+4)# parameters = 2 italic_V italic_D + 13 italic_D start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT italic_L + italic_D ( 11 italic_L + 4 ) where V𝑉Vitalic_V = 50277 is the vocabulary size, D𝐷Ditalic_D represents the Model Dimension and L𝐿Litalic_L corresponds to the number of layers. FLOPs is for a forward pass for one token. It was calculated as 2⁢(2⁢V⁢D+13⁢D2⁢L)22𝑉𝐷13superscript𝐷2𝐿2(2VD+13D^{2}L)2 ( 2 italic_V italic_D + 13 italic_D start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT italic_L ), which is the twice (add and multiply) the number of parameters in linear layers. The backwards pass FLOPs can be approximated as twice that of the forward pass, giving a total of 6⁢(2⁢V⁢D+13⁢D2⁢L)62𝑉𝐷13superscript𝐷2𝐿6(2VD+13D^{2}L)6 ( 2 italic_V italic_D + 13 italic_D start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT italic_L ) FLOP per token. Notably, this matches the standard formula for FLOP calculations in transformers Kaplan et al. (2020): FLOP=6⋅[# tokens]⋅[# parameters]FLOP⋅6delimited-[]# tokensdelimited-[]# parameters\text{FLOP}=6\cdot[\text{\# tokens}]\cdot[\text{\# parameters}]FLOP = 6 ⋅ [ # tokens ] ⋅ [ # parameters ].

4.1 Additional Training Details

For training, we use the standard Adam optimizer without weight decay, use bfloat16 precision, and train with a context length of 1024 tokens. Further details on hyperparameters are in Appendix G. Diverting from standard practice for transformers, we apply exponential decay to our learning rate. We also incorporate the auxiliary loss introduced by PaLM Chowdhery et al. (2022), supplementing the standard cross-entropy loss function. This auxiliary loss encourages the softmax normalizer to approximate zero closely. As for the learning rate schedule, it remains constant for the initial iterations, and subsequently decays exponentially.

4.2 Scaling Laws

Scaling laws Kaplan et al. (2020); Henighan et al. (2020); Hoffmann et al. (2022); Muennighoff et al. (2023) in language models refer to the mathematical relationships that describe how the performance of a language model changes with respect to various factors. These factors can include the model size (N𝑁Nitalic_N), dataset size (D𝐷Ditalic_D), or the optimally allocated compute budget (Cminsubscript𝐶minC_{\rm min}italic_C start_POSTSUBSCRIPT roman_min end_POSTSUBSCRIPT). Scaling laws are important for two primary reasons: they allow us to make predictions and plans regarding the costs and performance of large models before they are trained via interpolation and extrapolation Black et al. (2022); Le Scao et al. (2022) and the contexts in which they fail provides rich feedback on important areas for future research Wei et al. (2022a); Biderman et al. (2023a).

Previous work on scaling laws for RNNs has claimed that LSTMs do not strictly follow the same log-log linear scaling that transformers do Kaplan et al. (2020). We train 45 RWKV models for a variety of pairs (dataset, parameters) and find that RWKV does follow the same general form of the scaling law that is well established for transformers. Figure 4 shows our results for loss as a function of compute, with the linear fit to the Pareto optimal points holding an r2superscript𝑟2r^{2}italic_r start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT value of 0.9940.9940.9940.994. Even when we extrapolate our curve an additional order of magnitude (blue), we find an extremely good fit with an r2superscript𝑟2r^{2}italic_r start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT of 0.8750.8750.8750.875.

Refer to caption

Figure 4: Scaling laws curves for RWKV models

5 Evaluations

Having demonstrated the scalability of RWKV models in the previous section, we now turn our attention to their competitiveness with traditional transformers. We focus on two questions:

Competitiveness

Is RWKV competitive against quadratic transformer architectures with the same amount of compute?

Long Context

Does increasing the context length of RWKV yield better language modeling loss when RWKV models are trained for context lengths that most open-sourced quadratic transformers cannot efficiently process?

5.1 NLP Evaluations

Refer to caption

(a) ARC (Challenge)

Refer to caption

(b) HellaSwag

Refer to caption

(c) LAMBADA (OpenAI)

Refer to caption

(d) OpenBookQA

Refer to caption

(e) ReCoRD

Refer to caption

(f) Winogrande

Figure 5: Zero-Shot Performance of RWKV on common language modeling evaluation benchmarks. Additional plots can be found in Appendix J.

To demonstrate that RWKV is competitive with traditional transformers at NLP tasks, we compare with similarly sized models trained for a similar number of tokens (Pythia Biderman et al. (2023b), OPT Zhang et al. (2022) and BLOOM Scao et al. (2022)). All RWKV models were trained for one epoch on the Pile (330B tokens), which is close but not identical to the amount of tokens the Pythia, OPT, and BLOOM models were trained for. Consequently, we compare our models on a FLOP-matched basis. We avoid comparing with model trained in the Chinchilla-optimal regime (Hoffmann et al., 2022) or the overtrained regime (Touvron et al., 2023) to ensure the most equitable comparison.

We report results on ARC (both Easy and Challenge) Clark et al. (2018), BoolQ Clark et al. (2019), COPA Roemmele et al. (2018), HeadQA Vilares and Gómez-Rodríguez (2019), HellaSwag Zellers et al. (2019), LAMBADA Paperno et al. (2016), OpenBookQA Mihaylov et al. (2018), PIQA Bisk et al. (2020), ReCoRD Zhang et al. (2018), SciQ Johannes Welbl Nelson F. Liu (2017), and Winogrande Zellers et al. (2020). Figure 1 shows the average results across all benchmarks. Some individual benchmarks are shown in Fig 5, with the rest in Appendix J.

Additionally, we carried out comparative studies on RWKV and ChatGPT / GPT-4, see Appendix L. They revealed that RWKV is very sensitive to prompt engineering. When the prompts were adjusted (re-ordered) from the ones used for GPT to more suitable for RWKV, the performance (F1) increased even from 44.2% to 74.8%. For sarcasm detection, RWKV outperformed ChatGPT, but was still slightly worse than the SOTA solution.

5.2 Extended Context Finetuning

Unlike transformers, RNNs do not have a pre-defined sequences length when they are created. However in order to efficient make use of compute we nevertheless need to preprocess the training data into contexts of the same length. We find that we are able to teach the model how to efficiently handle substantially larger batch sizes by finetuning with progressively increasing sequence length. Specifically, we first double the sequence length from 1024 to 2048 and finetune for 10B tokens from the original pretraining corpus, then we double again to 4096 for 100B tokens from the same corpus, and finally double to 8192 tokens for another 100B tokens from the same corpus. In Fig. 6 we show that increasing context length leads to lower test loss on the Pile, an indication that RWKV can make effective use of long contextual information.

Refer to caption

Figure 6: RWKV shows decreasing mean test loss as a function of context length on the Pile Gao et al. (2020)

5.3 Long Context Benchmarks

Additionally, we evaluate our model’s ability to handle very long sequences by comparing to state-of-the-art long sequence models on the Long-Range Arena (LRA) benchmark Tay et al. (2021). LRA is designed to assess the performance of models in handling lengthy context situations. It includes a collection of tasks with sequences ranging from 1,000 to 16,000 tokens, covering various types of data like text, natural language, synthetic images, and mathematical expressions. We apply RWKV on the LRA benchmark and the results are in Appendix J.2. The results show that RWKV performs second only to the S4 model in five datasets.

6 Inference Experiments

We benchmark inference requirements according to size and family. Specifically, we evaluate text generation speed and memory requirements on typical compute platforms including CPU (x86) and GPU (NVIDIA A100 80 GB). For all of our inference experiments we use float32 precision and the HuggingFace Transformers Wolf et al. (2020). We include all model parameters in the parameter count, including both embedding and non-embedding layers. Performance under different quantization setups is left to further work. See Appendix K for more results.

Refer to caption

Figure 7: Cumulative time on text generation for LLMs. Unlike transformers, RWKV exhibits linear scaling.

7 Future Work

There are several promising directions for future work on the RWKV architecture. Work can be done to increase model expressivity by enhancing the time-decay formulations and exploring initial model states while maintaining efficiency.

The RWKV computational efficiency can be further improved by applying a parallel scan in the w⁢k⁢vt𝑤𝑘subscript𝑣𝑡wkv_{t}italic_w italic_k italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT step to reduce the computational cost to O⁢(B⁢log⁡(T)⁢d)𝑂𝐵𝑇𝑑O(B\log(T)d)italic_O ( italic_B roman_log ( italic_T ) italic_d ).

The mechanisms used in RWKV can be applied to encoder-decoder architectures, potentially replacing the cross-attention mechanism. This could be applicable in seq2seq or multimodal settings, thereby enhancing efficiency during both training and inference.

RWKV’s state (or context) can be leveraged for interpretability, predictability in sequence data, and safety. Manipulating the hidden state could also guide behavior and allow greater customizability through prompt tuning.

The RWKV architecture is not perfect, and can be improved via many aspects, such as modifying the formulae or implementing larger internal states. Larger states can enhance the model’s memory to previous context and improve performance over various tasks.

8 Conclusions

We introduced RWKV, a new approach to RNN models exploiting the potential of time-based mixing components. RWKV introduces several key strategies that allow it to capture locality and long-range dependencies while addressing limitations of current architectures by: (1) replacing the quadratic QK attention with a scalar formulation at linear cost, (2) reformulating recurrence and sequential inductive biases to enable efficient training parallelization and efficient inference, and (3) enhancing training dynamics using custom initializations.

We benchmark the proposed architecture in a wide variety of NLP tasks and show comparable performance to SoTA with reduced cost. Further experiments on expressivity, interpretability, and scaling showcase the model capabilities and draw parallels in behavior between RWKV and other LLMs.

RWKV opens a new route for scalable and efficient architectures to model complex relationships in sequential data. While many alternatives to Transformers have been proposed with similar claims, ours is the first to back up those claims with pretrained models with tens of billions of parameters.

9 Limitations

While our proposed RWKV model has demonstrated promising results regarding training and memory efficiency during inference, some limitations should be acknowledged and addressed in future work.

First, the linear attention of RWKV leads to significant efficiency gains but still, it may also limit the model’s performance on tasks that require recalling minutiae information over very long contexts. This is due to the funneling of information through a single vector representation over many time steps, compared with the full information maintained by the quadratic attention of standard Transformers. In other words, the model’s recurrent architecture inherently limits its ability to “look back” at previous tokens, as opposed to traditional self-attention mechanisms. While learned time decay helps prevent the loss of information, it is mechanistically limited compared to full self-attention.

Another limitation of this work is the increased importance of prompt engineering in comparison to standard Transformer models. The linear attention mechanism used in RWKV limits the information from the prompt that will be carried over to the model’s continuation. As a result, carefully designed prompts may be even more crucial for the model to perform well on tasks.

The above RWKV property was confirmed by studies on prompt engineering presented in Appendix L. By changing the order of the information pieces, we were even able to almost double the RWKV performance for some tasks.

10 Ethics Statement

In this paper, we present a novel architecture for sequential data processing and prove its effectiveness by building a series of LLMs trained on publicly released pretraining data Gao et al. (2020); Biderman et al. (2022) and later fine-tuned on publicly available instructions Taori et al. (2023); Chaudhary (2023); Cheung (2023); Anand et al. (2023); Anonymous (2023); Yang (2023); Ji et al. (2023a, b).

As a novel architecture for sequential data, RWKV has the potential to improve sequence-based models across different applications ranging from natural language processing to biomedical data processing or climate modelling. Since the training code is released open source, RWKV contributes to the democratization of AI, levels the playing field, and empowers members of the Open Source community to inspect, study, and finetune RWKV in particular tasks. Moreover, it contributes to advancing the understanding of LLMs capabilities and limitations. A significant amount of work has been devoted to increasing the efficiency of RWKV training so as to minimize its cost and promote accessibility.

As LLMs trained on public data, RWKV’s lower inference cost compared to Transformer alternatives makes it more suitable for deployment in consumer and edge hardware, which is a step towards the democratization and distribution of LLMs to the general public, creating better privacy and ownership incentives. It also lowers the resource barrier to Chat assistants and text generation for small and/or underrepresented communities. PreTrained model weights for different sizes ranging from 0.1B to 14B parameters trained on multiple languages are released to increase ease of adoption and allow for the study of emergent phenomena.

On the other hand, with lower resource barriers, the spreading of AI-generated text might become more prevalent. Current RWKV LLMs may exhibit and/or reproduce biases and potentially harmful content present in the data used for training. Nonetheless, mitigation and finetuning strategies discussed for other, large Transformer models should be applicable to RWKV as well.

Acknowledgements

We thank StabilityAI for the compute used to train our models and for technical support in development of RWKV. We also thank the members of the RWKV and EleutherAI Discord servers for their help and work on further extending the applicability of RWKV to different domains.

References

Appendix A Author Contributions

All authors contributed to the drafting of this paper. Eric Alcaide and Quentin Anthony organized the paper and its experiments and were involved in all phases of the development process.

Model Design and Development

Bo Peng (lead), Matteo Grella, Xuzheng He, Haowen Hou, Jiaming Kong, Johan S. Wind

Model Training

Bo Peng

Scaling Laws Analysis

Stella Biderman, Bo Peng

Benchmark Evaluations

Stella Biderman (lead), Kranthi Kiran GV, Krishna Sri Ipsit Mantri, Atsushi Saito, Qihang Zhao, Peng Zhou, Rui-Jie Zhuåç

Long Context Experiments

Xingjian Du, Rui-Jie Zhu, Bolun Wang, Ruichong Zhang, Jian Zhu, Rui-Jie Zhu

Inference Speed Experiments

Samuel Arcadinho, Przemysław Kazienko, Qinghua Zhou

Information Flow Experiments

Huanqi Cao, Michael Chung, Matteo Grella, Ferdinand Mom, Zhenyuan Zhang

Chat Experiments

Jan Kocoń (lead), Przemysław Kazienko, Bartłomiej Koptyra, Hayden Lau, Xiangru Tang, Stanisław Woźniak, Zhenyuan Zhang

Ethics and Broader Impacts

Stella Biderman, Guangyu Song

Appendix B Author Contributions

Bo Peng

Original RWKV idea, original code, performance optimizations, original experiments, and trained RWKV models from 0.1B to 14B.

Eric Alcaide

Manuscript (initial draft sections 1, C; sections 3, 7 and 8; revision and proofreading; final version ). Figures (2, 3, 3, 8). Experiments section 6. Appendices E, K. Contributions to Appendix M.

Quentin Anthony

Manuscript (organization, initial draft sections 1, C, 2; revision and proofreading; final version).

Alon Albalak

Manuscript (abstract and sections 1, 9, and 7; proofreading and revision).

Samuel Arcadinho

Contributions to Figures 7, 13, and 14. Contributions to Appendix K.

Stella Biderman

Performed the scaling laws analysis and evaluated competitor models on benchmark tasks.

Huanqi Cao

Manuscript (contributions to 3.2 and 3.3; proofreading and revision). Experiments for Appendix I.

Xin Cheng

Manuscript (proofreading and revision). Contributions to Appendix M, J.

Michael Chung

Manuscript (contributions to section I; proofreading and revision).

Xingjian Du

Evaluation on Long Range Arena Benchmark (TBD until 5.31).

Matteo Grella

Manuscript (sections H, I, 8; contributions to sections 1, 7 and 9; proofreading and revision). Contributions to Appendix D.

Kranthi Kiran GV

Manuscript (sections C and 5; contributions to section 2; revision and proofreading). Tables K and K. Appendix 4.

Xuzheng He

Manuscript (contributions to section 2; proofreading and revision). Contributions to Figure8. Appendix I. Contributions to appendix H.

Haowen Hou

Figure 9. Appendix F.

Jiaju Lin

RWKV on LRA benchmarking

Przemysław Kazienko

Manuscript (proofreading and revision). Contributions to Section 6, 9, and Appendix L.

Jan Kocon

Manuscript (Section 1; proofreading and revision). Contributions to Appendix L.

Jiaming Kong

Manuscript (revision and proofreading). Appendix H.

Bartłomiej Koptyra

Manuscript (revision and proofreading) Contributions to Appendix L.

Hayden Lau

Manuscript (contributions to section 1 and 9; proofreading and revision). Contributions to Appendix M.

Krishna Sri Ipsit Mantri

Figure 12

Ferdinand Mom

Manuscript (contributions to section 1, C, 3.3, I; proofreading and revision). Contributions to Appendix D.

Atsushi Saito

Manuscript (sections 2 and 5; contributions to section C). Contributions to Appendix J

Guangyu Song

Manuscript (rewrote section 3; final version). Initial draft Ethics Statement).

Xiangru Tang

Manuscript (sections C and 2; contributions to abstract; revision and proofreading). Contributions to Appendix M.

Bolun Wang

Contributions to Tables 1.

Johan S. Wind

RWKV performance optimizations (CUDA), Contributions to Appendix 4.

Stanisław Woźniak

Contributions to Appendix L.

Ruichong Zhang

Manuscript (proofreading and revision); Contributions to Figure 6 and Appendix M.

Zhenyuan Zhang

Manuscript (revision and proofreading). Figure 3. Experiments Appendix I. Contributions to Appendices D and M.

Qihang Zhao

Manuscript (proofreading and revision). Contributions to Table 5.

Peng Zhou

Contributions to Tables 1 and Table 5.

Qinghua Zhou

Manuscript (Proofreading and revision of section 3; Add missing citations in 3.3). Revision of Figures 2 and 12.

Jian Zhu

Manuscript (section C; proofreading and revision). Figures 3 and 6.

Rui-Jie Zhu

Tables 1 and 5. Experiments for table 5.

Recently, a number of techniques have been proposed to address the limitations of transformers.

Optimizing Attention Mechanism

Many transformer variants (“x-formers”) have been introduced to reduce the complexity of transformers Tay et al. (2022), including sparse attention Beltagy et al. (2020); Kitaev et al. (2020); Guo et al. (2022), approximating the full attention matrix Wang et al. (2020); Ma et al. (2021); Choromanski et al. (2020), combining chunked attention with gating Ma et al. (2023) and other efficient methods Katharopoulos et al. (2020); Jaegle et al. (2021).

Some recent works like FlashAttention Dao et al. (2022a) and others Rabe and Staats (2022); Jang et al. (2019) share similarities with RWKV’s chunked computation scheme. Despite being memory-efficient, their time complexity remains quadratic or contains chunk size as a hidden factor. In contrast, RWKV achieves better space and time complexity during inference by formulating a linear attention as an RNN.

Attention Free Models

Another line of research replaces the attention mechanism with other modules to scale to long sequences. MLP-Mixer and others Tolstikhin et al. (2021); Liu et al. (2021) propose replacing attention by Multi-Layer Perceptrons (MLPs) in computer vision tasks. The Attention Free Transformer (AFT) (Zhai et al., 2021) and HrrFormer (Alam et al., 2023) replaces dot-product self-attention with a computationally efficient alternative. None of these models have been successfully scaled to the point where drawing comparisons with transformer-based large language models makes sense.

There has also been substantial research into state space models (SSM) (Gu et al., 2021) and its variants (Dao et al., 2022b; Gupta et al., 2022; Poli et al., 2023). In contrast to the preceding models, SSM and its successors have shown substantial progress towards efficient scaling. Simultaneously with this work, Poli et al. (2023) train SSM-based models with 125 million and 355 million parameters and show that the performance is on-par with a transformer that uses a mix of local and global attention (Black et al., 2021).

Advances in RNNs

Inspired by the success of transformers, RNN-style Hochreiter and Schmidhuber (1997); Chung et al. (2014) recursive components have also been modified to increase context length, such as the Recurrent Memory Transformer Bulatov et al. (2022, 2023) and Linear Recurrent Units Orvieto et al. (2023). Most similar to our work, the Quasi-Recurrent neural network (QRNN) Bradbury et al. (2017) uses both convolutional layers and recurrent pooling functions across timesteps and channels. While QRNN utilizes convolutional filters with fixed sizes, RWKV employs a time-mixing module as an attention mechanism with time-decaying factors. Different from the element-wise pooling in QRNN, RWKV includes a parametrized channel-mixing module that is parallelizable.

Appendix D Time-Mixing Block as an RNN Cell

As stated in 3.3, the RWKV time-mixing block can be formulated as an RNN, as the W⁢K⁢V𝑊𝐾𝑉WKVitalic_W italic_K italic_V computation can be written in such a recursive form:

a0,b0subscript𝑎0subscript𝑏0\displaystyle a_{0},b_{0}italic_a start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT , italic_b start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT =0,absent0\displaystyle=0,= 0 , (19)
w⁢k⁢vt𝑤𝑘subscript𝑣𝑡\displaystyle wkv_{t}italic_w italic_k italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =at−1+eu+kt⊙vtbt−1+eu+kt,absentsubscript𝑎𝑡1direct-productsuperscript𝑒𝑢subscript𝑘𝑡subscript𝑣𝑡subscript𝑏𝑡1superscript𝑒𝑢subscript𝑘𝑡\displaystyle=\frac{a_{t-1}+e^{u+k_{t}}\odot v_{t}}{b_{t-1}+e^{u+k_{t}}},= divide start_ARG italic_a start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_u + italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_POSTSUPERSCRIPT ⊙ italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_ARG start_ARG italic_b start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_u + italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_POSTSUPERSCRIPT end_ARG , (20)
atsubscript𝑎𝑡\displaystyle a_{t}italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =e−w⊙at−1+ekt⊙vt,absentdirect-productsuperscript𝑒𝑤subscript𝑎𝑡1direct-productsuperscript𝑒subscript𝑘𝑡subscript𝑣𝑡\displaystyle=e^{-w}\odot a_{t-1}+e^{k_{t}}\odot v_{t},= italic_e start_POSTSUPERSCRIPT - italic_w end_POSTSUPERSCRIPT ⊙ italic_a start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_POSTSUPERSCRIPT ⊙ italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , (21)
btsubscript𝑏𝑡\displaystyle b_{t}italic_b start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =e−w⊙bt−1+ekt.absentdirect-productsuperscript𝑒𝑤subscript𝑏𝑡1superscript𝑒subscript𝑘𝑡\displaystyle=e^{-w}\odot b_{t-1}+e^{k_{t}}.= italic_e start_POSTSUPERSCRIPT - italic_w end_POSTSUPERSCRIPT ⊙ italic_b start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_POSTSUPERSCRIPT . (22)

Refer to caption

Figure 8: RWKV time-mixing block formulated as an RNN cell. Color codes: yellow (μ𝜇\muitalic_μ) denotes the token shift, red (1) denotes the denominator, blue (2) denotes the numerator, and pink (3) denotes the fraction computations in 16. hℎhitalic_h denotes the numerator-denominator tuple.

The dataflow of the RNN-like time-mixing is shown in Fig. 8, where the hidden states hℎhitalic_h is the numerator-denominator tuple (a,b)𝑎𝑏(a,b)( italic_a , italic_b ). To avoid overflow in calculating ektsuperscript𝑒subscript𝑘𝑡e^{k_{t}}italic_e start_POSTSUPERSCRIPT italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_POSTSUPERSCRIPT, a numerical trick is used in the official implementation. Noticing that a1=ek1⊙v1subscript𝑎1direct-productsuperscript𝑒subscript𝑘1subscript𝑣1a_{1}=e^{k_{1}}\odot v_{1}italic_a start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT = italic_e start_POSTSUPERSCRIPT italic_k start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT end_POSTSUPERSCRIPT ⊙ italic_v start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT and b1=ek1subscript𝑏1superscript𝑒subscript𝑘1b_{1}=e^{k_{1}}italic_b start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT = italic_e start_POSTSUPERSCRIPT italic_k start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT end_POSTSUPERSCRIPT, we set a1′=v1,b1′=1,p1=k1formulae-sequencesubscriptsuperscript𝑎′1subscript𝑣1formulae-sequencesubscriptsuperscript𝑏′11subscript𝑝1subscript𝑘1a^{\prime}_{1}=v_{1},b^{\prime}_{1}=1,p_{1}=k_{1}italic_a start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT = italic_v start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT , italic_b start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT = 1 , italic_p start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT = italic_k start_POSTSUBSCRIPT 1 end_POSTSUBSCRIPT, where ptsubscript𝑝𝑡p_{t}italic_p start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT stores the shared exponents of atsubscript𝑎𝑡a_{t}italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT and btsubscript𝑏𝑡b_{t}italic_b start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT. Now the above recursion can be converted into a numerical safe version, for each time step t>1𝑡1t>1italic_t > 1:

q𝑞\displaystyle qitalic_q :=max⁡(pt−1,u+kt),assignabsentsubscript𝑝𝑡1𝑢subscript𝑘𝑡\displaystyle:=\max(p_{t-1},u+k_{t}),:= roman_max ( italic_p start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT , italic_u + italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) , (23)
w⁢k⁢vt𝑤𝑘subscript𝑣𝑡\displaystyle wkv_{t}italic_w italic_k italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =ept−1−q⊙at−1′+eu+kt−q⊙vtept−1−q⊙bt−1′+eu+kt−q.absentdirect-productsuperscript𝑒subscript𝑝𝑡1𝑞subscriptsuperscript𝑎′𝑡1direct-productsuperscript𝑒𝑢subscript𝑘𝑡𝑞subscript𝑣𝑡direct-productsuperscript𝑒subscript𝑝𝑡1𝑞subscriptsuperscript𝑏′𝑡1superscript𝑒𝑢subscript𝑘𝑡𝑞\displaystyle=\frac{e^{p_{t-1}-q}\odot a^{\prime}_{t-1}+e^{u+k_{t}-q}\odot v_{% t}}{e^{p_{t-1}-q}\odot b^{\prime}_{t-1}+e^{u+k_{t}-q}}.= divide start_ARG italic_e start_POSTSUPERSCRIPT italic_p start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT - italic_q end_POSTSUPERSCRIPT ⊙ italic_a start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_u + italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT - italic_q end_POSTSUPERSCRIPT ⊙ italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_ARG start_ARG italic_e start_POSTSUPERSCRIPT italic_p start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT - italic_q end_POSTSUPERSCRIPT ⊙ italic_b start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_u + italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT - italic_q end_POSTSUPERSCRIPT end_ARG . (24)

The update to at′,bt′subscriptsuperscript𝑎′𝑡subscriptsuperscript𝑏′𝑡a^{\prime}_{t},b^{\prime}_{t}italic_a start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , italic_b start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT, and their shared exponent is also carried out in a similar fashion:

q′superscript𝑞′\displaystyle q^{\prime}italic_q start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT :=max⁡(pt−1−w,kt),assignabsentsubscript𝑝𝑡1𝑤subscript𝑘𝑡\displaystyle:=\max(p_{t-1}-w,k_{t}),:= roman_max ( italic_p start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT - italic_w , italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) , (25)
at′subscriptsuperscript𝑎′𝑡\displaystyle a^{\prime}_{t}italic_a start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =ept−1−w−q′⊙at−1′+ekt−q′⊙vt,absentdirect-productsuperscript𝑒subscript𝑝𝑡1𝑤superscript𝑞′subscriptsuperscript𝑎′𝑡1direct-productsuperscript𝑒subscript𝑘𝑡superscript𝑞′subscript𝑣𝑡\displaystyle=e^{p_{t-1}-w-q^{\prime}}\odot a^{\prime}_{t-1}+e^{k_{t}-q^{% \prime}}\odot v_{t},= italic_e start_POSTSUPERSCRIPT italic_p start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT - italic_w - italic_q start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT end_POSTSUPERSCRIPT ⊙ italic_a start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT - italic_q start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT end_POSTSUPERSCRIPT ⊙ italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , (26)
bt′subscriptsuperscript𝑏′𝑡\displaystyle b^{\prime}_{t}italic_b start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =ept−1−w−q′⊙bt−1′+ekt−q′,absentdirect-productsuperscript𝑒subscript𝑝𝑡1𝑤superscript𝑞′subscriptsuperscript𝑏′𝑡1superscript𝑒subscript𝑘𝑡superscript𝑞′\displaystyle=e^{p_{t-1}-w-q^{\prime}}\odot b^{\prime}_{t-1}+e^{k_{t}-q^{% \prime}},= italic_e start_POSTSUPERSCRIPT italic_p start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT - italic_w - italic_q start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT end_POSTSUPERSCRIPT ⊙ italic_b start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT - italic_q start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT end_POSTSUPERSCRIPT , (27)
ptsubscript𝑝𝑡\displaystyle p_{t}italic_p start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =q′.absentsuperscript𝑞′\displaystyle=q^{\prime}.= italic_q start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT . (28)

The RWKV model has an internal state that stores some previous information. In each layer, the internal state consists five parts, each of which is a vector with D𝐷Ditalic_D numbers, where D𝐷Ditalic_D is the model dimension. The five parts are:

Which yields a total size of 5⁢D⁢L5𝐷𝐿5DL5 italic_D italic_L parameters. It is worth noting that in an algebraic context with infinite precision, the helper state ptsubscript𝑝𝑡p_{t}italic_p start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT can be ignored, and the W⁢K⁢V𝑊𝐾𝑉WKVitalic_W italic_K italic_V numerator and denominator can be computed directly using equations (21) and (22), reducing the size of the internal state to 4⁢D⁢L4𝐷𝐿4DL4 italic_D italic_L.

Appendix E Parameter initializations

We describe the specific parameter initializations below and motivate the design choices. Parameters belonging to residual blocks are often adjusted by layer depth and total number of layers. Let ##\## denote the vocabulary size, s𝑠sitalic_s denote the embedding dimension, d𝑑ditalic_d denote the hidden size (we use d=4⁢s𝑑4𝑠d=4sitalic_d = 4 italic_s), L𝐿Litalic_L the number of layers, l𝑙litalic_l the layer index (from 0 to L−1𝐿1L-1italic_L - 1), we use the following initializations:

Appendix F Small Init Embedding

This section presents the experimental validation of small initialization embedding. The experimental setup is as follows. In the baseline configuration, the parameters are initialized using a normal distribution with a mean of 0.0 and a standard deviation of 0.02, which is a commonly used initialization method in models like BERT and GPT. On the other hand, in the small initialization of the embedding (small init emb) experiment, the parameters are initialized using a uniform distribution with a range of 1e-4, which is slightly different from RWKV where a normal distribution with a standard deviation of 1e-4 is used. However, this difference is negligible and does not affect our conclusions. The experiments were conducted with a batch size of 400. As depicted in Figure 9, the loss curve for the small init emb exhibits a faster rate of decrease and convergence compared to the traditional initialization using a normal distribution.

Refer to caption

Figure 9: Effect of small initialization embedding.

Appendix G Hyperparameters

Table 3: Hyperparameters for our learning rate (LR) schedule of the pretrained models.

To train the models mentioned, we use ϵ=(0.9,0.99)italic-ϵ0.90.99\epsilon=(0.9,0.99)italic_ϵ = ( 0.9 , 0.99 ) without weight decay for the Adam optimizer, and switch batch size dynamically between 128 or 256 sequences, each of 1024 tokens. We further organize the training into multiple mini-epochs, each of 40320 samples, to guide our learning rate schedule. The training process takes 8043 mini-epochs to make one pass over the Pile. The initial warming up mini-epochs have a constant learning rate of “Init LR”. After the warming up mini-epochs, the learning rate exponentially decays until in the last mini-epoch, in which the model finishes training on the entire Pile, the learning rate arrives at the “End LR”. The related hyperparameters are shown in Table 3.

Appendix H Gradient Stability in RWKV

In this section, we present a mathematical description of the gradient stability property in RWKV, focusing specifically on the time-mixing block. By gradient stability we mean that if the inputs xtsubscript𝑥𝑡x_{t}italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT are bounded and the model parameters are fixed, then the gradients with respect to Wksubscript𝑊𝑘W_{k}italic_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT and Wvsubscript𝑊𝑣W_{v}italic_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT are uniformly bounded for all T𝑇Titalic_T (thus not exploding). Consequently, we can control the amount each xtsubscript𝑥𝑡x_{t}italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT contributes to the gradient at T𝑇Titalic_T in a naturally decaying fashion by the weight decay mechanism w𝑤witalic_w (thus not vanishing unless desired).

First, we make the simplification that there are no token shifts, this will not affect the final conclusion. In this scenario, w⁢k⁢vT𝑤𝑘subscript𝑣𝑇wkv_{T}italic_w italic_k italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT can be written as

w⁢k⁢vT=∑t=1TKte⊙vt∑t=1TKte=E⁢(vt)=S⁢(vt)S⁢(1),𝑤𝑘subscript𝑣𝑇superscriptsubscript𝑡1𝑇direct-productsubscriptsuperscript𝐾𝑒𝑡subscript𝑣𝑡superscriptsubscript𝑡1𝑇subscriptsuperscript𝐾𝑒𝑡Esubscript𝑣𝑡Ssubscript𝑣𝑡S1\displaystyle wkv_{T}=\frac{\sum_{t=1}^{T}K^{e}_{t}\odot v_{t}}{\sum_{t=1}^{T}% K^{e}_{t}}=\mathrm{E}(v_{t})=\frac{\mathrm{S}(v_{t})}{\mathrm{S}(1)},italic_w italic_k italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT = divide start_ARG ∑ start_POSTSUBSCRIPT italic_t = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_T end_POSTSUPERSCRIPT italic_K start_POSTSUPERSCRIPT italic_e end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ⊙ italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_ARG start_ARG ∑ start_POSTSUBSCRIPT italic_t = 1 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_T end_POSTSUPERSCRIPT italic_K start_POSTSUPERSCRIPT italic_e end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT end_ARG = roman_E ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) = divide start_ARG roman_S ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) end_ARG start_ARG roman_S ( 1 ) end_ARG , (29)

where

vt=Wv⁢xt,∂(vt)i∂(Wv)i,j=(xt)j,formulae-sequencesubscript𝑣𝑡subscript𝑊𝑣subscript𝑥𝑡subscriptsubscript𝑣𝑡𝑖subscriptsubscript𝑊𝑣𝑖𝑗subscriptsubscript𝑥𝑡𝑗v_{t}=W_{v}x_{t},\quad\frac{\partial(v_{t})_{i}}{\partial(W_{v})_{i,j}}=(x_{t}% )_{j},italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT = italic_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT , divide start_ARG ∂ ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG start_ARG ∂ ( italic_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT end_ARG = ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_j end_POSTSUBSCRIPT ,
Kte=eWk⁢xt+wT,t,∂(Kte)i∂(Wk)i,j=(xt)j⁢(Kte)i,formulae-sequencesubscriptsuperscript𝐾𝑒𝑡superscript𝑒subscript𝑊𝑘subscript𝑥𝑡subscript𝑤𝑇𝑡subscriptsubscriptsuperscript𝐾𝑒𝑡𝑖subscriptsubscript𝑊𝑘𝑖𝑗subscriptsubscript𝑥𝑡𝑗subscriptsubscriptsuperscript𝐾𝑒𝑡𝑖K^{e}_{t}=e^{W_{k}x_{t}+w_{T,t}},\quad\frac{\partial(K^{e}_{t})_{i}}{\partial(% W_{k})_{i,j}}=(x_{t})_{j}(K^{e}_{t})_{i},italic_K start_POSTSUPERSCRIPT italic_e end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT = italic_e start_POSTSUPERSCRIPT italic_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT + italic_w start_POSTSUBSCRIPT italic_T , italic_t end_POSTSUBSCRIPT end_POSTSUPERSCRIPT , divide start_ARG ∂ ( italic_K start_POSTSUPERSCRIPT italic_e end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG start_ARG ∂ ( italic_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT end_ARG = ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_j end_POSTSUBSCRIPT ( italic_K start_POSTSUPERSCRIPT italic_e end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ,

and S⁢(⋅)S⋅\mathrm{S}(\cdot)roman_S ( ⋅ ) and E⁢(⋅)E⋅\mathrm{E}(\cdot)roman_E ( ⋅ ) are shorthand for denoting sums and averages over weights Ktesubscriptsuperscript𝐾𝑒𝑡K^{e}_{t}italic_K start_POSTSUPERSCRIPT italic_e end_POSTSUPERSCRIPT start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT.

The loss function at position T𝑇Titalic_T can be written as

LT=l⁢(f⁢(w⁢k⁢vT),yT).subscript𝐿𝑇𝑙𝑓𝑤𝑘subscript𝑣𝑇subscript𝑦𝑇\displaystyle L_{T}=l(f(wkv_{T}),y_{T}).italic_L start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT = italic_l ( italic_f ( italic_w italic_k italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT ) , italic_y start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT ) . (30)

Because w⁢k⁢vT𝑤𝑘subscript𝑣𝑇wkv_{T}italic_w italic_k italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT relates to (Wk)i,jsubscriptsubscript𝑊𝑘𝑖𝑗(W_{k})_{i,j}( italic_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT and (Wv)i,jsubscriptsubscript𝑊𝑣𝑖𝑗(W_{v})_{i,j}( italic_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT only through the i𝑖iitalic_i-th channel (w⁢k⁢vT)isubscript𝑤𝑘subscript𝑣𝑇𝑖(wkv_{T})_{i}( italic_w italic_k italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT, we have

∂LT∂(Wv)i,j=∂LT∂(w⁢k⁢vT)i⁢∂(w⁢k⁢vT)i∂(Wv)i,j.subscript𝐿𝑇subscriptsubscript𝑊𝑣𝑖𝑗subscript𝐿𝑇subscript𝑤𝑘subscript𝑣𝑇𝑖subscript𝑤𝑘subscript𝑣𝑇𝑖subscriptsubscript𝑊𝑣𝑖𝑗\displaystyle\frac{\partial L_{T}}{\partial(W_{v})_{i,j}}=\frac{\partial L_{T}% }{\partial(wkv_{T})_{i}}\frac{\partial(wkv_{T})_{i}}{\partial(W_{v})_{i,j}}.divide start_ARG ∂ italic_L start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT end_ARG start_ARG ∂ ( italic_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT end_ARG = divide start_ARG ∂ italic_L start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT end_ARG start_ARG ∂ ( italic_w italic_k italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG divide start_ARG ∂ ( italic_w italic_k italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG start_ARG ∂ ( italic_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT end_ARG . (31)

The first part of the above equation contains trivial operations like output layers, and other layers of time-mixing, which can be proven inductively. The second part of the above equation can be bounded as

| |∂(w⁢k⁢vT)i∂(Wv)i,j|subscript𝑤𝑘subscript𝑣𝑇𝑖subscriptsubscript𝑊𝑣𝑖𝑗\displaystyle\left|\frac{\partial(wkv_{T})_{i}}{\partial(W_{v})_{i,j}}\right|| divide start_ARG ∂ ( italic_w italic_k italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG start_ARG ∂ ( italic_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT end_ARG | | =|∂Ei⁢[(vt)i]∂(Wv)i,j|absentsubscriptE𝑖delimited-[]subscriptsubscript𝑣𝑡𝑖subscriptsubscript𝑊𝑣𝑖𝑗\displaystyle=\left|\frac{\partial\mathrm{E}_{i}[(v_{t})_{i}]}{\partial(W_{v})% _{i,j}}\right|= | divide start_ARG ∂ roman_E start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT [ ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ] end_ARG start_ARG ∂ ( italic_W start_POSTSUBSCRIPT italic_v end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT end_ARG | | | | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | - | ---- | | =|Ei⁢[(xt)j]|≤maxt⁡|(xt)j|,absentsubscriptE𝑖delimited-[]subscriptsubscript𝑥𝑡𝑗subscript𝑡subscriptsubscript𝑥𝑡𝑗\displaystyle=|\mathrm{E}_{i}[(x_{t})_{j}]|\leq\max_{t}|(x_{t})_{j}|,= | roman_E start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT [ ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_j end_POSTSUBSCRIPT ] | ≤ roman_max start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT | ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_j end_POSTSUBSCRIPT | , | (32) |

which is irrelevant to T𝑇Titalic_T. Similarly,

∂(w⁢k⁢vT)i∂(Wk)i,jsubscript𝑤𝑘subscript𝑣𝑇𝑖subscriptsubscript𝑊𝑘𝑖𝑗\displaystyle\frac{\partial(wkv_{T})_{i}}{\partial(W_{k})_{i,j}}divide start_ARG ∂ ( italic_w italic_k italic_v start_POSTSUBSCRIPT italic_T end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT end_ARG start_ARG ∂ ( italic_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT end_ARG =∂Si⁢[(vt)i]Si⁢(1)/∂(Wk)i,jabsentsubscriptS𝑖delimited-[]subscriptsubscript𝑣𝑡𝑖subscriptS𝑖1subscriptsubscript𝑊𝑘𝑖𝑗\displaystyle=\partial\frac{\mathrm{S}_{i}[(v_{t})_{i}]}{\mathrm{S}_{i}(1)}/% \partial(W_{k})_{i,j}= ∂ divide start_ARG roman_S start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT [ ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ] end_ARG start_ARG roman_S start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( 1 ) end_ARG / ∂ ( italic_W start_POSTSUBSCRIPT italic_k end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i , italic_j end_POSTSUBSCRIPT
=Si⁢[(xt)j⁢(vt)i]Si⁢(1)−Si⁢[(xt)j]⁢Si⁢[(vt)i]Si⁢(1)2absentsubscriptS𝑖delimited-[]subscriptsubscript𝑥𝑡𝑗subscriptsubscript𝑣𝑡𝑖subscriptS𝑖1subscriptS𝑖delimited-[]subscriptsubscript𝑥𝑡𝑗subscriptS𝑖delimited-[]subscriptsubscript𝑣𝑡𝑖subscriptS𝑖superscript12\displaystyle=\frac{\mathrm{S}_{i}[(x_{t})_{j}(v_{t})_{i}]}{\mathrm{S}_{i}(1)}% -\frac{\mathrm{S}_{i}[(x_{t})_{j}]\mathrm{S}_{i}[(v_{t})_{i}]}{\mathrm{S}_{i}(% 1)^{2}}= divide start_ARG roman_S start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT [ ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_j end_POSTSUBSCRIPT ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ] end_ARG start_ARG roman_S start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( 1 ) end_ARG - divide start_ARG roman_S start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT [ ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_j end_POSTSUBSCRIPT ] roman_S start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT [ ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ] end_ARG start_ARG roman_S start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( 1 ) start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT end_ARG
=Ei⁢[(xt)j⁢(vt)i]−Ei⁢[(xt)j]⁢Ei⁢[(vt)i]absentsubscriptE𝑖delimited-[]subscriptsubscript𝑥𝑡𝑗subscriptsubscript𝑣𝑡𝑖subscriptE𝑖delimited-[]subscriptsubscript𝑥𝑡𝑗subscriptE𝑖delimited-[]subscriptsubscript𝑣𝑡𝑖\displaystyle=\mathrm{E}_{i}[(x_{t})_{j}(v_{t})_{i}]-\mathrm{E}_{i}[(x_{t})_{j% }]\mathrm{E}_{i}[(v_{t})_{i}]= roman_E start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT [ ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_j end_POSTSUBSCRIPT ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ] - roman_E start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT [ ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_j end_POSTSUBSCRIPT ] roman_E start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT [ ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ]
=covi⁢((xt)j,(vt)i)absentsubscriptcov𝑖subscriptsubscript𝑥𝑡𝑗subscriptsubscript𝑣𝑡𝑖\displaystyle=\mathrm{cov}_{i}((x_{t})_{j},(v_{t})_{i})= roman_cov start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ( ( italic_x start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_j end_POSTSUBSCRIPT , ( italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT ) start_POSTSUBSCRIPT italic_i end_POSTSUBSCRIPT ) (33)

can also be bounded. Note that w⁢k⁢v𝑤𝑘𝑣wkvitalic_w italic_k italic_v’s softmax operation contains at least two non-zero terms (u𝑢uitalic_u and w𝑤witalic_w), so the above “covariance” will not degenerate into 0.

Appendix I Model Behavior Visualization

Refer to caption

Figure 10: Model behavior visualizations of RWKV.

The right plot illustrates the time decays (e−wsuperscript𝑒𝑤e^{-w}italic_e start_POSTSUPERSCRIPT - italic_w end_POSTSUPERSCRIPT) in each layer of the RWKV-169M model, sorted along the channel axis. Notably, several decays in the last layers are very close or equal to one, implying that certain information is preserved and propagated throughout the model’s temporal context. Meanwhile, many decays in the initial layer are close to zero, which corresponds to local operations in w⁢k⁢v𝑤𝑘𝑣wkvitalic_w italic_k italic_v (16), likely to be associated with tasks such as text parsing or lexical analysis. (Note that the local operations in w⁢k⁢v𝑤𝑘𝑣wkvitalic_w italic_k italic_v are due to the extra parameter u𝑢uitalic_u, when e−wsuperscript𝑒𝑤e^{-w}italic_e start_POSTSUPERSCRIPT - italic_w end_POSTSUPERSCRIPT is degenerated into 0.) These patterns of time decays are partly learned, but also come from parameter initialization as it speeds up training.

The plot below shows the information retrieval and propagation path in the RWKV-430M model. The experiment follows the causal trace method introduced by Meng et al. (2022), where we

    1. Run the model once, and record all states and activation of each layer during the computation;
    1. Corrupt the input embeddings of the subject using noise (“The Eiffel Tower” in this example);
    1. Restore the states and activation of a certain layer at a certain token during the computation, and record the log-probability of the model outputting the correct answer (“Paris”).

Unlike transformers, RWKV relies on the recursive propagation of information in the time dimension. In this case, the fact that the Eiffel Tower is located in Paris is retrieved in layer 4 just after the model sees “The Eiffel”. It is then passed down to the subsequent layers. In layer 20, mostly, the information is propagated through time until reaching where it is needed. Finally, at the token “of”, it is passed down to the last layer for outputting the answer.

Refer to caption

Figure 11: Model behavior visualizations of the RWKV model.

Appendix J Additional Evaluations

J.1 Further details on NLP tasks

We evaluate on the following tasks:

ARC Clark et al. (2018)

A dataset designed for multiple-choice question answering, encompassing science exam questions ranging from third grade to ninth grade. It has Easy and Challenge subsets that we report results on separately.

BoolQ Clark et al. (2019)

A binary yes/no question answering benchmark.

COPA Roemmele et al. (2018)

A dataset to evaluate achievement in open-domain commonsense causal reasoning.

HeadQA Vilares and Gómez-Rodríguez (2019)

A benchmark consisting of graduate-level questions encompassing various fields such as medicine, nursing, biology, chemistry, psychology, and pharmacology.

HellaSwag

Zellers et al. (2019) A novel benchmark for commonsense Natural Language Inference (NLI) which is build by adversarial filtering against transformer models.

LAMBADA Paperno et al. (2016)

A benchmark dataset that evaluates the model’s contextual reasoning and language comprehension abilities by presenting context-target pairs, where the objective is to predict the most probable target token. We follow standard practice and use the untokenized version created by OpenAI (Brown et al., 2020).

OpenBookQA Mihaylov et al. (2018)

A QA dataset to evaluate human comprehension of a subject by incorporating open book facts, scientific knowledge, and perceptual common sense, drawing inspiration from open book exams.

PIQA Bisk et al. (2020)

A benchmark for the task of physical common sense reasoning, which consists of a binary choice task that can be better understood as a set of two pairs, namely (Goal, Solution).

ReCoRD Zhang et al. (2018)

A benchmark for evaluating commonsense reasoning in reading comprehension by generating queries from CNN/Daily Mail news articles and requiring text span answers from corresponding summarizing passages.

SciQ Johannes Welbl Nelson F. Liu (2017)

A multiple-choice QA dataset which was created using an innovative approach to gather well-crafted multiple-choice questions that are focused on a specific domain.

Winogrande Zellers et al. (2020)

A dataset designed to evaluate the acquisition of common sense reasoning by neural language models, aiming to determine whether we are accurately assessing the true capabilities of machine common sense.

Refer to caption

(a) ARC (Challenge)

Refer to caption

(b) ARC (Easy)

Refer to caption

(c) BoolQ

Refer to caption

(d) COPA

Refer to caption

(e) HeadQA

Refer to caption

(f) HellaSwag

Refer to caption

(g) LAMBADA (OpenAI)

Refer to caption

(h) OpenBookQA

Refer to caption

(i) PiQA

Refer to caption

(j) ReCoRD

Refer to caption

(k) SciQ

Refer to caption

(l) Winogrande

Figure 12: Zero-Shot Performance of RWKV on common language modeling evaluation benchmarks.

J.2 Evaluation on Long Range Arena

The Long-Range Arena (LRA) benchmark (Tay et al., 2021) is designed to assess the performance of models in handling lengthy context situations. It includes a collection of tasks with sequences ranging from 1,000 to 16,000 tokens, covering various types of data like text, natural language, synthetic images, and mathematical expressions. We apply RWKV on the LRA benchmark and the report results are in Table 4. Other models’ performances are directly cited from Gu et al. (2022); Alam et al. (2023).

Table 4: Evaluation on Long Range Arena. Other models reported in the literature (Gu et al., 2022; Alam et al., 2023). Bolded values are the best.

Model ListOps Text Retrieval Image Pathfinder Path-X Avg
Transformer 36.37 64.27 57.46 42.44 71.40 53.66
Reformer 37.27 56.10 53.40 38.07 68.50 50.56
BigBird 36.05 64.02 59.29 40.83 74.87 54.17
Linear Trans. 16.13 65.90 53.09 42.34 75.30 50.46
Performer 18.01 65.40 53.82 42.77 77.05 51.18
FNet 35.33 65.11 59.61 38.67 77.80 54.42
Nyströmformer 37.15 65.52 79.56 41.58 70.94 57.46
Luna-256 37.25 64.57 79.29 47.38 77.72 59.37
Hrrformer 39.98 65.38 76.15 50.45 72.17 60.83
S4 59.60 86.82 90.90 88.65 94.20 96.35 86.09
RWKV 55.88 86.04 88.34 70.53 58.42 72.07

The results show that RWKV performs second only to the S4 model in five datasets. While RWKV substantially underpreforms S4 on Image, Pathfinder, and Path-X, on the problems related to natural language and computer code processing RWKV performs on par with S4 or nearly so.

J.3 Enwik8 Perplexity

We also evaluate our model in terms of perplexity on the Enwik8 dataset. Baseline comparisons are made with Reformer (Kitaev et al., 2020), Synthesizer (Tay et al., 2020) (the best performing dense version), Linear Transformer (Katharopoulos et al., 2020), Performer (Choromanski et al., 2020). L,d,𝐿𝑑L,d,italic_L , italic_d , and T𝑇Titalic_T denote the number of blocks (network depth), dimension of features, and sequence length, respectively. Both Linear Transformer and Performer are implemented with customized CUDA kernels (github.com/idiap/fast-transformers), and all other models are implemented in native Pytorch. 11{}^{1}start_FLOATSUPERSCRIPT 1 end_FLOATSUPERSCRIPT No weight decay nor dropout was used. 22{}^{2}start_FLOATSUPERSCRIPT 2 end_FLOATSUPERSCRIPT Trained with AdamW and weight decay set to 0.1, dropout of 0.1, batch size of 16, and initial learning rate of 6e-4.

Method L d T Train bpc Test bpc Time Complexity Space Complexity
Transformer 12 512 1024 0.977 1.137 O⁢(T2⁢d)𝑂superscript𝑇2𝑑O(T^{2}d)italic_O ( italic_T start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT italic_d ) O⁢(T2+T⁢d)𝑂superscript𝑇2𝑇𝑑O(T^{2}+Td)italic_O ( italic_T start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT + italic_T italic_d )
Transformer 24 256 1024 1.039 1.130 O⁢(T2⁢d)𝑂superscript𝑇2𝑑O(T^{2}d)italic_O ( italic_T start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT italic_d ) O⁢(T2+T⁢d)𝑂superscript𝑇2𝑇𝑑O(T^{2}+Td)italic_O ( italic_T start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT + italic_T italic_d )
Reformer 12 512 1024 1.040 1.195 O⁢(T⁢log⁡T⁢d)𝑂𝑇𝑇𝑑O(T\log{T}d)italic_O ( italic_T roman_log italic_T italic_d ) O⁢(T⁢log⁡T+T⁢d)𝑂𝑇𝑇𝑇𝑑O(T\log{T}+Td)italic_O ( italic_T roman_log italic_T + italic_T italic_d )
Synthesizer 12 512 1024 0.994 1.298 O⁢(T2⁢d)𝑂superscript𝑇2𝑑O(T^{2}d)italic_O ( italic_T start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT italic_d ) O⁢(T2+T⁢d)𝑂superscript𝑇2𝑇𝑑O(T^{2}+Td)italic_O ( italic_T start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT + italic_T italic_d )
Linear Transformer 12 512 1024 0.981 1.207 O⁢(T⁢d2)𝑂𝑇superscript𝑑2O(Td^{2})italic_O ( italic_T italic_d start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT ) O⁢(T⁢d+d2)𝑂𝑇𝑑superscript𝑑2O(Td+d^{2})italic_O ( italic_T italic_d + italic_d start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT )
Performer 12 512 1024 1.002 1.199 O⁢(T⁢d2⁢log⁡d)𝑂𝑇superscript𝑑2𝑑O(Td^{2}\log{d})italic_O ( italic_T italic_d start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT roman_log italic_d ) O⁢(T⁢d⁢log⁡d+d2⁢log⁡d)𝑂𝑇𝑑𝑑superscript𝑑2𝑑O(Td\log{d}+d^{2}\log{d})italic_O ( italic_T italic_d roman_log italic_d + italic_d start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT roman_log italic_d )
AFT-simple 12 512 1024 1.046 1.209 O⁢(T⁢d)𝑂𝑇𝑑O(Td)italic_O ( italic_T italic_d ) O⁢(T⁢d)𝑂𝑇𝑑O(Td)italic_O ( italic_T italic_d )
RWKV-RNN11{}^{1}start_FLOATSUPERSCRIPT 1 end_FLOATSUPERSCRIPT 6 512 1024 0.720 - O⁢(𝐓𝐝)𝑂𝐓𝐝O(\mathbf{Td})italic_O ( bold_Td ) O⁢(𝐝)𝑂𝐝O(\mathbf{d})italic_O ( bold_d )
RWKV-RNN22{}^{2}start_FLOATSUPERSCRIPT 2 end_FLOATSUPERSCRIPT 12 512 1024 1.010 1.178 O⁢(𝐓𝐝)𝑂𝐓𝐝O(\mathbf{Td})italic_O ( bold_Td ) O⁢(𝐝)𝑂𝐝O(\mathbf{d})italic_O ( bold_d )

Table 5: Enwik8 results, measured in bits per character (bpc).

Appendix K Inference results

Figures 13 and 14 illustrate, respectively, the results on time (s) and memory (RAM, VRAM) requirements for LLM inference in float32 precision. We benchmark the following model families and sizes:

Refer to caption

Refer to caption

Figure 13: Text generation inference memory (CPU RAM, GPU VRAM) for LLMs. Model parameters are not accounted.

Refer to caption

Refer to caption

Figure 14: Text generation inference time for LLMs.

Appendix L Importance of prompt construction and comparison to GPT models

Inspired by Kocoń et al. (2023), we compared the zero-shot performance of the RWKV-4-Raven-14B with ChatGPT (access in February 2023) and GPT-4 using several known NLP tasks, i.e., recognizing textual entailment (RTE), Winograd Natural Language Inference (WNLI), and recognizing emotions elicited in readers (GoEmotions and PolEmo2). Each model got the same prompts manually chosen to receive proper responses from the ChatGPT model. As shown in Tab. 6, RWKV performs significantly worse than ChatGPT and GPT-4 in several specific tasks. We suspect that this disparity is likely caused by the choice of prompts used to generate the answers since the prompts are written in natural language and do not take into account that RWKV, as an RNN, is unable to look back inside an instruction.

Table 6: ChatGPT, GPT-4 and RWKV-4-Raven-14B reasoning performance comparison in RTE Wang et al. (2019), WNLI Wang et al. (2018), GoEmotions Demszky et al. (2020), and PolEmo2 Kocoń et al. (2019) benchmarks. RWKV GPT prompts were primarily used for ChatGPT in Kocoń et al. (2023). SOTA is provided as a supplementary reference.

When the instruction style was adapted (re-ordered) to respect that RNNs are not capable of "retrospective processing", the quality may significantly change, e.g., for RTE Wang et al. (2019) F1 Macro increased from 44.2% to 74.8%. We hypothesize that RWKV models are more sensitive to the position of the components in the context, as RNN-based architectures cannot look back and readjust the weight of previous information. For better performance, the desired information should be placed after the main question.

Can you tell me if the hypothesis is entailment or is not entailment to the premise? premise: hypothesis:

Table 7: ChatGPT and RWKV-4-Raven-14B performance comparison in Aggresion Wulczyn et al. (2017), Sarcasm Siddiqui (2019), Unhealthy Price et al. (2020), MathQA Cobbe et al. (2021), and TweetSent Barbieri et al. (2020) benchmarks. SOTA is provided as a supplementary reference.

While separating the instruction from the input is relatively easy to do, some other aspects of prompt engineering are harder to quantify. For that purpose, we also tested the approach of stating the input after the question on multiple other tasks, i.e., aggression and sarcasm detection, classification of unhealthy (offensive) texts, mathematical Q&A, and sentiment analysis, see Tab. 7. The results suggest that better prompts might reduce the disparity between models. Raven achieves comparable results to ChatGPT on unhealthy conversation detection and even surpasses it on the sarcasm detection dataset. While such an approach to prompting looks necessary, it is not enough in itself to replace the capability of having free access to the whole context. Therefore, prompt engineering seems to be significantly more important for the RNN models rather than for standard transformers. It is entirely possible that good prompts to RNN models do not mean additional restrictions, but should simply be constructed using completely different guidelines. The authors of the aforementioned paper Kocoń et al. (2023)444This is in line with the idea discussed in Wei et al. (2022b) perform chain-of-thought to improve results on the MathQA dataset. Even including this approach, the Raven model achieved a very low accuracy of 5.43%. Without it, the model performed even worse, performing only very basic and simple calculations and achieving 4.13% accuracy. Raven struggled with questions that required intermediate results. It is likely that the order of information presented in the math questions inside the dataset poses a challenge for the RWKV model. It is yet to be seen if prompt engineering can address this issue. This further emphasizes the importance of the order of information the model receives.

Write the answer to the math question in the format: "Answer: ".
The question is:

Appendix M Cases

In this part, we present a few instances of outputs produced by RWKV model using a Chat interface555ttps://github.com/BlinkDL/ChatRWKV666https://huggingface.co/spaces/BlinkDL/ChatRWKV-gradio.

Tell me about ravens. Ravens are large, black birds with a distinctive white head and neck. They are found in most parts of the world, including North America, Europe, Asia, and Australia. Ravens are known for their intelligence and problem-solving abilities. They are also considered to be symbols of death and transformation in many cultures. Ravens are often associated with the afterlife or death because they have been known to eat carrion or even other birds. In some cultures, ravens are also believed to be messengers of the gods or guardians of treasure.

Explain the concept of machine learning to a 10-year-old. Machine learning is like a super-smart computer that can learn on its own. It can recognize patterns and make decisions based on what it has learned. Just like how you might remember something from the past, the computer can remember what it’s learned and use it to make predictions or decisions in the future. It’s like having a smart friend who always knows what to do!

Write a short poem describing nature’s beauty in springtime. Nature’s beauty in springtime is rare, A symphony of colors, a gentle surprise. The flowers bloom, the birds sing and sing, As the sun shines down, warming our skin. The trees sway in the gentle breeze, As we take in their beauty with ease. The sky is blue, the clouds are white, A sight to behold, it’s truly divine.

Below is the numerical stable formula for updating the state variable in RWKV: at*superscriptsubscript𝑎𝑡\displaystyle a_{t}^{*}italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT start_POSTSUPERSCRIPT * end_POSTSUPERSCRIPT =ep−q⁢at−1+eu+kt−q⁢vtabsentsuperscript𝑒𝑝𝑞subscript𝑎𝑡1superscript𝑒𝑢subscript𝑘𝑡𝑞subscript𝑣𝑡\displaystyle=e^{p-q}a_{t-1}+e^{u+k_{t}-q}v_{t}= italic_e start_POSTSUPERSCRIPT italic_p - italic_q end_POSTSUPERSCRIPT italic_a start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_u + italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT - italic_q end_POSTSUPERSCRIPT italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT bt*superscriptsubscript𝑏𝑡\displaystyle b_{t}^{*}italic_b start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT start_POSTSUPERSCRIPT * end_POSTSUPERSCRIPT =ep−q⁢bt−1+eu+kt−qabsentsuperscript𝑒𝑝𝑞subscript𝑏𝑡1superscript𝑒𝑢subscript𝑘𝑡𝑞\displaystyle=e^{p-q}b_{t-1}+e^{u+k_{t}-q}= italic_e start_POSTSUPERSCRIPT italic_p - italic_q end_POSTSUPERSCRIPT italic_b start_POSTSUBSCRIPT italic_t - 1 end_POSTSUBSCRIPT + italic_e start_POSTSUPERSCRIPT italic_u + italic_k start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT - italic_q end_POSTSUPERSCRIPT w⁢k⁢vt𝑤𝑘subscript𝑣𝑡\displaystyle wkv_{t}italic_w italic_k italic_v start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT =at*bt*absentsuperscriptsubscript𝑎𝑡superscriptsubscript𝑏𝑡\displaystyle=\frac{a_{t}^{*}}{b_{t}^{*}}= divide start_ARG italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT start_POSTSUPERSCRIPT * end_POSTSUPERSCRIPT end_ARG start_ARG italic_b start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT start_POSTSUPERSCRIPT * end_POSTSUPERSCRIPT end_ARG Please write out the corresponding pytorch code to do so, where both atsubscript𝑎𝑡a_{t}italic_a start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT and btsubscript𝑏𝑡b_{t}italic_b start_POSTSUBSCRIPT italic_t end_POSTSUBSCRIPT are torch.Tensor([n]) Here is the PyTorch code to update the state variable in RWKV using the given formula:

import torch

def update_state(a, b, u, v, p, q):

a_t = torch.exp(p - q) * a + torch.exp(u + k - q) * v

b_t = torch.exp(p - q) * b + torch.exp(u + k - q)

return a_t, b_t

a, b = torch.randn(2, 3)

u = torch.randn(2)

v = torch.randn(2)

wkv = torch.randn(2)

a_t, b_t = update_state(a, b, u, v, wkv, torch.randn(2))

print("a_t:", a_t)

print("b_t:", b_t)

This code defines a function update_state that takes in the state variables a and b, the control variables u and v, and the parameters wkv and torch.randn. It then uses the PyTorch torch.exp function to compute the exponential of the control variables and the log of the normalization factor. Finally, it uses these values to update the state.

[h] Explain the following metaphor: Life is like cats. The metaphor "Life is like cats" suggests that life can be unpredictable and full of unexpected twists and turns, just like a cat’s behavior. It implies that we cannot fully control or predict the outcome of our lives, just as we cannot control the movements of a cat. The metaphor also implies that life is complex and multifaceted, with many different variables and factors at play. Just as a cat can have multiple personalities or moods, life can have many twists and turns that are difficult to predict or understand. Overall, the metaphor suggests that life is complex and unpredictable, much like a cat’s behavior.

[ht] Test the convergence of ∑i=0∞2−isuperscriptsubscript𝑖0superscript2𝑖\sum_{i=0}^{\infty}2^{-i}∑ start_POSTSUBSCRIPT italic_i = 0 end_POSTSUBSCRIPT start_POSTSUPERSCRIPT ∞ end_POSTSUPERSCRIPT 2 start_POSTSUPERSCRIPT - italic_i end_POSTSUPERSCRIPT. The given series is a geometric series with the first term a0=2subscript𝑎02a_{0}=2italic_a start_POSTSUBSCRIPT 0 end_POSTSUBSCRIPT = 2 and common ratio r=12𝑟12r=\frac{1}{2}italic_r = divide start_ARG 1 end_ARG start_ARG 2 end_ARG. We can find the sum of this series using the formula for a geometric series: ∑_i=0^∞a_i = a01-r = 121-12 = 11-12 = 21 = 2 So, the sum of the given series is 22\boxed{2}2.

I would like to cook some Chinese food at home. Do you have any suggestions on the meal I could choose? Yes, of course! If you are not very experienced in cooking Chinese food, you could start with making noodles or rice, then make a dish using one of those. I also think making stir-fry with some meat is an excellent idea. I want to make something more special. What about red braised pork belly? Yes, red braised pork belly is also very popular and tasty. Tell me the steps to cook. You will need a few ingredients such as pork belly, star anise, sugar, ginger, scallions and wine. Put all of these in a pot, cook on medium heat and simmer until the meat is tender. After that, remove the meat from the pot and reserve the sauce. Slice the meat and add some ginger slices into it, then boil until it’s almost done. What do I need to be careful with the most during the cooking process? To cook pork belly properly, it is very important to monitor the heat of the stove and be patient with the cooking process. Overcooking it can cause the pork to be tough, dry and difficult to chew. You can add a few ingredients like soy sauce or black beans to give the pork more flavor. Tell me more about the proper heat and timing. If the temperature of the water is too high, it can damage the collagen and connective tissues, resulting in tough pork belly. It is essential to control the temperature carefully, especially when braising the pork. Additionally, keep the water simmering but not boiling to maintain a low simmer, as high heat will dry out the pork belly and result in tough and unpalatable meat. Lastly, monitor the heat during cooking and make sure the heat is medium, as this is what is most appropriate for slow cooking pork belly. I see. Thank you for your advice. You’re welcome! It was a pleasure chatting with you. Let me know if you need any further assistance in the future.

[h] If all apples are fruits and all fruits are sweet, are all apples sweet? Please explain your answer in detail. Yes, all apples are fruits, and all fruits are sweet. Therefore, it is safe to assume that all apples are sweet. If birds can fly. Penguins are birds. Can penguins fly? Please explain your answer in detail. No, penguins cannot fly. Penguins are birds, but they are not capable of flight like other birds. They have wings and can glide for short distances, but they cannot take off or fly long distances like other birds. The day before two days after the day before tomorrow is Saturday. What day is it today? Please explain your answer in detail. Today is Friday. Two days after tomorrow is Sunday. Tomorrow is Saturday. If five cats can catch five mice in five minutes, how long will it take one cat to catch one mouse? Please explain your answer in detail. If five cats can catch five mice in five minutes, it means that each cat can catch one mouse in one minute. To find out how long it will take one cat to catch one mouse, we can use the concept of proportionality. Let’s assume that it takes x minutes for a cat to catch a mouse. Then, the number of mice caught by a cat in one minute is equal to the number of cats multiplied by the number of mice caught by each cat: x = 5 * 1 = 5 So, it will take 5 minutes for one cat to catch one mouse.

Summarize the following paragraph. The paper presents a novel architecture, Receptance Weighted Key Value (RWKV), that combines the efficient parallelizable training of Transformers with the efficient inference of RNNs. The model is able to scale to tens of billions of parameters and exhibits linear computational complexity during training and inference, making it a promising alternative to Transformers for sequence processing tasks. Please further summarize it in one sentence. The researchers have developed a new architecture that combines the efficient parallelizable training of Transformers with the efficient inference of RNNs, which allows for linear scaling in memory and computational requirements.