RWKV论文燃爆！将RNN崛起进行到底！可扩百亿级参数，与Transformer表现相当！

zenRRan

发布于 2023-08-22 14:08:49

5910

发布于 2023-08-22 14:08:49

文章被收录于专栏：深度学习自然语言处理

深度学习自然语言处理原创 作者：鸽鸽

万众期待的RWKV论文来啦！

这股RNN崛起的“清流”，由民间开源组织发起，号称是第一个可扩展到百亿级参数的非transformer架构！

RWKV结合了RNN和Transformer的优势：一方面，抛弃传统的点积自注意力、使用线性注意力，解决transformer内存和计算复杂度随序列增长呈平方缩放的瓶颈；另一方面，突破了RNN梯度消失、并行化和可扩展性等限制。

居然实现O(Td)的时间复杂度和O(d)的空间复杂度！

表：不同transformer的复杂度比较。T是序列长度，d是特征维度

今天我们基于这篇论文，讲讲RWKV背后的注意力、时间混合和通道混合模块的原理与组成。

论文：RWKV: Reinventing RNNs for the Transformer Era 地址：https://arxiv.org/pdf/2305.13048.pdf 代码: https://github.com/BlinkDL/RWKV-LM 模型：https://huggingface.co/BlinkDL/rwkv-4-raven

回顾RNN

流行的RNN架构（如LSTM）公式如下：

\begin{aligned} f_{t} & =\sigma_{g}\left(W_{f} x_{t}+U_{f} h_{t-1}+b_{f}\right), \\ i_{t} & =\sigma_{g}\left(W_{i} x_{t}+U_{i} h_{t-1}+b_{i}\right), \\ o_{t} & =\sigma_{g}\left(W_{o} x_{t}+U_{o} h_{t-1}+b_{o}\right), \\ \tilde{c}_{t} & =\sigma_{c}\left(W_{c} x_{t}+U_{c} h_{t-1}+b_{c}\right), \\ c_{t} & =f_{t} \odot c_{t-1}+i_{t} \odot \tilde{c}_{t}, \\ h_{t} & =o_{t} \odot \sigma_{h}\left(c_{t}\right) . \end{aligned}

下图(a)展示了RNN的数据流，尽管RNN可以分解为两个线性块（

和

）和一个特定于RNN的块，但对于先前时间步的数据依赖阻止了RNN的并行化。

图:RWKV与QRNN和RNN（Vanilla、LSTM、GRU等）的架构对比。橙色表示时间混合、卷积或矩阵乘法，连续块表示这些计算可以同时进行；蓝色表示沿着通道或特征维度（按元素）同时操作的无参函数；绿色表示通道混合

RWKV 注意力进化史

要知道 RWKV 的线性注意力怎么来的，我们先看经典的transformer自注意力（省略multi-head和缩放因子

\frac{1}{\sqrt{d_{k}}}

）：

\operatorname{Attn}(Q, K, V)=\operatorname{softmax}\left(Q K^{\top}\right) V

其中核心的

Q K^{\top}

乘法是序列中每两两词元 (token) 之间成对的注意力分数的集合，对于每个时间步 t，分解为向量操作：

\operatorname{Attn}(Q, K, V)_{t}=\frac{\sum_{i=1}^{T} e^{q_{t}^{\top} k_{i}} v_{i}}{\sum_{i=1}^{T} e^{q_{t}^{\top} k_{i}}}

在Apple提出的 Attention Free Transformer (AFT) 中，引入成对位置偏差来替换点积，可以看作是每个特征维度对应一个头的多头注意力。上面的等式被写作：

\operatorname{Attn}^{+}(W, K, V)_{t}=\frac{\sum_{i=1}^{t} e^{w_{t, i}+k_{i}} v_{i}}{\sum_{i=1}^{t} e^{w_{t, i}+k_{i}}}

其中每个

w_{t, i}

是一个标量，组合成代表成对位置偏差的可学习参数矩阵

\left\{w_{t, i}\right\} \in R^{T \times T}

在 RWKV 中，参数

w_{t, i}

进一步简化为一个通道级时间衰减向量，乘以相对位置：

w_{t, i}=-(t-i) w

其中

w \in\left(R_{\geq 0}\right)^{d}

，

是通道数（特征维度）。

RWKV 模型架构

The Receptance Weighted Key Value (RWKV) 的名字来自于时间混合 (time-mixing) 和通道混合 (channel-mixing) 块中使用的四个主要元素：

(Receptance) ：接受过去信息的接受向量；

(Weight)：位置权重衰减向量（可训练的模型参数）；

(Key) ：键是类似于传统注意力中

的向量；

(Value)：值是类似于传统注意力中

的向量。

每个时间步，主要元素之间通过乘法进行交互。RWKV 架构如图所示：

图：RWKV块元素（左）和带有用于语言建模的最终头部的RWKV残差块（右）

RWKV 架构由一系列堆叠的残差块组成，每个残差块由具有循环结构的时间混合和通道混合子块组成。该循环通过将当前输入和上一时间步的输入之间进行线性插值来实现（作者称为token shift），如下图中对角线所示。这个线性插值可以针对每个输入 embedding 的线性投影（例如时间混合中的

R,K,V

和通道混合中的

R,K

）独立调整。

图：用于语言建模的RWKV架构

时间混合块的公式如下：

\begin{aligned} r_{t} & =W_{r} \cdot\left(\mu_{r} x_{t}+\left(1-\mu_{r}\right) x_{t-1}\right) \\ k_{t} & =W_{k} \cdot\left(\mu_{k} x_{t}+\left(1-\mu_{k}\right) x_{t-1}\right) \\ v_{t} & =W_{v} \cdot\left(\mu_{v} x_{t}+\left(1-\mu_{v}\right) x_{t-1}\right) \\ w k v_{t} & =\frac{\sum_{i=1}^{t-1} e^{-(t-1-i) w+k_{i}} v_{i}+e^{u+k_{t}} v_{t}}{\sum_{i=1}^{t-1} e^{-(t-1-i) w+k_{i}}+e^{u+k_{t}}} \\ o_{t} & =W_{o} \cdot\left(\sigma\left(r_{t}\right) \odot w k v_{t}\right) \end{aligned}

其中，

W K V

的计算

w k v_{t}

扮演了Transformer中

\operatorname{Attn}(Q, K, V)

的角色，其中的时间衰减机制（方程中的

e^{-w}

）保持了对序列元素间位置关系的敏感性，逐渐减少过去信息对当前时间的影响。直观地说，随着时间

的增加，向量

o_{t}

依赖于一个长时间历史，由越来越多的项的总和所表示。对于目标位置

，RWKV在位置间隔

[1,t]

执行加权求和，然后乘以接受度

\sigma(r)

。因此，交互在给定时间步内是乘法，并在不同的时间步上求和。

通道混合块采用平方ReLU激活，公式如下：

\begin{aligned} r_{t} & =W_{r} \cdot\left(\mu_{r} x_{t}+\left(1-\mu_{r}\right) x_{t-1}\right), \\ k_{t} & =W_{k} \cdot\left(\mu_{k} x_{t}+\left(1-\mu_{k}\right) x_{t-1}\right), \\ o_{t} & =\sigma\left(r_{t}\right) \odot\left(W_{v} \cdot \max \left(k_{t}, 0\right)^{2}\right), \end{aligned}

注意在时间混合和通道混合中，Receptance取sigmoid的情况下，可以直观地将其视为“遗忘门”，以消除不必要的历史信息。

在最后一个块之后，使用由LayerNorm和线性投影组成的简单输出投影头来获取在下一个token预测任务中使用的logits和计算训练期间的交叉熵损失。训练采用时间并行模式，而自回归推理和聊天则采用时间顺序模式。

评估

实验表明，与具有相同参数和训练token数量的传统transformer架构（Pythia、OPT、BLOOM、GPT-Neo）相比，RWKV在六个基准测试（Winogrande、PIQA、ARC-C、ARC-E、LAMBADA和SciQ）上均具有竞争力。RWKV甚至在四项任务中超越了Pythia和GPT-Neo.

图：零样本表现：横轴是参数数量，竖轴是准确率

并且，增加上下文长度会导致Pile上的测试loss降低，这表明RWKV可以有效地利用长期的上下文信息。

图：增加上下文长度有助于在Pile上降低测试损失

有趣的是，在RWKV-4和ChatGPT / GPT-4的比较研究显示，RWKV-4对提示工程非常敏感。当将指令风格从适合GPT调整为更适合RWKV时，RTE的F1性能甚至从44.2％增加到74.8％。作者猜想是因为RNN不能回溯处理 ( retrospective processing) 来重新调整先前信息的权重。因此为了让性能更好，期望信息应该在问题之后展示。