【自然语言处理】【大模型】RWKV：基于RNN的LLM

相关博客
【自然语言处理】【大模型】RWKV：基于RNN的LLM
【自然语言处理】【大模型】CodeGen：一个用于多轮程序合成的代码大语言模型
【自然语言处理】【大模型】CodeGeeX：用于代码生成的多语言预训练模型
【自然语言处理】【大模型】LaMDA：用于对话应用程序的语言模型
【自然语言处理】【大模型】DeepMind的大模型Gopher
【自然语言处理】【大模型】Chinchilla：训练计算利用率最优的大语言模型
【自然语言处理】【大模型】大语言模型BLOOM推理工具测试
【自然语言处理】【大模型】GLM-130B：一个开源双语预训练语言模型
【自然语言处理】【大模型】用于大型Transformer的8-bit矩阵乘法介绍
【自然语言处理】【大模型】BLOOM：一个176B参数且可开放获取的多语言模型
【自然语言处理】【大模型】PaLM：基于Pathways的大语言模型
【自然语言处理】【chatGPT系列】大语言模型可以自我改进
【自然语言处理】【ChatGPT系列】FLAN：微调语言模型是Zero-Shot学习器
【自然语言处理】【ChatGPT系列】ChatGPT的智能来自哪里？

RWKV：基于RNN的LLM

​ 基于Transformer的LLM已经取得了巨大的成功，但是其在显存消耗和计算复杂度上都很高。RWKV是一个基于RNN的LLM，其能够像Transformer那样高效的并行训练，也能够像RNN那样高效的推理。

一、背景知识

1. RNN

​ RNN是指一类神经网络模型结构，其中最具有代表性的是LSTM：
f t = σ g ( W f x t + U f h t − 1 + b f ) i t = σ g ( W i x t + U i h t − 1 + b i ) o t = σ g ( W o x t + U o h t − 1 + b o ) c ~ t = σ c ( W c x t + U c h t − 1 + b c ) c t = f t ⊙ c t − 1 + i t ⊙ c ~ t h t = o t ⊙ σ h ( c t ) ft=σg(Wfxt+Ufht−1+bf)it=σg(Wixt+Uiht−1+bi)ot=σg(Woxt+Uoht−1+bo)˜ct=σc(Wcxt+Ucht−1+bc)ct=ft⊙ct−1+it⊙˜ctht=ot⊙σh(ct)

\\ ft​it​ot​c~t​ct​ht​​=σg​(Wf​xt​+Uf​ht−1​+bf​)=σg​(Wi​xt​+Ui​ht−1​+bi​)=σg​(Wo​xt​+Uo​ht−1​+bo​)=σc​(Wc​xt​+Uc​ht−1​+bc​)=ft​⊙ct−1​+it​⊙c~t​=ot​⊙σh​(ct​)​(1)(2)(3)(4)(5)(6)​
其中，$x_t$是当前时间步的输入，$h_{t-1}$是上一个时间步的隐藏状态，所有的$W$、$U$和$b$都是可学习参数，$\sigma$表示$\text{sigmoid}$函数。$f_t$是“遗忘门”，用来控制前一个时间步上传递信息的比例；$i_t$是“输入门”，用于控制当前时间步保留的信息比例；$o_t$是"输出门"，用于产生最终的输出。

2. Transformers和AFT

​ Transformer是NLP中主流的一种模型架构，其依赖于注意力机制来捕获所有输入和输出tokens的关系：
$$\begin{gathered} \text{Attn}(Q,K,V)=\text{softmax}(Q{K}^{\top })V \\ \text{\hspace{1em}(7)} \end{gathered}$$
为了简洁，这里忽略了多头和缩放因子$\frac{1}{\sqrt{d_k}}$。$Q{K}^{\top }$是序列中每个token之间的成对注意力分数，其能够被分解为向量表示：
$$\begin{gathered} \text{Attn}(Q,K,V{)}_t=\sum _{i=1}^T\frac{e^{q_t^{\top }{k}_i}}{{\sum }_{i=1}^T{e}^{q_t^{\top }{k}_i}}{v}_i=\frac{{\sum }_{i=1}^T{e}^{q_t^{\top }{k}_i}{v}_i}{{\sum }_{i=1}^T{e}^{q_t^{\top }{k}_i}} \\ \text{\hspace{1em}(8)} \end{gathered}$$
在AFT中，设计了一种注意力变体：
$$\begin{gathered} {\text{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}} \\ \text{\hspace{1em}(9)} \end{gathered}$$
其中， { w t , i } ∈ R T × T \{w_{t,i}\}\in R^{T\times T} {wt,i​}∈RT×T是可学习的位置偏差，每个$w_{t,i}$是一个标量。

​ 受AFT启发，在RWKV中的$w_{t,i}$是一个乘以相对位置的时间衰减向量：
$$\begin{gathered} w_{t,i}=-(t-i)w \\ \text{\hspace{1em}(10)} \end{gathered}$$
其中，$w\in (R_{\ge 0}{)}^d$，$d$是通道数。这里需要$w$是非负来保证$e^{w_{t,i}}\le 1$并且每个信道随时间衰减。

二、RWKV(Receptance Weighted Key Value)

​ RWKV由一系列的基本Block组成，每个Block则由time-mixing block和channel-mixing block组成的(如上图所示)。

​ RWKV递归的形式可以看做是当前输入和前一个时间不输入的线性插值，如上图所示。

1. Time-mixing block

​ Time-mixing block的作用同Self-Attention相同，就是提供全局token的交互。细节如下：
r t = W r ⋅ ( μ r x t + ( 1 − u r ) x t − 1 ) k t = W k ⋅ ( μ k x t + ( 1 − u k ) x t − 1 ) v t = W v ⋅ ( μ v x t + ( 1 − μ v ) x t − 1 ) w k v t = ∑ i = 1 t − 1 e − ( t − 1 − i ) w + k i v i + e u + k t v t ∑ i = 1 t − 1 e − ( t − 1 − i ) w + k i + e u + k t o t = W o ⋅ ( σ ( r t ) ⊙ w k v t ) rt=Wr⋅(μrxt+(1−ur)xt−1)kt=Wk⋅(μkxt+(1−uk)xt−1)vt=Wv⋅(μvxt+(1−μv)xt−1)wkvt=∑t−1i=1e−(t−1−i)w+kivi+eu+ktvt∑t−1i=1e−(t−1−i)w+ki+eu+ktot=Wo⋅(σ(rt)⊙wkvt)

\\ rt​kt​vt​wkvt​ot​​=Wr​⋅(μr​xt​+(1−ur​)xt−1​)=Wk​⋅(μk​xt​+(1−uk​)xt−1​)=Wv​⋅(μv​xt​+(1−μv​)xt−1​)=∑i=1t−1​e−(t−1−i)w+ki​+eu+kt​∑i=1t−1​e−(t−1−i)w+ki​vi​+eu+kt​vt​​=Wo​⋅(σ(rt​)⊙wkvt​)​(11)(12)(13)(14)(15)​
所有的$\mu$和$W$都是可训练参数，$r_t$、$k_t$和$v_t$是当前输入$x_t$和上一个时间步输入$x_{t-1}$的加权投影。

公式(14)中，$w$和$u$是可训练参数，分子的第一项${\sum }_{i=1}^{t-1}{e}^{-(t-1-i)w+k_i}{v}_i$表示前$t-1$步的加权结果，$-(t-1-i)w+k_i$是随相对距离逐步衰减；$e^{u+k_t}{v}_t$则是当前时间步的结果。

公式(15)中，则通过$\sigma (r_t)$控制最终输出的比例。

2. Channel-mixing block

​ Channel-mixing block类似于Transformer中的FFN部分，细节如下：
r t = W r ⋅ ( μ r x t − ( 1 − μ r ) x t − 1 ) k t = W k ⋅ ( μ k x t − ( 1 − μ k ) x t − 1 ) o t = σ ( r t ) ⊙ ( W v ⋅ max ⁡ ( k t , 0 ) 2 ) rt=Wr⋅(μrxt−(1−μr)xt−1)kt=Wk⋅(μkxt−(1−μk)xt−1)ot=σ(rt)⊙(Wv⋅max(kt,0)2)

\\ rt​kt​ot​​=Wr​⋅(μr​xt​−(1−μr​)xt−1​)=Wk​⋅(μk​xt​−(1−μk​)xt−1​)=σ(rt​)⊙(Wv​⋅max(kt​,0)2)​(16)(17)(18)​

三、并行训练和序列解码

​ RWKV可以类似Transformer那样高效的并行。设batch size为B、seq_length为T、channels为d，计算量主要来自于矩阵乘法 W □ , □ ∈ { r , k , v , o } W_\square,\square\in \{r,k,v,o\} W□​,□∈{r,k,v,o}，单层的时间复杂度为$O(BT{d}^2)$。此外，更新注意力分数$wk{v}_t$需要顺序扫描，其时间复杂度为$O(BTd)$。矩阵乘法可以像Transformer那样并行，但是WKV的计算是依赖时间步的，所以只能在其他维度上并行。

​ RWKV具有类似RNN的结构，解码时将$t$步的输出作为$t+1$步的输入。相比于自注意力机制随着序列长度，计算复杂度呈平方次增长，RWKV则是与序列长度呈线性关系。因此，RWKV能够更高效的处理更长的序列。