RWKV：在 Transformer 时代重塑 RNN

2026/05/23 18:56:26

AI架构·12 min read

RWKV RNN Linear Attention Architecture LLM

背景

为什么需要 RWKV

Transformer 自 2017 年以来几乎统治了所有 NLP 任务，但其核心自注意力机制的计算和内存复杂度随序列长度呈 二次增长（ $$O(T^2d)$$ ），在长序列场景下成为致命瓶颈。循环神经网络（RNN）虽然具备线性复杂度的理论优势，但长期受制于梯度消失问题和训练阶段无法并行化的缺陷，始终难以 scale 到十亿参数级别。

RWKV（Receptance Weighted Key Value）的核心命题是：能否在不引入近似的前提下，将 Transformer 的并行训练能力与 RNN 的线性推理效率结合起来？

核心问题

RNN 与 Transformer 的各自瓶颈

RNN 的困境

梯度消失：长序列训练时反向传播路径过长，信号衰减严重
无法并行训练： $h_t = \alpha_t \odot h_{t-1} + (1-\alpha_t) \odot x_t$ 每一步依赖前一步，time dimension 必须串行
可扩展性差：难以堆叠深层网络，限制了模型容量

Transformer 的困境

二次复杂度：自注意力计算 $$O(T^2d)$$ ，长序列时计算和内存成本爆炸
推理效率低：每个新 token 生成都需要重新计算对所有历史 token 的注意力
内存随序列增长：KV Cache 随序列长度线性增长

寻址记忆 vs 联想记忆

论文提出了一个深刻的类比：Transformer 是一种 寻址记忆（Addressable Memory）——通过 Query 向历史位置发起精确查询；而 RWKV 是一种 联想记忆（Associative Memory）——通过时间衰减权重对历史信息进行加权求和，无需显式寻址。这从根本上改变了对"记忆"的理解方式。

前身工作

从 AFT 到 RWKV

RWKV 并非凭空出现，其设计直接继承了 AFT（Attention-Free Transformer） 的核心思想。AFT 将传统注意力中的 Query 替换为一个可学习的偏置矩阵 $w_{t,i}$ ，提出：

\text{Attn}^+(W, K, V)_t = \frac{\sum_{i=1}^{t} e^{w_{t,i}+k_i} v_i}{\sum_{i=1}^{t} e^{w_{t,i}+k_i}}

这一公式绕过了对 Q 的依赖，使注意力计算不再需要 $$O(T^2)$$ 的 token-pair 交互。但 AFT 的 $w_{t,i}$ 是全局标量，缺乏对不同时间尺度的建模能力。

RWKV 的关键突破在于将 $w_{t,i}$ 重新参数化为一个 channel-wise 的时间衰减向量：

w_{t,i} = -(t-i)w

其中 $$w$$ 是一个与隐藏维度同大小的可学习向量。这意味着每个通道（channel）都有独立的时间衰减率——某些通道快速遗忘，另一些通道保留长程依赖。这种设计让模型能够同时捕获短程和长程依赖，无需引入任何近似。

核心机制

WKV Operator：线性注意力的精确形式

将上述参数化代入 AFT 公式，得到 RWKV 的核心——WKV Operator：

\text{Attn}^+(W, K, V)_t = \frac{\sum_{i=1}^{t-1} e^{-(t-1-i)w + k_i} v_i + e^{k_t} v_t}{\sum_{i=1}^{t-1} e^{-(t-1-i)w + k_i} + e^{k_t}}

这里引入的四个向量：

R（Receptance）：控制过去信息接受程度的向量，充当"输入门"
W：位置权重衰减向量，可训练的模型参数
K（Key）：键向量，类似于传统注意力中的 K
V（Value）：值向量，类似于传统注意力中的 V

计算本质：分母中的

\sum_{i=1}^{t-1} e^{-(t-1-i)w + k_i}

和分子中的加权求和项可以递归展开为对数空间中的累积量。这意味着推理时不需要存储完整的注意力权重历史，只需维护一个固定大小的累积状态

a_{t-1}

，从而实现 $$O(1)$$ 推理时间。

核心机制

Token Shift：时间与通道的混合

RWKV 另一个关键设计是 Token Shift——将当前 token $$x_t$$ 与前一个 token $x_{t-1}$ 做线性插值。这一操作被应用于两个地方：

Time-Mixing（时间混合）

Time-Mixing 用 Sigmoid 门控控制当前 token 与前序 token 的混合比例：

\text{TimeMix}(x_t, x_{t-1}) = (1-\sigma(r_t)) \odot x_{t-1} + \sigma(r_t) \odot x_t

其中 $$r_t$$ 是从输入计算出的 receptance 向量。这使得模型能够自适应地决定在每个时间步保留多少"历史记忆"，相当于为每个通道配备了独立的信息保留门。

Channel-Mixing（通道混合）

Channel-Mixing 采用 ReLU 门控：

\text{ChannelMix}(x_t, x_{t-1}) = \text{ReLU}(r_t) \odot x_{t-1} + (1-\text{ReLU}(r_t)) \odot x_t

其作用类似于 Transformer 中的 FFN 层，但通过 Token Shift 引入了局部上下文信息。

架构

完整 Block 与复杂度对比

一个 RWKV Block 的前向流程为：

Time-Mixing：通过 Token Shift 混合 $$x_t$$ 和 $x_{t-1}$ ，计算 $$r_t, k_t, v_t$$
WKV Operator：利用时间衰减权重更新累积状态
Channel-Mixing：通过 FFN-like 操作输出

模型	推理时间	推理空间
Transformer	$$O(T^2d)$$	$$O(Td + d^2)$$
Linear Transformers	$$O(Td^2)$$	$$O(Td + d^2)$$
AFT-local	$$O(Tsd)$$	$$O(Td)$$
RWKV	$$O(Td)$$	$$O(d)$$

RWKV 是目前唯一达到 $$O(Td)$$ 时间 和 $$O(d)$$ 空间 的精确线性注意力模型（无任何近似假设）。

训练

Scaling Law 与模型族

论文训练了 45 个不同规模的 RWKV 模型，发现 RWKV 遵循与 Transformer 相同的 log-log 线性 Scaling Law：

Pareto 最优拟合点： $$R^2 = 0.994$$
外推到 10× 计算量： $$R^2 = 0.875$$ （依然优异）

这推翻了"RNN 无法遵循 Transformer Scaling Law"的固有认知。

模型规模	参数量	单 Token FLOPs
RWKV-169M	1.69 亿	$2.985 \times 10^9$
RWKV-430M	4.3 亿	$7.393 \times 10^9$
RWKV-1B	10 亿	$1.437 \times 10^{10}$
RWKV-7B	70 亿	$7.393 \times 10^9$
RWKV-14B	140 亿	$1.415 \times 10^{10}$

RWKV-14B 是迄今为止训练的最大 dense RNN。训练配置：context length 1024 tokens（通过渐进式 finetuning 扩展到 8192）、bfloat16 精度、Adam 优化器（无 weight decay）、指数衰减学习率。

长上下文扩展

通过渐进式 context length finetuning（1024→2048→4096→8192），RWKV 能够有效利用长上下文信息——测试损失随 context 长度增加而持续下降，证明模型确实习得了长程依赖建模能力。

实验

评估结果

NLP Benchmarks（12 项任务）

在 ARC（Easy/Challenge）、BoolQ、COPA、HellaSwag、LAMBADA、OpenBookQA、PIQA、ReCoRD、SciQ、Winogrande 等基准上，RWKV 与同规模、同训练 tokens 量的 Pythia、OPT、BLOOM 进行了 FLOP-matched 对比。

结果：RWKV 在 12 项 NLP 任务上与同等规模 Transformer 性能持平。

Long-Range Arena

在序列长度 1K~16K 的 LRA 基准上，RWKV 在 5 个数据集上排名仅次于 S4 模型，显著优于其他线性注意力方法。

推理效率

文本生成累计时间曲线清晰显示：Transformer 呈二次增长，而 RWKV 呈线性增长。14B 参数的 RWKV 在 CPU 和 GPU 上的推理速度均优于同规模 Transformer。

Prompt Engineering 敏感性

论文发现 RWKV 对 prompt 格式敏感。调整 prompt 顺序后，某项任务的 F1 从 44.2% 提升至 74.8%。

深层解读

为什么 RWKV 的设计是合理的

1. Channel-wise Decay 的表达能力

与 S4/Mamba 等 State Space Model 不同，RWKV 的时间衰减以向量形式作用于每个通道。这意味着不同频率的信息（高衰减通道捕获短程依赖、低衰减通道捕获长程依赖）可以被模型自主选择，而非由单一的系统动态方程决定。这增加了表达能力，但也意味着放弃了一些理论上的最优性质（如 S4 的 HiPPO 初始化保证）。

2. 精确线性注意力（无近似）

与 Linear Transformers 的核函数近似（random feature decomposition）不同，RWKV 的线性注意力是精确的。代价是 $w_{t,i}$ 的结构约束（必须是时间衰减形式），但这换来了可解释性和稳定的长程建模能力。

3. 联想记忆的归纳偏置

将 Transformer 的寻址记忆转化为联想记忆，其归纳偏置更接近人类认知中的"联想"机制——相关概念自动被唤醒，而非精确地址调用。这可能是 RWKV 在某些任务上对 prompt 格式更敏感的原因：信息的组织方式本身成为了影响检索效果的因素。

局限与未来

已知限制与开放问题

时间衰减函数固定：当前 $e^{-(t-i)w}$ 形式简单，可能无法捕获更复杂的时间依赖模式
初始化敏感：论文提到未来工作将探索更好的初始化策略
训练效率：虽然推理是 O(1)，但训练阶段的并行扫描仍未实现最优（计划 $O(B\log T d)$ ）
多模态扩展：作为纯序列模型，尚未探索视觉、语音等多模态场景

总结

核心贡献与影响

RWKV 在 2023 年的贡献不仅是一个新架构，更是一种重新审视 RNN 与 Transformer 关系的思路转变：

它证明了线性注意力可以在无近似的前提下实现，且在十亿参数规模上仍具竞争力
它揭示了寻址记忆 vs 联想记忆这一对称结构在序列建模中的深层含义
它打破了"RNN 无法 Scaling"的刻板印象，Scaling Law 实验提供了坚实的实证支撑

作为后续，RWKV 社区发展出了多代版本（RWKv4、RWKV-5、RWKV-6），在预训练和微调层面持续优化，并在多模态方向（视觉、语言、语音）展开探索。