相关博客
【自然语言处理】【大模型】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的智能来自哪里？
基于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 ) \begin{align} f_t&=\sigma_g(W_fx_t+U_f h_{t-1}+b_f) \tag*{(1)} \\ i_t&=\sigma_g(W_ix_t+U_i h_{t-1}+b_i) \tag*{(2)} \\ o_t&=\sigma_g(W_ox_t+U_o h_{t-1}+b_o) \tag*{(3)} \\ \tilde{c}_t&=\sigma_c(W_cx_t+U_c h_{t-1}+b_c) \tag*{(4)} \\ c_t&=f_t\odot c_{t-1}+i_t\odot\tilde{c}_t \tag*{(5)} \\ h_t&=o_t\odot\sigma_h(c_t) \tag*{(6)} \end{align} \\ ftitotc~tctht=σg(Wfxt+Ufht−1+bf)=σg(Wixt+Uiht−1+bi)=σg(Woxt+Uoht−1+bo)=σc(Wcxt+Ucht−1+bc)=ft⊙ct−1+it⊙c~t=ot⊙σh(ct)(1)(2)(3)(4)(5)(6)
其中， x t x_t xt是当前时间步的输入， h t − 1 h_{t-1} ht−1是上一个时间步的隐藏状态，所有的 W W W、 U U U和 b b b都是可学习参数， σ \sigma σ表示 sigmoid \text{sigmoid} sigmoid函数。 f t f_t ft是“遗忘门”，用来控制前一个时间步上传递信息的比例； i t i_t it是“输入门”，用于控制当前时间步保留的信息比例； o t o_t ot是"输出门"，用于产生最终的输出。
2. Transformers和AFT
Transformer是NLP中主流的一种模型架构，其依赖于注意力机制来捕获所有输入和输出tokens的关系：
Attn ( Q , K , V ) = softmax ( Q K ⊤ ) V (7) \text{Attn}(Q,K,V)=\text{softmax}(QK^\top)V \tag{7} \\ Attn(Q,K,V)=softmax(QK⊤)V(7)
为了简洁，这里忽略了多头和缩放因子 1 d k \frac{1}{\sqrt{d_k}} dk1。 Q K ⊤ QK^\top QK⊤是序列中每个token之间的成对注意力分数，其能够被分解为向量表示：
Attn ( Q , K , V ) t = ∑ i = 1 T e q t ⊤ k i ∑ i = 1 T e q t ⊤ k i v i = ∑ i = 1 T e q t ⊤ k i v i ∑ i = 1 T e q t ⊤ k i (8) \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}}\tag{8} \\ Attn(Q,K,V)t=i=1∑T∑i=1Teqt⊤kieqt⊤kivi=∑i=1Teqt⊤ki∑i=1Teqt⊤kivi(8)
在AFT中，设计了一种注意力变体：
Attn + ( W , K , V ) t = ∑ i = 1 t e w t , i + k i v i ∑ i = 1 t e w t , i + k i (9) \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}} \tag{9} \\ Attn+(W,K,V)t=∑i=1tewt,i+ki∑i=1tewt,i+kivi(9)
其中， { w t , i } ∈ R T × T \{w_{t,i}\}\in R^{T\times T} {wt,i}∈RT×T是可学习的位置偏差，每个 w t , i w_{t,i} wt,i是一个标量。
受AFT启发，在RWKV中的 w t , i w_{t,i} wt,i是一个乘以相对位置的时间衰减向量：
w t , i = − ( t − i ) w (10) w_{t,i}=-(t-i)w \tag{10} \\ wt,i=−(t−i)w(10)
其中， w ∈ ( R ≥ 0 ) d w\in (R_{\geq 0})^d w∈(R≥0)d， d d d是通道数。这里需要 w w w是非负来保证 e w t , i ≤ 1 e^{w_{t,i}}\leq 1 ewt,i≤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 ) \begin{align} r_t&=W_r\cdot(\mu_rx_t+(1-u_r)x_{t-1}) \tag*{(11)} \\ k_t&=W_k\cdot(\mu_kx_t+(1-u_k)x_{t-1}) \tag*{(12)} \\ v_t&=W_v\cdot(\mu_vx_t+(1-\mu_v)x_{t-1}) \tag*{(13)} \\ wkv_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}} \tag*{(14)} \\ o_t&=W_o\cdot(\sigma(r_t)\odot wkv_t) \tag*{(15)} \end{align} \\ rtktvtwkvtot=Wr⋅(μrxt+(1−ur)xt−1)=Wk⋅(μkxt+(1−uk)xt−1)=Wv⋅(μvxt+(1−μv)xt−1)=∑i=1t−1e−(t−1−i)w+ki+eu+kt∑i=1t−1e−(t−1−i)w+kivi+eu+ktvt=Wo⋅(σ(rt)⊙wkvt)(11)(12)(13)(14)(15)
所有的 μ \mu μ和 W W W都是可训练参数， r t r_t rt、 k t k_t kt和 v t v_t vt是当前输入 x t x_t xt和上一个时间步输入 x t − 1 x_{t-1} xt−1的加权投影。
公式(14)中， w w w和 u u u是可训练参数，分子的第一项 ∑ i = 1 t − 1 e − ( t − 1 − i ) w + k i v i \sum_{i=1}^{t-1}e^{-(t-1-i)w+k_i}v_i ∑i=1t−1e−(t−1−i)w+kivi表示前 t − 1 t-1 t−1步的加权结果， − ( t − 1 − i ) w + k i -(t-1-i)w+k_i −(t−1−i)w+ki是随相对距离逐步衰减； e u + k t v t e^{u+k_t}v_t eu+ktvt则是当前时间步的结果。
公式(15)中，则通过 σ ( r t ) \sigma(r_t) σ(rt)控制最终输出的比例。
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 ) \begin{align} r_t&=W_r\cdot(\mu_rx_t-(1-\mu_r)x_{t-1}) \tag*{(16)} \\ k_t&=W_k\cdot(\mu_kx_t-(1-\mu_k)x_{t-1}) \tag*{(17)} \\ o_t&=\sigma(r_t)\odot(W_v\cdot\max(k_t,0)^2) \tag*{(18)} \\ \end{align} \\ rtktot=Wr⋅(μrxt−(1−μr)xt−1)=Wk⋅(μkxt−(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 ( B T d 2 ) O(BTd^2) O(BTd2)。此外，更新注意力分数 w k v t wkv_t wkvt需要顺序扫描，其时间复杂度为 O ( B T d ) O(BTd) O(BTd)。矩阵乘法可以像Transformer那样并行，但是WKV的计算是依赖时间步的，所以只能在其他维度上并行。
RWKV具有类似RNN的结构，解码时将 t t t步的输出作为 t + 1 t+1 t+1步的输入。相比于自注意力机制随着序列长度，计算复杂度呈平方次增长，RWKV则是与序列长度呈线性关系。因此，RWKV能够更高效的处理更长的序列。
