在 Transformer 之前，自然语言理解（NLP）领域的主流算法是 Recurrent Neural Network (RNN) 和它的增强 Long Short Term Memory (LSTM)。

RNN/LSTM 的工作方式跟状态机有点类似，它内部有个隐状态（hidden state），根据这个隐状态和当前输入 token 共同计算得到输出 token，并根据当前输入的 token 来更新这个隐状态。随着不断的处理输入 token，这个隐状态相当于记录了过去看到过的 token 信息，也就是对过去序列的”记忆“。在更新隐状态的时候，我们需要有机制来控制哪些历史信息需要保留，哪些则需要遗忘，这个机制通过“forget gate”, “include gate”来实现。

典型的 RNN/LSTM 算法可以用下面一组公式来描述：

• 公式中的 σ_g 等函数，都是非线性函数，例如 sigmoid
• f_t: forget gate，控制哪些历史信息需要遗忘 (1)
• i_t: include gate，控制哪些历史信息需要保留 (2)
• o_t: 当前时刻 t 的输出，根据当前输入 x_t 和上一时刻的隐状态 h_t-1 计算出 (3)

为什么说 RNN 训练很难并行呢？我画了一个前后两个时刻各状态的依赖图来帮助理解：

可以看到，计算下一个隐状态 h_t 需要经过复杂的依赖关系，其中包括对上一个状态的多个嵌套的非线性依赖，而这些非线性函数又无法通过矩阵运算来表达。这意味着当我们训练时，面对一整个序列的 token，只能从前往后一个个位置计算隐状态，而不能像 Transformer 那样，通过矩阵运算把序列中每个位置的隐状态并行计算出来。

再看个简单的例子，假如我们的序列有 4 个输入 x1, x2, x3, x4，每个位置的隐状态就是当前输入的累加，也就是上一个隐状态与当前输入之和：

h0 = 0
h1 = h0 + x1 = x1
h2 = h1 + x2 = x1 + x2
h3 = h2 + x3 = x1 + x2 + x3
h4 = h3 + x4 = x1 + x2 + x3 + x4

这种情况，前后状态之间是线性依赖，我们只需要构造一个下三角矩阵：

1 0 0 0
1 1 0 0
1 1 1 0
1 1 1 1

把输入向量 [x1, x2, x3, x4] 和这个矩阵相乘，就可以把每个位置的隐状态并行计算出来。GPU 特别擅长这种矩阵运算，因此这种情况训练效率就很高。

但是，如果我们的隐状态和上一个状态是非线性依赖的，例如 h(t) = sigmoid(h(t-1)) + x(t) :

h0 = 0
h1 = sigmoid(h0) + x1
h2 = sigmoid(h1) + x2
h3 = sigmoid(h2) + x3
h4 = sigmoid(h3) + x4

像这种情况，我们就没办法用矩阵来并行计算每个位置的隐状态了，只能从 h1 开始一个个计算。RNN 训练难以并行，关键就在这里（我的理解）。

Transformer 的注意力机制可以用下面这个公式来表达（关于 Transformer 算法的细节已经有很多论述，这里就不重复了）：

上面这个 Softmax 展开后，可以写成这样的迭代形式：

也就是说，在位置 t 的注意力，是 Attn(Q,K,V)_t ，这是一个大小为 T 的向量，因此整个 Attn(Q,K,V)给出的是一个 TxT 的矩阵（T是序列长度）。 (这里 7-8 的展开细节，以后会另写一篇来讨论)

RWKV 的灵感来自于 2021 年 Apple 的一篇论文『Attention Free Transformer』（AFT），它里面把注意力表达成这种形式：

其中 w 是一个 TxT 的矩阵，通过训练学习得到。

RWKV 里面，把每个 w_t,i 变成了一个 channel-wise time decay vector，与位置的乘积：

其中的 w 是一个大小为 d 的向量（d是 channel 数，也就是 Transformer 模型训练代码里面 B,T,C 里的 C）。

RWKV 算法的核心思想，其实就这些，后面我们会看到它的应用和带来的变化。

RWKV 模型的架构，和 Transformer 非常类似，也是由多个 RWKV block 组成，最后加一个 language modeling head 输出下一个 token 的分布概率。每个 RWKV block 内部，有一个 channel mix 和一个 time mix 模块。

RWKV block 内部的 Time Mixing 和 Channel Mixing 模块：

RWKV block 的 Time Mixing 模块，可以用下面的一组公式来描述：

首先，输入经过 LayerNorm 后，将当前位置和前一个位置的输入按权重做一个 mix，然后分别投影成 R, K, V (公式 11,12,13)。可以看到，这里就是在投影前把历史信息 mix 起来了，越久的历史，其权重就越小（衰减速率由 \miu 控制）。

R 通过一个非线性函数（文中说是 sigmoid），得到的结果叫做 receptance，我的理解有点类似于 RNN 里的 forget gate 。

WKV 这块（公式 14）是 Time Mix 的核心，它的作用就是前面提到的 RWKV 的”注意力“机制。 WKV 的计算有点类似于 Softmax，但是在分子和分母上分别加了一个当前位置 v_t 的项。 和注意力公式差不多， WKV_t 可以理解为：位置 t 相对于它之前各个位置的相关性（”注意力”），从公式看，WKV_t 应该是一个大小为 C 的向量（C是 channel 数）。

最后， receptance 和 WKV_t 相乘（element-wise product，两边都是大小为 C 的向量），得到位置 t 的输出 o_t （也是大小为 C 的向量）。显然，o_t 的计算包含了历史信息，随着 t 的增加，o_t 会依赖于越来越长的历史。

Channel Mixing 模块就简单多了，R和 K 同前面的 Time Mixing 一样，同样的 receptance，与 K 相乘，但在相乘之前，先要做一个 ReLU 的平方（squared ReLU）。这里 receptance 也是用作 forget gate，用于过滤掉不需要的历史信息。

相比以前的 RNN/LSTM，RWKV 最大的改进就是在训练时能够像 Transformer 那样并行计算。这方面是我最感兴趣的，但是很遗憾论文里面只是简单提了一句，没有仔细讲。这里根据论文里的公式，结合我自己的理解来分析一下。

先从简单的看起，R，K，V这三个向量的计算，都是把当前输入和前一个输入按一定权重混合。在训练的时候，我们处理的是整个序列，因此这种混合，可以通过构造两个序列，两者相差 1 个位置，然后并行混合即可。论文中也提到，这种 shift-by-one 的序列，可以用 PyTorch 的 nn.ZeroPad2d((0,0,1,-1)) 来构造。

有了各个位置的 R，K，Channel Mixing 的输出就可以并行计算，因为它没有时间上的依赖（公式 18）。

Time Mixing 部分，核心是 WKV 计算（公式 14)。

这个公式其实是一个类似 Softmax 的计算，只是分子和分母各加了一项，两者都是 weighted sum，和 Transformer 有点类似，因此我们可以用类似 Transformer 的方法来并行计算。论文中没有提到具体的做法，实际肯定也不止一种，我根据自己的理解提一个：

不过我看了 RWKV-LM 的几个版本，好像没有一个是完全按照论文来做的，当然，也可能是我没看懂他的 CUDA 代码 🤦

RNN 的一大优势就是在推理时只需要上一个状态，而不是全部的历史序列，因此生成每个 token 的计算和内存都是常量。

RWKV 算法可以写成 RNN 那样的递归形式，因此也具有 RNN 推理时的高效（这里只讲了 Time-Mixing 部分，Channel-Mixing 本身就是递归形式，就不讨论了）：

其中 \(a_t, b_t\) 就是上一个 WKV 状态的分子和分母，也就是模型的隐状态 h。整个过程可以描述成这样：

由于 RWKV 生成每个 token 的计算是常量，因此生成时的计算量是随着序列长度线性增长的，和 Transformer 相比，后者是随着序列长度的平方增长。同时 RWKV 在推理时的内存消耗也是固定的，不需要 KV cache。因此 RWKV 在推理效率和成本上有明显的优势。