Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ARWKV: Pretrain is not what we need, an RNN-Attention-Based Language Model Born from Transformer

Lin Yueyu, Li Zhiyuan, Peter Yue, Liu Xiao

TL;DR

ARWKV proposes an RNN-attention hybrid by adapting RWKV-7 time-mixing to an architecture initially based on Transformer, with the aim of reducing pretraining cost while enhancing expressiveness. The authors implement a three-stage pipeline—Stage 1 time-mixing attention alignment, Stage 2 knowledge distillation from a 32B teacher to a 7B student, and Stage 3 SFT and DPO to extend context and align with user preferences—to transfer transformer-style patterns into an RWKV-based model. They demonstrate that a 7B ARWKV can be trained on a single A100-80G and a 32B variant on multi-GPU setups, and find FP16 inference improves performance over BF16, while aggressive distillation without gating can hinder accuracy. The work offers a practical path toward efficient, expressive LLMs and suggests broad generalization to MoE, multimodal, and compression scenarios.

Abstract

As is known, hybrid quadratic and subquadratic attention models in multi-head architectures have surpassed both Transformer and Linear RNN models , with these works primarily focusing on reducing KV complexity and improving efficiency. For further research on expressiveness, we introduce our series of models distilled from Qwen 2.5, based on pure native RWKV-7 attention, which aims to make RNN more expressive and demonstrates state tracking ability beyond transformers. We work with QRWK 32B based on RWKV-6 architecture, another approach that reduces the entire knowledge processing time to just 8 hours using 16 AMD MI300X GPUs while maintaining Qwen 2.5's performance. In fact, the distillation process can utilize any LLM, not just Qwen, and enables knowledge transfer from larger LLMs to smaller ones with more fewer tokens. We will explain the detailed process and share our insights on building more powerful foundation models. Please note that this is an ongoing work that will be updated continuously. The model checkpoints and source code are available at \href{https://github.com/yynil/RWKVInside}{https://github.com/yynil/RWKVInside}, \href{https://huggingface.co/RWKV-Red-Team/ARWKV-7B-Preview-0.1}{https://huggingface.co/RWKV-Red-Team/ARWKV-7B-Preview-0.1}.

ARWKV: Pretrain is not what we need, an RNN-Attention-Based Language Model Born from Transformer

TL;DR

ARWKV proposes an RNN-attention hybrid by adapting RWKV-7 time-mixing to an architecture initially based on Transformer, with the aim of reducing pretraining cost while enhancing expressiveness. The authors implement a three-stage pipeline—Stage 1 time-mixing attention alignment, Stage 2 knowledge distillation from a 32B teacher to a 7B student, and Stage 3 SFT and DPO to extend context and align with user preferences—to transfer transformer-style patterns into an RWKV-based model. They demonstrate that a 7B ARWKV can be trained on a single A100-80G and a 32B variant on multi-GPU setups, and find FP16 inference improves performance over BF16, while aggressive distillation without gating can hinder accuracy. The work offers a practical path toward efficient, expressive LLMs and suggests broad generalization to MoE, multimodal, and compression scenarios.

Abstract

As is known, hybrid quadratic and subquadratic attention models in multi-head architectures have surpassed both Transformer and Linear RNN models , with these works primarily focusing on reducing KV complexity and improving efficiency. For further research on expressiveness, we introduce our series of models distilled from Qwen 2.5, based on pure native RWKV-7 attention, which aims to make RNN more expressive and demonstrates state tracking ability beyond transformers. We work with QRWK 32B based on RWKV-6 architecture, another approach that reduces the entire knowledge processing time to just 8 hours using 16 AMD MI300X GPUs while maintaining Qwen 2.5's performance. In fact, the distillation process can utilize any LLM, not just Qwen, and enables knowledge transfer from larger LLMs to smaller ones with more fewer tokens. We will explain the detailed process and share our insights on building more powerful foundation models. Please note that this is an ongoing work that will be updated continuously. The model checkpoints and source code are available at \href{https://github.com/yynil/RWKVInside}{https://github.com/yynil/RWKVInside}, \href{https://huggingface.co/RWKV-Red-Team/ARWKV-7B-Preview-0.1}{https://huggingface.co/RWKV-Red-Team/ARWKV-7B-Preview-0.1}.
Paper Structure (10 sections, 5 equations, 5 figures, 1 table)

This paper contains 10 sections, 5 equations, 5 figures, 1 table.

Table of Contents

  1. Introduction
  2. Architecture
  3. From Transformer to RNN
  4. Stage 1 - Time Mixing module replacing Self-Attention
  5. Stage 2 - Knowledge Distillation
  6. Stage 3 - SFT and DPO
  7. Evaluation
  8. Conclusions
  9. Future Work
  10. Acknowledgements

Figures (5)

  • Figure 1: replace self-attention by RWKV-7 time mixing module
  • Figure 2: RWKV-7 architecture.capability of attention is the key for RNN-based LLMs, which in this case is Time mixing module
  • Figure 3: General Decoder Layer in transformer
  • Figure 4: We replace the standard Attention with an AttentionWrapper that contains both the original self-attention mechanism and a TimeMixer. The TimeMixer is trained to minimize the gap between its output and that of the self-attention module. The final output combines the hidden states from the original self-attention with the residual difference between self-attention and TimeMixer outputs. This architecture enables the model to optimize the TimeMixer to progressively reduce the discrepancy between self-attention and TimeMixer outputs.
  • Figure 5: Stage-1 loss, 18 hours with one 8*h800 80G , context length 2048, 4B tokens