RWKV : Reinventing RNN s for the Transformer Era (2023)

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.

## 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 @xcite @xcite . predominant among these techniques include rnns and transformers @xcite , 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 @xcite . transformers emerged as a powerful alternative, adept at managing local and long-range dependencies and supporting parallelized training @xcite . models such as @xcite , @xcite , llama @xcite , and chinchilla @xcite 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 @xcite @xcite . 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 @xcite 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 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: here we briefly review the fundamentals of rnns and transformers. ## rwkv the rwkv model architecture is defined by four fundamental elements that are intrinsic to the timemixing and channel-mixing blocks: these core elements interact multiplicatively at each timestep, as depicted in figure 2 . ## 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 the number of parameters for each model is computed using the formula: # @xmath0, d represents the model dimension and l corresponds to the number of layers. flops is for a forward pass for one token. it was calculated as 2(2v d + 13d 2 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(2v d + 13d 2 l) flop per token. notably, this matches the standard formula for flop calculations in transformers @xcite ## 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? ## 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 hug-gingface transformers @xcite . 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. figure 7 : cumulative time on text generation for llms. unlike transformers, rwkv exhibits linear scaling. ## 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 wkv t step to reduce the computational cost to o(b log(t )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. ## 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 longrange 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.