RWKV模型

RWKV

近年来，基于Transformer架构的大规模语言模型在自然语言处理领域取得了显著成功，但其自注意力机制存在内存和计算复杂度随序列长度平方级增长的问题，限制了其在长序列任务中的效率。同时，传统RNN虽然具有线性复杂度，但在训练并行性和长程依赖建模方面存在不足。

在此背景下，RWKV（Receptance Weighted Key-Value）架构被BlinkDL提出，旨在结合Transformer的高效训练能力（可并行化）和RNN的推理效率。RWKV-7 "Goose"作为该系列的最新进展，通过恒定内存占用和每token恒定推理时间的设计，解决了传统Transformer的长序列处理瓶颈。

2023 年 9 月 20 日，RWKV 开源项目正式加入 Linux 开源基金会。如今的 RWKV 项目是一个开源的非盈利组织，隶属于 Linux 基金会，部分算力曾由赞助商支持。

一、RWKV 模型简介

RWKV 是一个融合了 RNN 和 Transformer 优势的新型语言模型架构。
名称含义：
- R：Receptance，作为过去信息的接受程度的接受向量
- W：Weight，位置权重衰减向量，可训练的模型参数
- K & V：类似传统注意力机制中的 Key 和 Value 概念
核心特点：
- 拥有 Transformer 的并行训练能力
- 拥有 RNN 的线性推理效率

二、RWKV 模型结构

RWKV Block（基本结构）

LayerNorm

TimeMix
- 输入：x
- 拆分为：[math]\displaystyle{ x_k, x_v, x_r }[/math]
- 加入时间加权参数（Python代码）：

decay = exp(-exp(time_decay))
k = x_k * time_mix_k + x_k_prev * (1 - time_mix_k)

数学公式形式：

Decay定义：

[math]\displaystyle{ \mathrm{decay} = e^{-e^{\mathrm{time\_decay}}} }[/math]

k计算：

[math]\displaystyle{ k = x_k \times \mathrm{time\_mix\_k} + x_k^{\mathrm{prev}} \times \left(1 - \mathrm{time\_mix\_k}\right) }[/math]

ChannelMix
- 类似 Transformer 的 FFN 模块。
- 激活函数通常使用 ReLU，公式如下：

[math]\displaystyle{ \mathrm{output} = \mathrm{ReLU}(W_1 \times x + b_1) \times W_2 + b_2 }[/math]

注意力层简化表达

[math]\displaystyle{ x_{t+1} = \sigma(Rx_t) \cdot \frac{\exp(Kx_t) \cdot (Vx_t) + \exp(W) \cdot a_t}{\exp(Kx_t) + \exp(W) \cdot b_t} }[/math]

核心机制代码

def ref_fwd(r, w, k, v, a, b):
    r = r.view(B, T, H, N)
    k = k.view(B, T, H, N)
    v = v.view(B, T, H, N)
    a = a.view(B, T, H, N)
    b = b.view(B, T, H, N)
    w = torch.exp(-torch.exp(w.view(B, T, H, N)))
    out = torch.zeros((B, T, H, N), device=DEVICE)
    state = torch.zeros((B, H, N, N), device=DEVICE)
 
    for t in range(T):
        kk = k[:, t, :]
        rr = r[:, t, :]
        vv = v[:, t, :]
        aa = a[:, t, :]
        bb = b[:, t, :]
        sab = torch.einsum('bhik,bhk,bhj->bhij', state, aa, bb)
        state = state * w[: , t, :, None, :] + sab + torch.einsum('bhj,bhi->bhij', kk, vv)
        out[:, t, :] = torch.einsum('bhj,bhij->bhi', rr, state)
 
    return out.view((B, T, C))

三、RWKV架构特性

推理

类似 RNN，每步只依赖上一步状态
推理内存 O(1)，不需上下文缓存

训练

并行训练兼容 Transformer-style 训练器
对时间参数稳定性要求高（需梯度裁剪、精调）

RWKV 架构由一系列堆叠的残差块组成，每个残差块由具有循环结构的时间混合（timemixing）和通道混合（channel-mixing）子块组成，该循环通过将当前输入和上一时间步的输入之间进行线性插值来实现（在 RWKV-4 架构论文中，这一过程被称为 token shift）。RWKV 6 通过借鉴 LoRA 技术优化了 token shift 过程，使其 RWKV4/5 的简单线性插值（lerp）变成了数据依赖的、动态的线性插值（ddlerp）。

根据对不同模型的推理复杂度进行比较，Transformer 的时间复杂度为：O (T^2)、空间复杂度为：O (T^2)，因此推理速度将会越来越慢，也越来越耗内存。而 RWKV 的时间复杂度为：O(T)、空间复杂度为 O(1)。RWKV 大模型通过对计算流程的优化，实现了恒定的推理速度，极大地减少了推理过程中的时间消耗。

此外，RWKV 架构设计显著降低了显存占用，使得模型即使在标准配置的 CPU 或非专业级的 GPU上也能高效运行，无需依赖昂贵或高端的计算硬件。这一突破性的进展使得大规模深度学习模型不再受限于特定的硬件平台，拓宽了应用范围。

四、RWKV 架构演进时间线

版本	时间	核心特性/改进	备注
v1	2021 年底	初代结构，融合 Attention 和 RNN	实验性版本，训练不稳定
v2	2022 年初	TimeMix 和 ChannelMix 雏形初步支持并行训练	性能有明显提升
v3	2022 年中	明确使用 `time_decay` 机制引入 `time_mix_k` 等精细控制	架构逐渐稳定，接近 Transformer 效果
v4	2023 年初	优化 ChannelMix（ReLU 替换 SiLU）更强的泛化能力	支持更大规模模型训练
v5	2023 底 / 2024 初	正式提出 “Transformer as RNN” 理念可用标准 Transformer 训练器性能与效率大幅提升	当前主力架构，广泛用于 RWKV-LM 系列
v6	2024 年初	引入动态时间衰减机制（Dynamic Time Decay）增强对长上下文的记忆能力	进一步提升语言建模的表现
v7	2024 年中	模块化设计支持多模态任务（文本、图像、音频等）更高效的推理优化	扩展应用场景，迈向通用 AI

五、RWKV 与其他模型对比

特性	Transformer	RNN	RWKV
并行训练能力	✅	❌	✅
推理内存效率	❌ (O(n))	✅ (O(1))	✅ (O(1))
长序列建模能力	中等	较弱	强
上下文可控性	较强	差	中-强
参数共享压缩性	中等	强	强

六、RWKV 社区与生态

主要贡献者

BlinkDL（中国开发者，RWKV 创始人）

开源项目

RWKV-LM：主项目
RWKV.cpp：轻量 C++ 实现，适合本地部署
RWKV-Runner：便捷部署和体验框架
RWKV-Vision：拓展到图像、多模态领域

七、RWKV 架构总结

优势

资源效率高：运行和训练时对 VRAM、CPU、GPU 的需求显著低于 Transformer 架构。
计算量低：相较于处理较大上下文的 Transformer，RWKV 的计算需求低 10 倍到 100 倍。
线性扩展性强：上下文长度可线性扩展，Transformer 则为二次扩展（[math]\displaystyle{ O(n^2) }[/math]）。
答案质量高：在文本生成质量和泛化能力方面可媲美 Transformer 架构。
多语言能力优异：训练数据覆盖中文、日文等多种语言，在多语言任务中表现优于多数开源模型。

不足

提示词格式敏感：生成效果对 prompt 的格式非常敏感，稍微不同的提示可能导致结果差异很大。
任务顺序依赖性强：由于架构本身的递归特性，对于需要回顾历史信息的任务较弱，需要合理组织输入顺序（例如先给任务，再给内容）。