RWKV: Reinventing RNNs for the Transformer Era

Problem

Framing

Transformers train well but pay $O(T^2)$ memory and compute during inference over long contexts. RWKV closes this gap with a recurrent reformulation of attention-like mixing that stays parallel during training and linear at inference, while remaining competitive with similarly sized LMs on Pile-trained zero-shot benchmarks.

Currently Used Methods

Foundational

• @vaswaniAttentionAllNeed2017 — full self-attention gives strong sequence modeling with global token interactions.
  • Limitation in context: quadratic inference cost blocks long-context efficiency.
• @hochreiterLSTM1997 — recurrent state gives linear-time autoregressive inference.
  • Limitation in context: training is sequential and scales poorly.
• "An Attention Free Transformer" — replaces pairwise attention with kernelized content-position mixing.
  • Limitation in context: not cast as an efficient recurrent cell.
• "Reformer: The Efficient Transformer" — reduces attention cost with sparse and hashed approximations.
  • Limitation in context: still keeps transformer-style sequence memory overhead.
• "Linear Transformers Are Secretly Fast Weight Memory Systems" — linearizes attention with kernel feature maps.
  • Limitation in context: does not recover RWKV's simple recurrent update.

Proposed Method

Architecture

RWKV stacks pre-norm residual blocks with two sublayers: Time Mixing and Channel Mixing. Time Mixing computes receptance, key, value, and a decayed weighted key-value aggregation; Channel Mixing is a gated feed-forward block with token shift.

Loss / Objective

The model trains with standard next-token cross-entropy.

Sampling Rule / Algorithm

Autoregressive inference updates a recurrent WKV state and emits the next-token distribution from the LM head.

Training Procedure

• Optimizer: Adam, $\beta = (0.9, 0.99)$
• Weight decay: $0$
• Precision: bfloat16
• Context length: $1024$
• Batch size: 128 or 256 sequences
• Training data: one epoch on the Pile, 330B tokens
• Initial learning rate: $6\times 10^{-4}$ to $1\times 10^{-4}$ by model size
• End learning rate: about $1\times 10^{-5}$

Evaluation

Datasets

• The Pile pretraining corpus
• ARC-Easy, ARC-Challenge
• BoolQ, COPA, HeadQA
• HellaSwag, LAMBADA, OpenBookQA
• PiQA, ReCoRD, SciQ, Winogrande
• Enwik8
• Long Range Arena

Metrics

• Zero-shot accuracy
• Validation loss
• Bits per character
• Inference time
• RAM / VRAM usage

Headline results

• Scaling-law fit: Pareto linear fit reaches $r^2 = 0.994$.
• Scaling-law extrapolation: one-order extrapolation still reaches $r^2 = 0.875$.
• Pile-trained zero-shot NLP: RWKV tracks OPT, BLOOM, and Pythia at matched scale.
• Inference complexity: linear scaling in generation time and memory, unlike transformer baselines.
• Enwik8: competitive test bpc against efficient-transformer baselines.

Ablations

• Model size sweep: larger RWKV improves zero-shot accuracy across common NLP tasks.
• Compute scaling: loss follows transformer-style log-log scaling trends.
• Context-length training: longer contexts improve long-context language modeling.
• Inference benchmark: advantage widens as prompt length and model size grow.

Method Strengths and Weaknesses

Strengths

• Parallel training with recurrent inference targets the main transformer efficiency bottleneck.
• Scaling-law fit with $r^2=0.994$ suggests predictable large-model behavior.
• Zero-shot accuracy stays competitive with OPT, BLOOM, and Pythia at similar scales.
• Linear-time inference and lower memory directly matter for long-context deployment.

Weaknesses

• Core WKV mechanism is more specialized than standard attention and less widely validated.
• Reported competitiveness is mostly against open-source baselines, not strongest closed models.
• Training still uses fixed 1024-token chunks for pretraining.
• Evidence for long-context gains is narrower than the broad zero-shot suite.

Suggestions from the authors

• Extend RWKV to more modalities, including vision and audio.
• Study stronger long-context training and evaluation regimes.
• Improve kernels and systems support for faster large-scale training.
• Analyze theoretical properties of recurrent state and gradient flow.

Links

Prior Papers

• @vaswaniAttentionAllNeed2017 — RWKV keeps transformer-style parallel training while removing quadratic inference.
• @hochreiterLSTM1997 — RWKV inherits recurrent hidden-state efficiency without classic sequential training constraints.

Further Papers

• @guMamba2023 — another sequence model that replaces full attention with linear-time state-space style recurrence.