RWKV：结合 RNN 效率与 Transformer 表达力的架构

Attention(Q,K,V) = softmax(QK^T / √d_k) V

- Q（Query）、K（Key）、V（Value）由输入线性变换得到。
- softmax 操作使得每个查询需要与所有键计算相似度，形成 n×n 的注意力矩阵，导致 O(n²d) 计算和 O(n²) 存储。

### 2.2 线性注意力思想

线性注意力的核心是**将 softmax 替换为可拆分的核函数**，使得计算顺序可以改变：

Attention(Q,K,V) ≈ (φ(Q) φ(K)^T) V = φ(Q) (φ(K)^T V)

这里 φ(·) 是一个逐元素特征映射。通过先计算 K^T V（大小为 d×d，与序列长度无关），再乘 Q，复杂度降至 O(nd²)。当 d ≪ n 时，效率大幅提升。但早期线性注意力模型的表达能力往往不及 softmax 注意力。

### 2.3 RWKV 的 WKV 算子：带指数衰减的线性注意力

RWKV 并没有直接套用 φ(Q)φ(K)，而是设计了一个**受遗忘门控制的加权累积**机制，相当于一种状态空间模型。其核心递推公式可视为：

st = e^{-w} * st-1 + k_t^T v_t output_t = q_t * st

其中 w 是可学的通道级衰减速率（通过指数函数保证为正），决定模型对历史信息保留或遗忘的程度。训练时，得益于巧妙的并行化设计（使用一种 **“low-rank + 线性注意力”** 的组合），可以一次性算出全序列输出，避免串行循环。

## 3. RWKV 架构详解

RWKV 模型由多个相同的基本块堆叠而成，每个块包含时间混合（Time-Mixing）和通道混合（Channel-Mixing）两个子层，借鉴了 Transformer 的残差连接和前置 Norm，但彻底抛弃了注意力矩阵。

### 3.1 时间混合（Token Mixing）

作用类似自注意力，负责在序列维度上混合信息。给定输入 x_t（归一化后），时间混合计算如下：

- 接收向量：`r_t = W_r * x_t`  （用于接收最终信号）
- 键和值：`k_t = W_k * x_t`，`v_t = W_v * x_t`
- 查询和衰减：`q_t = W_q * x_t`，同时计算出位置相关的**动态衰减** `w_t`（由参数和位置嵌入共同决定）
- WKV 计算：通过递推累积 `a_t` 和 `b_t`（类似分子分母），或训练时直接用并行公式：

wkv_t = (sum_{i=1}^{t-1} e^{-(w_t + w_i)} * k_i^T v_i + e^{u} * k_t^T v_t) / (sum_{i=1}^{t-1} e^{-(w_t + w_i)} * k_i^T + e^{u} * k_t^T)

这可以理解为对过去令牌的加权平均，其中权重由相对位置和学习到的衰减控制。最终输出由 `r_t` 对 `wkv_t` 进行门控。

### 3.2 通道混合（Channel Mixing）

类似 Transformer 的 FFN，逐位置操作。但引入了 RNN 风格的遗忘门：

- 接收：`r'_t = W'_r * x_t`
- 键：`k'_t = W'_k * x_t`
- 值：`v'_t = ReLU(k'_t)`^2  （平方 ReLU 激活）
- 输出：`r'_t * (W'_v * v'_t)`，通过残差连接相加。

这种设计让每个令牌独立地混合自身特征，避免了时间维度计算开销。

## 4. 为什么 RWKV 能做到高效？

- **推理阶段**：序列长度为 n 时，Transformer 缓存前文 K/V 矩阵需要 O(n) 内存，且每次生成步骤计算注意力需要 O(n) 时间。RWKV 以固定大小的状态向量 s_t 递推，每步时间 O(d²)，内存 O(d²)，两者均与序列长度无关，非常适合流式处理与边缘设备。
- **训练阶段**：虽然递推形式便于推理，但训练时采用 **“teacher forcing”** 模式。RWKV 利用数学变换将递推过程转化为一种**并行扫描**（或称为“合并前缀和”），可充分利用 GPU 并行度，训练速度接近 Transformer，且复杂度为 O(nd²)。
- **状态大小固定**：与 Transformer-XL 等缓存机制不同，RWKV 的状态是固定维度的向量，不会随时间增长，因此可以轻松处理无限长上下文（理论上），实际训练中支持长达数万 token 的上下文。

## 5. RWKV 与其他模型的对比

| 特性 | Transformer | LSTM/GRU | RWKV |
|------|------------|-----------|------|
| 训练并行性 | 高（自注意矩阵并行） | 低（必须串行展开） | 高（并行化 WKV） |
| 推理复杂度 | O(n²) 逐 token 增长 | O(1) 每步，但长程受限 | O(1) 每步，长程建模强 |
| 长距记忆机制 | 位置编码 + 全连接注意力 | 门控记忆单元 | 指数遗忘 + 可学衰减 |
| 显存占用（推理） | 缓存 K/V 共 2nd | 隐藏状态 ~ d² | 状态向量 ~ d² |
| 表达力瓶颈 | 注意矩阵稠密，成本高 | 递推压缩历史，可能丢失信息 | 加权和形式，信息保留更完整 |

## 6. 动手理解：核心公式的 Python 伪代码

以下展示训练时并行计算 WKV 的简化版（假设未使用动态位置衰减，仅用固定衰减 w）：

```python
import torch

def wkv_parallel(q, k, v, w, u):
    # q, k, v: (batch, seq_len, dim)
    # w: 遗忘速率 (seq_len,)
    # u: 额外偏置
    B, T, C = q.shape
    k = k.float()
    v = v.float()

    # 递归形式需要变为并行扫描
    # 这里只是概念示意，实际涉及复杂前缀扫描
    output = torch.zeros_like(q)
    s = torch.zeros(B, C, C)
    for t in range(T):
        # 真实实现会调用高效CUDA kernel
        s = torch.exp(-w[t]) * s + k[:, t:t+1].transpose(-1, -2) @ v[:, t:t+1]  # (C, C)
        num = (q[:, t:t+1] @ s).squeeze(1) + torch.exp(u) * k[:, t] * v[:, t]
        den = (q[:, t:t+1] @ s.expand(B, C, 1)).squeeze(1).norm(2) # 简化
        output[:, t] = num / den
    return output