Skip to content
Published on

RWKV: Reinventing RNNs for the Transformer Era — 从 v4 到 v7 Goose

🇰🇷KO🇺🇸EN🇯🇵JA🇨🇳ZH
필사 모드
分享
Authors
RWKV Architecture

论文概述

  • 标题: RWKV: Reinventing RNNs for the Transformer Era
  • 作者: Bo Peng 等（RWKV Foundation）
  • 首次发表: 2023 年 5 月（arXiv: 2305.13048）,EMNLP 2023
  • 最新版本: RWKV-7 "Goose"（2025 年 3 月发布）
  • 代码: github.com/BlinkDL/RWKV-LM

动机:Transformer vs RNN 的两难

现代 LLM 几乎都以 Transformer 为基础,但存在根本性的局限:

  • O(N²) Self-Attention: 相对于序列长度呈二次复杂度
  • KV Cache 爆炸: 推理时内存占用与序列长度成正比
  • 长上下文成本: 在 128K+ 上下文下,成本与内存急剧攀升

相比之下,传统 RNN 具有:

  • O(N) 复杂度: 线性时间与内存
  • 固定状态: 恒定的推理成本
  • 但: 训练无法并行化,速度慢,捕捉长期依赖能力弱

RWKV 提出的问题:"能不能像 Transformer 一样并行训练,又像 RNN 一样高效推理?"

RWKV 核心架构

RWKV 这个名字本身就概括了架构的核心要素:

  • R: Receptance（接受门）— 决定接受多少过去信息
  • W: Weight（时间衰减）— 过去信息的衰减率
  • K: Key — 当前输入的键
  • V: Value — 当前输入的值

WKV（Weighted Key-Value）机制

RWKV 的核心运算 WKV 如下:

import torch

def wkv_vanilla(w, u, k, v):
    """
    RWKV 的 WKV 机制(纯 Python 实现)
    w: time decay(负值,绝对值越大衰减越快)
    u: bonus(当前 token 的加成)
    k: key
    v: value
    """
    T, C = k.shape
    output = torch.zeros_like(v)

    for c in range(C):
        # 按通道独立处理(O(T) per channel)
        a = 0.0  # 累积分子
        b = 0.0  # 累积分母
        p = -1e30  # 最大值(数值稳定性)

        for t in range(T):
            # 当前 token 的贡献
            e1 = torch.exp(torch.clamp(u[c] + k[t, c] - p, max=30))
            e2 = torch.exp(torch.clamp(w[c] + p - p, max=30))  # 前一次累积的衰减

            # WKV 计算
            wkv = (e1 * v[t, c] + e2 * a) / (e1 + e2 * b)
            output[t, c] = wkv

            # 状态更新(RNN 方式)
            new_p = max(w[c] + p, k[t, c])
            e1 = torch.exp(k[t, c] - new_p)
            e2 = torch.exp(w[c] + p - new_p)

            a = e2 * a + e1 * v[t, c]
            b = e2 * b + e1
            p = new_p

    return output

双模式:Transformer 模式 vs RNN 模式

# 训练时:Transformer 模式(并行)
# 一次性处理整个序列,O(T) 复杂度
def rwkv_parallel(x, w, u, k_proj, v_proj, r_proj):
    T, C = x.shape
    k = k_proj(x)  # (T, C)
    v = v_proj(x)  # (T, C)
    r = torch.sigmoid(r_proj(x))  # (T, C) 接受门

    # 并行 WKV 运算(CUDA kernel)
    wkv = parallel_wkv_cuda(w, u, k, v)

    return r * wkv  # 应用接受门


# 推理时:RNN 模式(顺序,恒定内存)
def rwkv_sequential(x_t, state, w, u, k_proj, v_proj, r_proj):
    """逐 token 处理,O(1) 复杂度"""
    k = k_proj(x_t)
    v = v_proj(x_t)
    r = torch.sigmoid(r_proj(x_t))

    # state = (a, b, p) — 固定大小!
    a, b, p = state

    e1 = torch.exp(u + k - p)
    e2 = torch.exp(w + p - p)

    wkv = (e1 * v + e2 * a) / (e1 + e2 * b)
    output = r * wkv

    # 状态更新
    new_p = torch.maximum(w + p, k)
    e1 = torch.exp(k - new_p)
    e2 = torch.exp(w + p - new_p)
    new_a = e2 * a + e1 * v
    new_b = e2 * b + e1
    new_state = (new_a, new_b, new_p)

    return output, new_state

RWKV 块结构

class RWKVBlock:
    """
    RWKV 的基本块结构

    与 Transformer Block 的对比:
    Transformer: LayerNorm → Attention → Add → LayerNorm → FFN → Add
    RWKV:        LayerNorm → TimeMix  → Add → LayerNorm → ChannelMix → Add
    """

    def __init__(self, dim, layer_id):
        # Time Mixing(替代 Attention)
        self.time_mix = TimeMixing(dim, layer_id)
        # Channel Mixing(替代 FFN)
        self.channel_mix = ChannelMixing(dim, layer_id)
        self.ln1 = LayerNorm(dim)
        self.ln2 = LayerNorm(dim)

    def forward(self, x, state):
        # Time Mixing(混合过去 token 的信息)
        dx, state = self.time_mix(self.ln1(x), state)
        x = x + dx

        # Channel Mixing(混合通道间的信息)
        dx = self.channel_mix(self.ln2(x))
        x = x + dx

        return x, state


class TimeMixing:
    """
    Token Shift + WKV
    对当前 token 与前一 token 做线性插值后使用
    """

    def __init__(self, dim, layer_id):
        self.mix_r = nn.Parameter(torch.ones(dim))  # 插值比例
        self.mix_k = nn.Parameter(torch.ones(dim))
        self.mix_v = nn.Parameter(torch.ones(dim))

    def forward(self, x, state):
        # Token Shift:当前 token 与前一 token 的加权平均
        x_prev = state.shift  # 前一 token
        xr = x * self.mix_r + x_prev * (1 - self.mix_r)
        xk = x * self.mix_k + x_prev * (1 - self.mix_k)
        xv = x * self.mix_v + x_prev * (1 - self.mix_v)

        r = torch.sigmoid(self.W_r(xr))
        k = self.W_k(xk)
        v = self.W_v(xv)

        wkv = compute_wkv(k, v, state)
        return r * wkv, new_state

各版本演进

RWKV-4（Eagle）

  • 引入基础 WKV 机制
  • 用 Token Shift 替代位置编码
  • 最大 14B 参数

RWKV-5（Eagle）

# RWKV-5: Multi-headed State
# 维护多个独立状态以提升表现力
class RWKV5_TimeMix:
    def __init__(self, dim, n_heads=8):
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        # 每个 head 拥有独立的 decay rate
        self.time_decay = nn.Parameter(torch.randn(n_heads, self.head_dim))

RWKV-6（Finch）

# RWKV-6: Data-dependent time decay
# 衰减率随输入动态变化(与 Mamba 的选择性机制类似)
class RWKV6_TimeMix:
    def forward(self, x, state):
        # 不再是固定 decay,而是依赖输入的 decay
        time_decay = self.W_decay(x)  # 每个输入的衰减率都不同!
        time_decay = torch.exp(-torch.exp(time_decay))

        # LoRA 风格的 decay 调制
        decay_lora = self.decay_lora_a(x)
        decay_lora = torch.tanh(decay_lora) @ self.decay_lora_b
        time_decay = time_decay + decay_lora

RWKV-7（Goose）— 2025 年 3 月

# RWKV-7: State Transition Matrix
# 用矩阵表示状态转移,实现更丰富的状态更新
class RWKV7_TimeMix:
    """
    RWKV-7 的核心创新:
    1. State Transition Matrix: 状态转移不再是标量,而是矩阵
    2. In-Context Learning 强化: 动态调整学习规则
    3. Improved Token Mixing: 更精细的 token 间信息流动
    """

    def forward(self, x, state):
        # 计算状态转移矩阵
        a = torch.sigmoid(self.W_a(x))  # (B, T, D)
        b = self.W_b(x)                 # (B, T, D)

        # 矩阵形式的状态更新
        # s_{t+1} = diag(a_t) @ s_t + k_t^T @ v_t
        k = self.W_k(x)
        v = self.W_v(x)

        # 状态:(B, H, D/H, D/H) 矩阵!
        for t in range(T):
            state = torch.diag(a[:, t]) @ state + \
                    k[:, t].unsqueeze(-1) @ v[:, t].unsqueeze(-2)

        return output, state

与 Transformer 的性能对比

按模型规模的性能(Pile dataset perplexity,↓ 越低越好):

| 模型规模  | Transformer | RWKV-4 | RWKV-6 | RWKV-7 |
|-----------|-------------|--------|--------|--------|
| 169M      | 17.2        | 18.1   | 17.4   | 17.0   |
| 430M      | 13.8        | 14.5   | 13.9   | 13.6   |
| 1.5B      | 11.2        | 11.8   | 11.3   | 11.0   |
| 3B        | 9.8         | 10.3   | 9.9    | 9.6    |
| 7B        | 8.5         | 9.0    | 8.6    | 8.3    |
| 14B       | 7.8         | 8.2    | 7.9    | 7.6    |

推理效率(token/秒,A100 GPU):

| 序列长度   | Transformer | RWKV-7  |
|-----------|-------------|---------|
| 1K        | 1000        | 1200    |
| 4K        | 800         | 1200    |
| 16K       | 300         | 1200    |
| 64K       | 50          | 1200    |
| 128K+     | OOM         | 1200    |

实战:使用 RWKV

在 HuggingFace 上使用

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载 RWKV-7 模型
model = AutoModelForCausalLM.from_pretrained(
    "RWKV/rwkv-7-world-3b",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(
    "RWKV/rwkv-7-world-3b",
    trust_remote_code=True
)

# 文本生成
prompt = "要在 Kubernetes 中实现 Pod 自动扩缩容"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
outputs = model.generate(
    **inputs,
    max_new_tokens=200,
    temperature=0.7,
    top_p=0.9
)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

用 RWKV Runner 本地运行

# 安装 RWKV Runner(GUI 工具)
git clone https://github.com/josStorer/RWKV-Runner
cd RWKV-Runner

# 或使用 ChatRWKV(CLI)
pip install rwkv

# 在 Python 中直接使用
python3 << 'EOF'
from rwkv.model import RWKV
from rwkv.utils import PIPELINE

model = RWKV(model='/path/to/RWKV-7-World-3B.pth', strategy='cuda fp16')
pipeline = PIPELINE(model, "rwkv_vocab_v20230424")

result = pipeline.generate(
    "请介绍一下韩国 AI 产业的现状。",
    token_count=200,
    temperature=0.8
)
print(result)
EOF

RWKV vs Mamba vs Transformer

| 特性                 | Transformer    | Mamba          | RWKV-7        |
|---------------------|----------------|----------------|---------------|
| 训练复杂度           | O(N²)         | O(N)           | O(N)          |
| 推理复杂度           | O(N) per token | O(1) per token | O(1) per token|
| 内存                | O(N) KV Cache  | O(1) 状态     | O(1) 状态     |
| 并行训练            | ✅             | ✅（scan）      | ✅（WKV）      |
| 长期依赖            | ✅ 强          | ⭕ 良好        | ⭕ 良好       |
| In-Context Learning | ✅ 强          | ⭕ 良好        | ✅ v7 中强化   |
| 实现复杂度           | 低             | 中（CUDA）     | 中（CUDA）    |
| 社区                | 庞大           | 成长中         | 活跃          |

局限与未来

当前局限

  1. 复杂检索任务: 与 Transformer 的 Full Attention 相比,在特定模式检索上较弱
  2. 依赖 CUDA kernel: 要达到最优性能需要定制 CUDA kernel
  3. 生态: 相较 Transformer,工具与库仍显不足

未来方向

  • 混合架构: 将 RWKV 与少量 Attention 结合
  • 硬件优化: 面向 Groq、Cerebras 等新芯片的优化
  • 多模态: 向视觉、音频等其他模态扩展

小测验

Q1. RWKV 中的 R、W、K、V 各自代表什么?

R: Receptance（接受门）,W: Weight（时间衰减）,K: Key,V: Value。

Q2. RWKV 在训练与推理时分别以哪种模式运行?

训练时以 Transformer 式的并行(parallel)模式一次性处理整个序列;推理时以 RNN 式的顺序(sequential)模式运行,每个 token 成本为 O(1)。

Q3. RWKV-6 引入的 Data-dependent time decay 是什么?

不再使用固定的衰减率,而是让衰减率随输入数据动态变化的机制。其思路与 Mamba 的选择性(Selective)机制类似。

Q4. RWKV-7 Goose 的核心创新是什么?

引入 State Transition Matrix,将状态转移用矩阵而非标量表示。由此实现更丰富的状态更新,并强化了 In-Context Learning。

Q5. 相较 Transformer,RWKV 最大的优势是什么?

无论序列长度多长,推理成本都保持恒定(O(1) per token)。即便在 128K+ token 下,也不会像 Transformer 那样出现 OOM。

Q6. Token Shift 机制的作用是什么?

将当前 token 与前一 token 做加权平均后混合,使得即便没有位置编码,也能在序列中传递位置信息。

Q7. RWKV 目前的局限之一是什么?

与 Transformer 的 Full Self-Attention 相比,在复杂模式检索或精确信息检索任务上,性能可能有所不及。

结语

RWKV 是一种颠覆"RNN 已死"这一成见的创新架构。它将 Transformer 的并行训练效率与 RNN 的推理效率结合起来,尤其在长序列处理与边缘设备部署上展现出优势。v7 Goose 引入 State Transition Matrix 后,与 Transformer 的性能差距已几乎消解。

参考资料

Discuss on Twitter • View on GitHub