Source-linked AI summary
RWKV: Reinventing RNNs for the Transformer Era
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, Rui-Jie Zhu
TL;DR
RNNs offer efficient memory use but lack Transformer-level parallelization and scalability, while Transformers incur quadratic costs on long sequences. RWKV addresses this gap with linear attention that supports both Transformer-style training and RNN-style inference, and performs comparably to state-of-the-art models at reduced cost. Its efficiency supports scaling and deployment, though very long-context recall and prompt sensitivity remain limitations.
Problem
RNNs have efficient memory use but limited training parallelization, whereas Transformers handle dependencies and parallelize training at quadratic computational and memory cost for long sequences.
Method
RWKV uses a channel-wise, recurrently formulated linear-attention WKV operator that combines Transformer-style parallelized computation with RNN-style processing.
Results
RWKV achieves comparable performance to state-of-the-art models across diverse NLP tasks with reduced cost.
Takeaways & Limitations
RWKV’s lower inference cost makes it more suitable for deployment on consumer and edge hardware and lowers resource barriers to language-model applications.
Takeaways & Limitations
Linear attention may limit recall of minutiae across very long contexts because information is funneled through a single vector representation over time.
Abstract
from arXiv · showhide
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.
1 Introduction
RWKV is introduced to combine RNN memory efficiency with Transformer parallelization and scalability while avoiding quadratic attention costs. The paper evaluates this architecture across NLP benchmarks and releases models ranging from 169 million to 14 billion parameters.
- RNNs use less memory but suffer from vanishing gradients and non-parallelizable training, limiting scalability.
- Transformers support local and long-range dependencies and parallelized training, but self-attention has quadratic computational and memory complexity for long sequences.
- RWKV uses linear attention and channel-directed interactions to reduce Transformer memory bottlenecks and quadratic scaling without approximation.
- RWKV combines RNN and Transformer strengths, targeting efficient processing of large-scale sequential models while retaining expressive capacity and parallelized training.
- The paper contributes benchmark experiments and releases pretrained models from 169 million to 14 billion parameters trained on the Pile.
2 Background
RNNs process sequences recurrently and therefore are difficult to parallelize, whereas Transformers use attention to relate tokens across a sequence. RWKV adapts attention weights into channel-wise, position-dependent decays that can be transformed into an RNN.
- Typical RNNs depend on previous time steps, which prevents parallelization despite separable linear and recurrent blocks.
- Transformers use attention mechanisms to capture relationships between all input and output tokens rather than operating step-by-step.
- Transformer attention includes pairwise token interactions represented through the QK^⊤ multiplication and vector operations.
- AFT introduces learned pairwise position biases, while RWKV modifies interaction weights so the computation can be transformed into an RNN.
- RWKV represents each interaction weight as a channel-wise time-decay vector scaled by relative position and traced backward from the current time.
- RWKV requires non-negative decay parameters so per-channel weights decay backward in time and remain at most one.
3 RWKV
RWKV combines time-mixing and channel-mixing residual blocks with recurrent, attention-like computation. Its design supports parallelized training while permitting recurrent inference through a time-sequential formulation.
- Architecture: R, W, K, and V represent receptance, trainable positional decay, key, and value components that interact multiplicatively at each timestep.
- Architecture: RWKV blocks contain time-mixing and channel-mixing sub-blocks that use recurrent structures to incorporate past information.
- Token Shift: RWKV produces projection vectors from linear combinations of current and previous block inputs, implementing a token shift across time.
- WKV Operator: The WKV operator uses a channel-wise, position-modified decay vector rather than AFT’s pairwise weight matrix, enabling recurrent updates.
- Computation Modes: Time-parallel processing has O(BTd^2) complexity for matrix multiplications and O(BTd) for the serial-scan WKV update.
- Computation Modes: During inference, RWKV can be recursively formulated in time-sequential mode, while its matrix multiplications and other operations remain parallelizable during training.
4 Trained Models and Computing Costs
The experiments train RWKV models from 169 million to 14 billion parameters and examine their computational scaling. RWKV follows the general scaling-law form established for Transformers with strong fitted relationships.
- Trained Models: Six RWKV models span 169 million to 14 billion parameters, each trained for one epoch comprising 330 billion Pile tokens.
- Computing Costs: Model parameter counts and FLOP estimates are calculated from vocabulary size, model dimension, and layer count using explicit architecture-based formulas.
- Training Setup: Training uses Adam without weight decay, bfloat16 precision, 1024-token contexts, exponential learning-rate decay, and an auxiliary loss.
- Scaling Laws: RWKV follows the general scaling-law form established for Transformers across 45 models trained with varied dataset and parameter sizes.
- Scaling Laws: r2 = 0.994 for the linear fit to Pareto-optimal loss-versus-compute points, while an additional order-of-magnitude extrapolation achieves r2 = 0.875.
5 Evaluations
The evaluations test RWKV’s competitiveness with similarly sized Transformers and its ability to use longer contexts. Results cover FLOP-matched language benchmarks, prompt-sensitive comparisons, and progressively extended context lengths.
- Competitiveness: RWKV is compared with Pythia, OPT, and BLOOM models of similar size and training-token scale on a FLOP-matched basis.
- Competitiveness: The benchmark suite includes ARC, BoolQ, COPA, HeadQA, HellaSwag, LAMBADA, OpenBookQA, PIQA, ReCoRD, SciQ, and Winogrande.
- Additional Comparisons: Reordered prompts increased RWKV’s F1 performance from 44.2% to 74.8%, demonstrating sensitivity to prompt engineering.
- Extended Context: Context length is progressively increased from 1024 to 2048, 4096, and 8192 tokens through finetuning on additional tokens from the Pile.
- Extended Context: Increasing context length leads to lower mean test loss on the Pile, indicating effective use of longer contextual information.
- Long-Sequence Evaluation: Long-Range Arena evaluates RWKV on sequences ranging from 1,000 to 16,000 tokens across text, natural-language, synthetic-image, and mathematical-expression tasks.
6 Inference Experiments
Inference experiments measure text-generation speed and memory requirements across model sizes and families on CPU and NVIDIA A100 hardware. The reported comparison highlights RWKV’s linear scaling during text generation.
- Evaluation Setup: Inference evaluation measures text-generation speed and memory requirements across model sizes and model families.
- Evaluation Setup: Experiments use x86 CPUs and NVIDIA A100 80 GB GPUs with float32 precision and HuggingFace Transformers.
- Evaluation Scope: All model parameters, including embedding and non-embedding layers, are included in parameter counts; quantization effects are left for further work.
- Results: RWKV exhibits linear cumulative time scaling for text generation, unlike Transformers.
7 Future Work
Future work targets greater RWKV expressivity and efficiency while exploring extensions to other architectures and uses of its recurrent state.
- Future work can enhance RWKV expressivity by improving time-decay formulations and exploring initial model states while maintaining efficiency.
- A parallel scan in the wkvt step could reduce RWKV computational cost to O(B log(T)d).
- RWKV mechanisms could replace cross-attention in encoder-decoder architectures for more efficient seq2seq or multimodal training and inference.
- RWKV’s recurrent state may support interpretability, predictability, safety, and behavior customization through hidden-state manipulation.
- Larger internal states could improve RWKV’s memory for previous context and performance across tasks.
8 Conclusions
RWKV combines linear-cost attention and recurrent computation with training parallelization, and is benchmarked as a scalable alternative to Transformers. The authors report comparable state-of-the-art performance at reduced cost and pretrained models reaching tens of billions of parameters.
- RWKV is introduced as an RNN approach using time-based mixing components to capture locality and long-range dependencies.
- RWKV replaces quadratic QK attention with a scalar formulation at linear cost.
- RWKV reformulates recurrence and sequential inductive biases to enable parallelized training and efficient inference, while custom initializations improve training dynamics.
- RWKV shows comparable performance to SoTA across diverse NLP tasks with reduced cost, alongside experiments on expressivity, interpretability, and scaling.
- RWKV provides pretrained models with tens of billions of parameters, supporting its proposed route toward scalable and efficient sequential-data architectures.
9 Limitations
RWKV’s efficiency gains come with limitations in long-context detail recall and prompt sensitivity. The paper also reports that prompt information ordering can substantially affect performance.
- Linear attention may limit recalling minute details over very long contexts because information is funneled through a single vector across time.Learned time decay helps preserve information but remains mechanistically limited compared with full self-attention.
- RWKV places greater importance on prompt engineering than standard Transformer models because less prompt information carries into the continuation.
- Changing the order of information pieces almost doubled RWKV performance on some tasks, confirming the sensitivity to prompt construction.
10 Ethics Statement
The paper presents RWKV as an openly released model family intended to lower deployment and research barriers. It also notes that broader access may increase exposure to generated text, bias, and harmful content.
- RWKV training code is released open source, enabling the Open Source community to inspect, study, and finetune the model for particular tasks.
- Pretrained RWKV weights from 0.1B to 14B parameters are released across multiple languages to ease adoption and support study of emergent phenomena.
- Lower inference cost makes RWKV more suitable for consumer and edge hardware and lowers resource barriers for chat assistants and text generation.
- Lower resource barriers may spread AI-generated text, while current RWKV models may reproduce biases and potentially harmful content from their training data.The paper states that mitigation and finetuning strategies developed for large Transformer models should also apply to RWKV.
J.2 Evaluation on Long Range Arena
RWKV is evaluated on the Long Range Arena, which tests models on sequences from 1,000 to 16,000 tokens across several data types. It performs strongly on natural-language and code-related tasks but is weaker on some image and path tasks.
- The Long Range Arena evaluates lengthy-context handling with sequences ranging from 1,000 to 16,000 tokens across text, natural language, synthetic images, and mathematical expressions.
- RWKV performs second only to S4 on five datasets in the reported Long Range Arena results.
- RWKV performs on par with or nearly so to S4 on natural-language and computer-code problems, but substantially underperforms S4 on Image, Pathfinder, and Path-X.
K Inference results
The inference evaluation compares RWKV with several language-model families across time and memory requirements, reasoning benchmarks, and prompt-sensitive tasks. Results indicate that RWKV’s performance depends strongly on prompt organization, with adapted prompts narrowing disparities on some tasks.
- Inference benchmarking: Figures 13 and 14 evaluate float32 inference time and memory requirements across RWKV, Bloom, OPT, GPT-Neo, and Pythia models.The benchmark includes model sizes ranging from roughly 125m to 14b parameters, depending on the family.
- Prompt sensitivity: 44.2% to 74.8%: RTE F1 Macro increased after re-ordering the instruction style to account for RNNs’ lack of retrospective processing.The authors hypothesize that RWKV is sensitive to component position because RNN-based architectures cannot revisit and reweight earlier information.
- Prompt sensitivity: Better prompts may reduce model disparities, with Raven comparable to ChatGPT on unhealthy-conversation detection and surpassing it on the sarcasm dataset.The authors state that prompt adaptation appears necessary but is not sufficient by itself to replace the capabilities being compared.
- Architectural implication: RWKV combines parallelizable Transformer-style training with efficient RNN-style inference and is described as having linear computational complexity during training and inference.The architecture is presented as an alternative to Transformers for sequence-processing tasks.