RWKV模型：线性注意力的演进

RWKV（pronounced “RwaKuv”）是一种将RNN与Transformer优势相结合的序列建模架构。它既能像Transformer一样并行训练，又能在推理时保持恒定的内存占用和O(1)时间复杂度。RWKV-7代号”Goose”代表了该系列的最新突破，引入了表达力更强的动态状态演化机制。

1. RWKV架构基础

1.1 核心设计哲学

RWKV的命名来自四个关键组件的首字母¹：

• Receptive field：接收域
• Weighted key-value：加权键值
• Kv-positional decay：位置衰减
• Value：值

其核心思想是：用线性注意力机制模拟RNN的归纳偏置，从而获得两者的最佳特性。

1.2 线性注意力机制

传统Transformer的自注意力计算为：

$$\text{Attention}(Q,K,V)=\text{softmax}(Q{K}^{\top })V$$

复杂度为 $O(T^2)$，其中 $T$ 是序列长度。

RWKV的线性注意力通过以下方式避免softmax的全局依赖：

$$\text{Attention}(x{)}_t=\frac{{\sum }_{i=1}^t{e}^{-(t-i)w}\cdot \exp (q_t^{\top }{k}_i)\cdot v_i}{{\sum }_{i=1}^t{e}^{-(t-i)w}\cdot \exp (q_t^{\top }{k}_i)}$$

其中 $w$ 是时间衰减因子，控制历史信息的重要性。

1.3 RWKV的时间混合模块

RWKV的核心是**时间混合（Time Mixing）**模块，用于融合不同时间步的信息：

RWKV Block (Layer t)
┌─────────────────────────────────────┐
│  x_t ──┬──→ [Time Mixing]            │
│        │                             │
│        └──→ [Channel Mixing]          │
│                                     │
│  Time Mixing:                       │
│  r_t = σ(W_r · x_t)                │
│  k_t = W_k · x_t                   │
│  v_t = W_v · x_t                   │
│                                     │
│  输出 = Σ_i^t [e^{-(t-i)w} · r_i·k_i·v_i] │
└─────────────────────────────────────┘

1.4 位置编码与衰减

RWKV使用指数衰减的位置编码：

$$w_t=\frac{u}{\sqrt{N}}+p_t$$

其中：

• $u$：可学习的衰减偏差
• $p_t$：位置嵌入
• $N$：模型维度

这种设计使得：

• 近期tokens权重高
• 远期tokens权重指数衰减
• 完全避免了三角函数位置编码

1.5 与Transformer的关键差异

特性	Transformer	RWKV
训练并行性	完全并行	完全并行
推理内存	$O(T)$ KV缓存	$O(1)$ 状态
推理时间	$O(T)$	$O(1)$
全局注意力	是	线性近似
RNN-like特性	否	是

2. RWKV-7 “Goose” 新特性

2.1 核心创新：动态状态演化

RWKV-7引入了**动态状态演化（Dynamic State Evolution, DSE）**机制，这是对之前delta规则的重要泛化²。

传统RWKV的更新规则：

$$\begin{array}{rl}{k}_t& =\sigma (W_k\cdot x_t+U_k\cdot {\text{rwkv}}_{t-1})\\ {\text{rwkv}}_t& =(1-k_t)\odot {\text{rwkv}}_{t-1}+k_t\odot x_t\end{array}$$

RWKV-7的动态状态演化：

$$\begin{array}{rl}{g}_t& =\sigma (W_g\cdot x_t+U_g\cdot {\text{state}}_{t-1})\text{// 门控}\\ {\alpha }_t& =\psi (W_{\alpha }\cdot x_t+U_{\alpha }\cdot {\text{state}}_{t-1})\text{// 上下文学习率}\\ {\delta }_t& =\varphi (W_{\delta }\cdot x_t+U_{\delta }\cdot {\text{state}}_{t-1})\text{// delta}\\ {\text{state}}_t& =(1-{\alpha }_t\odot g_t)\odot {\text{state}}_{t-1}+{\alpha }_t\odot g_t\odot {\delta }_t\end{array}$$

2.2 向量值门控

与标量门控不同，RWKV-7使用向量值门控（Vector-valued Gating）：

$$g_t\in {\mathbb{R}}^D\text{而非}{g}_t\in \mathbb{R}$$

这允许模型对不同维度施加不同程度的控制，增强表达能力。

2.3 上下文学习率

核心创新：每个token位置都有独立的上下文学习率 ${\alpha }_t$：

$${\alpha }_t:{\text{state}}_{t-1}\to [0,1{]}^D$$

这使得模型能够：

• 自适应地决定保留多少历史信息
• 在不同上下文中调整记忆强度
• 更灵活地进行信息选择性遗忘

2.4 放松的值替换规则

传统RWKV使用硬性的值替换：

$${\text{state}}_t=(1-k_t)\cdot {\text{state}}_{t-1}+k_t\cdot v_t$$

RWKV-7引入软性插值：

$${\text{state}}_t=\underset{\text{保留项}}{\underbrace{(1-{\alpha }_t\odot g_t)}}\odot {\text{state}}_{t-1}+\underset{\text{更新项}}{\underbrace{{\alpha }_t\odot g_t}}\odot {\delta }_t$$

这种设计更加平滑，避免了突变式的状态更新。

2.5 表达能力提升

RWKV-7的关键理论优势是能够识别所有正则语言²：

定理：RWKV-7可以执行状态追踪并识别所有正则语言，这在标准复杂性猜想下超越了Transformer的能力（TC⁰）。

这意味着RWKV-7能够：

• 维护精确的计数器
• 识别正则表达式定义的语言
• 执行复杂的状态转移逻辑

2.6 训练效率

尽管增强了表达能力，RWKV-7保持了完全可并行训练的特性：

• 前向传播：$O(T)$ 矩阵运算
• 反向传播：标准反向模式
• 与RWKV-6相同的计算图结构

3. 技术实现

3.1 RWKV-7架构参数

模型规模	参数量	多语言SOTA	英语SOTA
RWKV-7-168M	168M	基础	基础
RWKV-7-421M	421M	提升	提升
RWKV-7-1.47B	1.47B	新SOTA	持平
RWKV-7-2.9B	2.9B	新SOTA	持平

3.2 核心代码结构

class RWKV7Block(nn.Module):
    def __init__(self, dim, head_dim=64):
        super().__init__()
        # 时间混合参数
        self.time_mix = nn.Parameter(torch.zeros(1, 1, dim))
        self.key = nn.Linear(dim, dim)
        self.value = nn.Linear(dim, dim)
        self.receptance = nn.Linear(dim, dim)
        self.gate = nn.Linear(dim, dim)  # RWKV-7 新增
        
        # 动态状态参数
        self.alpha = nn.Linear(dim, dim)  # 上下文学习率
        self.delta = nn.Linear(dim, dim)  # delta
 
    def forward(self, x, state=None):
        B, T, C = x.shape
        # 动态状态演化
        g = torch.sigmoid(self.gate(x))
        alpha = torch.sigmoid(self.alpha(x))
        delta = self.delta(x)
        
        if state is None:
            state = torch.zeros_like(x[:, 0])
        
        # 更新规则
        new_state = (1 - alpha * g) * state + alpha * g * delta
        # ... 继续处理

3.3 与其他线性注意力的对比

架构	状态演化	门控	表达能力
Mamba	线性	标量	高
RetNet	固定	标量	中
RWKV-6	线性	标量	中
RWKV-7	动态	向量	高

4. 实践应用

4.1 部署场景

RWKV特别适合以下场景：

场景	优势
无限上下文	恒定内存，支持任意长度
实时生成	O(1)推理延迟
资源受限设备	内存效率高
流式处理	自然支持增量生成

4.2 使用示例

from rwkv import RWKVModel
 
model = RWKVModel("RWKV-7-2.9B")
 
# 无限上下文生成
state = None
for chunk in stream:
    output, state = model.forward(chunk, state)
    # state始终恒定大小

4.3 与Transformer的互补性

实践中，RWKV和Transformer可以互补使用：

• RWKV：处理长距离依赖的”记忆”
• Transformer：精细的全局推理

5. 理论分析

5.1 表达能力界限

RWKV-7的表达能力由其状态维度 $D$ 和非线性函数决定：

• 表达能力：理论上可以逼近任意连续函数
• 形式语言：可识别所有正则语言
• 时序建模：优于标准TC⁰复杂度类

5.2 与Transformer的关系

当RWKV-7的状态维度趋向无穷时：

$$\underset{D\to \infty }{\lim }\text{RWKV-7}\approx \text{Transformer}$$

但实际应用中，RWKV-7通过有限的维度实现了接近Transformer的表达能力。

6. 未来方向

1. 更大规模：继续扩展到更大参数规模
2. 多模态：扩展到视觉、音频等领域
3. 理论深化：更严格的表达能力分析
4. 硬件优化：针对性kernel加速

---

参考资料

Footnotes

1. Peng B, et al. RWKV: Reinventing RNNs for the Transformer Era. arXiv:2305.13048, 2023. ↩

2. Goldstein D, et al. RWKV-7 “Goose” with Expressive Dynamic State Evolution. arXiv:2503.14456, 2025. ↩ ↩²