Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

RWKV-X: A Linear Complexity Hybrid Language Model

Haowen Hou, Zhiyi Huang, Kaifeng Tan, Rongchang Lu, Fei Richard Yu

TL;DR

Transformers incur $O(N^2)$ self-attention costs that hinder long-context processing. RWKV-X hybridizes RWKV blocks with Top-$k$ Chunk Sparse Attention to achieve $O(N)$ training and $O(1)$ decoding, enabling ultra-long contexts while preserving short-context performance. Key contributions include the Top-$k$ Chunk Sparse Attention with KV-cache compression, a block-expansion training regime, and long-context continual pretraining at 64K with Long-context Cross-Entropy loss, achieving near-perfect pass-key retrieval on 64K sequences and stable decoding up to 1M tokens. This work provides a scalable, efficient backbone for long-context language modeling and offers public checkpoints and code to accelerate further research.

Abstract

In this paper, we introduce RWKV-X, a novel hybrid architecture that combines the efficiency of RWKV for short-range modeling with a sparse attention mechanism designed to capture long-range context. Unlike previous hybrid approaches that rely on full attention layers and retain quadratic complexity, RWKV-X achieves linear-time complexity in training and constant-time complexity in inference decoding. We demonstrate that RWKV-X, when continually pretrained on 64K-token sequences, achieves near-perfect accuracy on the 64K passkey retrieval benchmark. It consistently outperforms prior RWKV-7 models on long-context benchmarks, while maintaining strong performance on short-context tasks. These results highlight RWKV-X as a scalable and efficient backbone for general-purpose language modeling, capable of decoding sequences up to 1 million tokens with stable speed and memory usage. To facilitate further research and analysis, we have made the checkpoints and the associated code publicly accessible at: https://github.com/howard-hou/RWKV-X.

RWKV-X: A Linear Complexity Hybrid Language Model

TL;DR

Transformers incur self-attention costs that hinder long-context processing. RWKV-X hybridizes RWKV blocks with Top- Chunk Sparse Attention to achieve training and decoding, enabling ultra-long contexts while preserving short-context performance. Key contributions include the Top- Chunk Sparse Attention with KV-cache compression, a block-expansion training regime, and long-context continual pretraining at 64K with Long-context Cross-Entropy loss, achieving near-perfect pass-key retrieval on 64K sequences and stable decoding up to 1M tokens. This work provides a scalable, efficient backbone for long-context language modeling and offers public checkpoints and code to accelerate further research.

Abstract

In this paper, we introduce RWKV-X, a novel hybrid architecture that combines the efficiency of RWKV for short-range modeling with a sparse attention mechanism designed to capture long-range context. Unlike previous hybrid approaches that rely on full attention layers and retain quadratic complexity, RWKV-X achieves linear-time complexity in training and constant-time complexity in inference decoding. We demonstrate that RWKV-X, when continually pretrained on 64K-token sequences, achieves near-perfect accuracy on the 64K passkey retrieval benchmark. It consistently outperforms prior RWKV-7 models on long-context benchmarks, while maintaining strong performance on short-context tasks. These results highlight RWKV-X as a scalable and efficient backbone for general-purpose language modeling, capable of decoding sequences up to 1 million tokens with stable speed and memory usage. To facilitate further research and analysis, we have made the checkpoints and the associated code publicly accessible at: https://github.com/howard-hou/RWKV-X.
Paper Structure (31 sections, 6 equations, 7 figures, 8 tables)

This paper contains 31 sections, 6 equations, 7 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Linear Complexity Language Models
  4. Hybrid Language Models
  5. Sparse Attention
  6. Method
  7. Preliminaries
  8. Top-$k$ Chunk Sparse Attention
  9. KV Cache Management
  10. Complexity Analysis
  11. RWKV-X
  12. Block Expansion Method
  13. Long-context Continual Pretraining
  14. Experiments
  15. Experiment Setup
  16. ...and 16 more sections

Figures (7)

  • Figure 1: Passkey retrieval performance of RWKV-X models on documents up to 64K tokens. Results are shown for: (a) RWKV-7 pretrained with a 4K context length; (b) RWKV-7 after continual pretraining with a 128K context length; and (c) RWKV-X trained with continual pretraining on a 64K context length.
  • Figure 2: The architecture of RWKV-X, a hybrid model that combines RWKV-7 blocks with Sparse Attention blocks.
  • Figure 3: Prefill latency comparison between RWKV-X and a full-attention Transformer.
  • Figure 4: Decoding latency comparison between RWKV-7-2.9B and RWKV-X-3.6B models. The horizontal axis represents the context length (log scale).
  • Figure 5: Validation loss vs. percentage of attention layers for 124M-parameter RWKV-X models (12 layers). 0% = RWKV-7, 100% = Fully Sparse-Attention Transformer.
  • ...and 2 more figures