RWKV模型

RWKV （读作 RWaKuV）是一种具有 GPT 级大型语言模型（LLM）性能的 RNN，也可以像 GPT Transformer 一样直接训练（可并行化）。

RWKV 结合了 RNN 和 Transformer 的最佳特性：出色的性能、恒定的显存占用、恒定的推理生成速度、"无限" ctxlen 和免费的句嵌入，而且 100% 不含自注意力机制。

---

RWKV结构示意图

RWKV 模型简介

• RWKV:
  • 近年来，基于Transformer架构的大规模语言模型在自然语言处理领域取得了显著成功，但其自注意力机制存在内存和计算复杂度随序列长度平方级增长的问题，限制了其在长序列任务中的效率。同时，传统RNN虽然具有线性复杂度，但在训练并行性和长程依赖建模方面存在不足。在此背景下，RWKV（Receptance Weighted Key-Value）架构被BlinkDL提出，旨在结合Transformer的高效训练能力（可并行化）和RNN的推理效率。
  • RWKV-7 "Goose"作为该系列的最新进展，通过恒定内存占用和每token恒定推理时间的设计，解决了传统Transformer的长序列处理瓶颈。
  • 2023 年 9 月 20 日，RWKV 开源项目正式加入 Linux 开源基金会。如今的 RWKV 项目是一个开源的非盈利组织，隶属于 Linux 基金会，部分算力曾由赞助商支持。
• 名称含义：
  • R：Receptance，作为过去信息的接受程度的接受向量
  • W：Weight，位置权重衰减向量，可训练的模型参数
  • K & V：类似传统注意力机制中的 Key 和 Value 概念
• 核心特点：
  • 拥有 Transformer 的并行训练能力
  • 拥有 RNN 的线性推理效率

---

RWKV 模型结构

RWKV 是 RNN（循环神经网络） 的一个变体。在深入探讨 RWKV 之前，我们有必要先了解 RNN 的基本原理，特别是其 隐藏状态（Hidden State） 的概念。

RNN 架构与隐藏状态

循环神经网络（RNN, Recurrent Neural Network） 是一种广泛应用于序列数据（如文本、语音、时间序列等）的深度学习模型。

RNN 的核心机制：隐藏状态

RNN 在运行过程中会维护一个 隐藏状态（Hidden State），可以将其类比为 RNN 的“记忆”或“心理状态”。

• 就像人类在思考时，大脑会保留与当前任务最相关的“关键信息”一样，RNN 的隐藏状态存储了网络在处理序列时的上下文信息。
• 每当 RNN 处理一个新的输入（如一个 token），它会结合 当前隐藏状态 计算输出，并更新自身的隐藏状态，以便影响后续的计算。

RNN 的工作流程

1. 输入处理：RNN 依次接收序列中的每个 token（如单词或字符）。
2. 状态更新：对于每个 token，RNN 结合 当前隐藏状态 计算输出，并生成 新的隐藏状态。
3. 循环反馈：更新后的隐藏状态会被传递到下一步，用于处理下一个 token。
4. 输出预测（可选）：如果有任务需求（如语言建模），RNN 可以在每一步生成预测结果（如下一个可能的 token）。

这一过程不断循环，直到整个序列处理完毕或达到终止条件，如下图所示：

RNN结构示意图

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))

• Delta Rule

RWKV-7 中的 广义 Delta Rule 是其架构的核心创新之一，是对传统神经网络中 Delta Rule 的扩展和优化。以下是其核心要点：

传统 Delta Rule 的基础 Delta Rule 是神经网络中的经典学习规则，通过计算目标输出与实际输出的差异（即误差 delta），调整模型参数以减少误差 。这一规则通常用于监督学习中的梯度下降优化。 广义 Delta Rule 的扩展 RWKV-7 将 Delta Rule 从标量形式推广到向量形式，并引入动态状态演化机制。具体来说： 向量化门控（Vector-Valued State Gating） ：原始 Delta Rule 中的衰减系数 a 被定义为向量而非标量，允许每个状态通道独立控制信息的保留与更新，从而增强模型对状态信息的精细操控能力 。 动态状态调整 ：通过结合输入序列的历史状态与当前误差信号，广义 Delta Rule 能够更高效地调整状态表示，实现长序列依赖建模 。

---

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))
长序列建模能力	中等	较弱	强
上下文可控性	较强	差	中-强
参数共享压缩性	中等	强	强

优势

• 运行和训练时资源使用更低（VRAM、CPU、GPU 等）。
• 与较大上下文的 Transformer 相比，计算需求低 10 倍到 100 倍。
• 支持以线性方式扩展到任何上下文长度（Transformer 以二次方扩展）。
• 在答案质量和泛化能力方面的表现与 Transformer 架构一样好。
• RWKV 模型的训练数据包括英文以外的其他语言（例如中文、日文等），多语言能力比大多数现有的开源模型要好。

不足

• RWKV 基底模型对提示词（prompt）的格式非常敏感，提示词的格式对生成结果影响较大。
• 出于架构设计，RWKV 模型在需要回顾的任务上较弱，因此需要恰当地对提示词进行排序。比如先给模型提供任务指令，再提供需要执行任务的材料文本。

---

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