1.1 Transformer 的二次复杂度瓶颈#

标准自注意力的计算复杂度为 O(n²)，这带来了两个核心问题：

1. 训练成本：长序列训练的时间和内存随长度二次增长
2. 推理效率：KV Cache 随序列长度线性增长，限制批处理大小

1.2 线性注意力的基本思想#

标准注意力的计算可以写为：

1
Attention(Q, K, V) = softmax(QK^T / √d) × V

其中 softmax 使得 Q 和 K 必须成对计算，导致 O(n²)。如果去掉 softmax（或用核函数近似），计算可以重排为：

1
LinearAttention(Q, K, V) = (Q × (K^T × V)) / (Q × (K^T × 1))

这样 K^T × V 可以先计算（O(n × d²)），然后再与 Q 相乘（O(n × d²)），总复杂度降为 O(n × d²)。当 d << n 时，这比 O(n²) 高效得多。

更重要的是，这种重排使得推理时可以维护一个累积状态（类似 RNN），实现 O(1) 的逐 Token 推理。

2.1 Token Shift#

RWKV 的第一个关键设计是 Token Shift（Token 移位）：将当前 Token 和前一个 Token 的线性插值作为输入。

1
# Token Shift 机制
2
# 对于时间步 t，输入不是单纯的 x_t
3
# 而是 x_t 和 x_{t-1} 的线性插值
4

5
def token_mixing(x_t, x_t_minus_1, ratio):
6
    """Token Shift：线性插值当前和前一个 Token"""
7
    return ratio * x_t + (1 - ratio) * x_t_minus_1

Token Shift 的作用是让模型在处理当前 Token 时直接获取前一个 Token 的信息，类似于 n-gram 模型中的上下文效果，但更加灵活。

2.2 时间衰减机制（Time Decay）#

RWKV 的第二个核心设计是时间衰减。在标准注意力中，所有位置的关系通过 softmax 归一化。RWKV 用可学习的衰减参数替代 softmax：

1
# 简化的 RWKV 注意力计算
2
# 每个通道有独立的衰减参数 w
3
w = channel_decay  # [hidden_dim]，可学习参数，0 < w < 1
4

5
# 累积状态随时间衰减
6
# state_t = w * state_{t-1} + k_t * v_t
7
# output_t = receptance(wkv_state_t)
8

9
# 等价于：对每个位置 j <= t
10
# weight_{t,j} = w^(t-j)  # 越远的位置权重越小

关键特性：

• 每个通道有独立的衰减参数 w，而非全局共享
• w 接近 1 表示”长记忆”（类似长距离注意力）
• w 接近 0 表示”短记忆”（只关注最近几个 Token）

2.3 WKV 计算#

RWKV 的核心计算被称为 WKV（Weighted Key Value），其 RNN 形式如下：

1
def rwkv_wkv(r, k, v, w, u, state):
2
    """
3
    RWKV 的 WKV 计算（单步）
4

5
    r: receptance，决定如何读取状态 [hidden_dim]
6
    k: key，用于更新状态 [hidden_dim]
7
    v: value，用于更新状态 [hidden_dim]
8
    w: time decay，衰减系数 [hidden_dim]
9
    u: time bonus，当前 Token 的额外权重 [hidden_dim]
10
    state: 累积的 KV 状态 [hidden_dim]
11
    """
12
    # 当前 Token 对输出的贡献（通过 u 加权）
13
    current = torch.exp(u) * k * v
14

15
    # 衰减后的历史状态 + 当前贡献
16
    wkv = state * torch.exp(-w) + current
17

18
    # 更新状态
19
    new_state = wkv
20

21
    # 通过 receptance 门控输出
22
    output = torch.sigmoid(r) * wkv
23

24
    return output, new_state

3.1 整体架构#

3.2 Time-Mixing 模块#

Time-Mixing 替代了 Transformer 的 Self-Attention：

1. 对输入应用 Token Shift，得到当前 Token 和前一个 Token 的混合
2. 通过线性投影生成 R、K、V、W、U 五个向量
3. 通过 WKV 计算得到输出
4. 输出经过门控（Receptance 门）和残差连接

3.3 Channel-Mixing 模块#

Channel-Mixing 替代了 Transformer 的 FFN：

1
def channel_mixing(x, x_prev, W1, W2):
2
    """
3
    Channel-Mixing 模块
4
    使用前一个 Token 的信息（通过 Token Shift）
5
    """
6
    # Token Shift
7
    x_mixed = token_shift(x, x_prev)
8

9
    # 非线性变换（类似 Square ReLU）
10
    hidden = torch.square(torch.relu(W1(x_mixed)))
11
    output = W2(hidden)
12
    return output

3.4 与 Transformer 的对比#

4.1 核心改进#

1. 多头公式化：类似 Transformer 的多头注意力，RWKV-5 将隐藏维度分为多个头，每个头有独立的参数
2. 数据依赖衰减：衰减参数 w 不再是固定的可学习参数，而是由输入数据动态生成
3. GroupNorm：在 WKV 计算中加入 GroupNorm，提高训练稳定性

1
# RWKV-5 的数据依赖衰减
2
# w 不再是固定参数，而是由输入 x 生成
3
w = softplus(time_decay_proj(x))  # 每步动态计算
4

5
# 这使得模型可以根据输入内容调整"记忆长度"
6
# 对需要长距离依赖的内容 → w 接近 1（长记忆）
7
# 对只需局部上下文的内容 → w 接近 0（短记忆）

5.1 核心改进#

1. Lerp-based Token Shift：用数据依赖的插值比替代固定的 Token Shift
2. 细粒度注意力路由：每个 Token 可以动态选择关注不同的历史信息
3. 改进的初始化：更合理的参数初始化策略

1
# RWKV-6 的数据依赖 Token Shift
2
# ratio 不再是固定的，而是由输入动态决定
3
ratio = sigmoid(lerp_proj(x_t))
4
x_shifted = ratio * x_t + (1 - ratio) * x_{t-1}
5

6
# 这使得模型可以根据当前内容决定
7
# 多少信息来自当前 Token，多少来自前一个 Token

5.2 模型规模#

RWKV-6 已训练了多种规模的模型：

7.1 与同规模 Transformer 的对比#

7.2 推理效率对比#

RWKV 的推理速度和内存占用不随序列长度变化，这是其相对于 Transformer 的核心优势。

8.1 RWKV Foundation#

RWKV 项目由 RWKV Foundation 维护，是一个活跃的开源社区：

1. 模型权重：所有规模模型权重完全开源（Apache 2.0）
2. 训练框架：提供完整的训练和微调代码
3. 多语言支持：社区贡献了中文、日文、韩文等多语言版本
4. 生态工具：WebUI、API 服务、量化工具等

8.2 社区贡献#

• RWKV-World：多语言版本，支持 100+ 种语言
• RWKV-Music：音乐生成模型
• RWKV-Visual：视觉语言模型探索
• ChatRWKV：类似 ChatGPT 的对话界面

8.1 Q1： RWKV 能完全替代 Transformer 吗？#

目前还不能。RWKV 在语言建模上接近同规模 Transformer，但在需要精确长距离依赖的任务上（如需要回溯全文的阅读理解）仍有一定差距。不过，对于大多数实际应用（对话、文本生成、代码补全），RWKV 已经足够好用。

8.2 Q2： RWKV 和 Mamba 应该选哪个？#

选择取决于场景：

• 需要成熟的生态和多语言支持 → RWKV
• 需要极致的硬件优化和学术支持 → Mamba
• 混合架构（Transformer + 线性层）→ 两者都可用
• 纯推理场景、超长序列 → 两者都适合

8.3 Q3： RWKV 的训练效率如何？#

RWKV 的训练可以像 Transformer 一样并行化（使用时间维度的并行扫描），但训练效率略低于 Transformer（约慢 10-20%），因为 RWKV 的计算模式对 GPU 不如标准注意力友好。RWKV 的优势主要在推理端。

8.4 Q4： RWKV 支持多模态吗？#

目前 RWKV 的多模态支持还在早期阶段。社区有 RWKV-Visual 等探索性项目，但还没有像 LLaVA 那样成熟的多模态方案。这是 RWKV 生态的一个待完善方向。

8.5 Q5： RWKV 如何处理超长上下文？#

RWKV 天然支持超长上下文，因为其推理状态是固定大小的。理论上，RWKV 可以处理无限长度的序列（状态不随序列增长）。但实际效果受限于时间衰减机制——太远的信息会被衰减到几乎为零。RWKV-6 的数据依赖衰减在一定程度上缓解了这个问题。