RWKV:一种鱼和熊掌兼得的线性transformer模型

曾经终败给现在 2023-10-12 15:47 23阅读 0赞

众所周知，现在transformer及其变种是NLP和CV领域已经杀疯了。但其中最核心的self-attention机制因为其O(N2)的时间复杂度（二次依赖问题）被诟病。

在不改变transformer block这个整体架构的前提下，现在学术界解决二次依赖问题的主要是两个思路。一种是实现self-attention的线性化。这方面的工作是很多的，比如Performer\[5\]、Reformer\[6\]、Linformer\[7\]、Nyströmformer\[9\]、AdaMRA\[10\]等。关于这部分工作更多的内容大家可以在苏剑林的博客中了解到\[8\].虽然关于线性attention的工作很多，但参考AdaMRA\[10\]论文的图。只有Nyströmformer\[9\]和AdaMRA\[10\]相较于Transformer能获得速度和效果的双重提升，其他的大多需要付出效果的代价才能获取一定的速度提升。但就是这哥俩由于用了平均池化作为特征聚类，因此无法mask未来信息从而丧失了自回归的能力。因此通过替换线性attention从而提升transformer速度这一思路是必须付出代价的。

![389581d4382ca7d3bac9e67d37e0b55c.png][]

另一种思路将self-attention换成其他线性复杂度的部件。比如前段时间谷歌发现用膨胀卷积取代self-attention也能取到不错的效果\[1\]。而在CV领域杀疯的MLP-Mixer\[2\]，兼具CV和NLP能力的gMLP、aMLP,\[3\]MLP-Mixer的NLP版本Synthesizer\[4\]。但都有或多或少的缺点，就比如Synthesizer和gMLP在NLP领域相较于self-attention还是差了点的。而aMLP虽然效果好了吧，但其实还是要用到self-attention，提速的目的还是没达到。不过今年暑假那会，苹果提出的AFT模型\[11\]号称自己是最快的transformer模型。

![207478c9b1cbfe4720453d4923eca057.png][]

上述是标准AFT的公式，其中σ是sigmoid函数，QKV就是sefl-attention的那一套，w是一个训练出来的参数矩阵。不难看出AFT是通过点乘的方式实现的注意力，在做自回归时只需要对W矩阵进行mask即可。并且W矩阵是自带位置信息的，不仅解决了部分线性attention不能做自回归的问题，还顺便把transformer里位置编码的问题给解决了。可以说AFT实现了一举三得。但成也萧何败也萧何，W矩阵是AFT成功的核心也是AFT的最大缺点。一般来说W应该是一个\[max\_len,max\_len\]大小的方阵。换而言之AFT所能处理的文本长度受限于W矩阵的大小，如果想要处理一万字的长文本，W矩阵的参数量就快赶上Bert了。为了解决这个问题，下面该本文的主角RWKV出场了。RWKV的原文在[RWKV is all you need？一种新语言模型，改进 Transformer - 知乎][RWKV is all you need_ Transformer -]，不过原文实在过于简短了不便阅读和理解。因此笔者写了此文介绍一下RWKV是怎么实现鱼和熊掌兼得的。

### RWKV ###

**整体结构** RWKV的整体结构依然采用的是transformer block的思路,其整体结构如图所示。相较于原始transformer block的结构，RWKV将self-attention替换为Position Encoding和TimeMix，将FFN替换为ChannelMix。其余部分与transfomer一致的。

![4e025381e9095daf16e3db08ae92d4d3.png][]

**Position Matrix** RWKV采用的位置编码类似于AliBi编码\[12\]的形式。原文作者并没有给他的位置编码命名，为了便于介绍参考该位置编码主要考虑距离衰减的特性，本文将其命名为distance编码。对于第i个head的第j个token而言,其位置编码如下述公式所示。其中n*head表示头的数量，max*\_len表示为所允许的最大长度。

![626d77aa7011009cc7104f9e644a5c3e.png][]

目前学术界的主流观点是RNN结构是天然的时序结构，不需要transformer模型必须的位置编码。而如果我们查看RNN的计算流程，可以发现RNN只考虑到当前token及之前的信息，而随着距离的延长前面的信息会逐渐减少。而distance位置编码便是参考RNN时序特点所设计的。

不过RWKV模型中，不会直接对输入的X进行上述计算。而是得到类似AFT中的W矩阵参与后续Time-Mix计算。其中W矩阵的形状为\[n\_head,seq\_len,seq\_len\]。因此对于W矩阵中的而言，其数值如下述公式所示。

![d25dd40ff68d40f6cebf4a6275a58632.png][]

从这里不难看出，AFT中的W矩阵在RWKV中是通过公式得到而不是训练得到的，因此解决了AFT中无法解决长文本，或解决长文本时参数爆炸的问题。

当然，在处理的任务文本长度有限的情况下。比如机器翻译，或者是RWKV目前应用的ai写小说这类应用场景。在这类应用场景中，由于不会面临长文本的情况，因此可以为W矩阵添加更多的位置信息。参考公式如下

![8fd06942db3772d3fa4b62f94967214b.png][]

其中和分别为形状\[n\_head,seq\_len,1\]和\[n\_head,1,seq\_len\]的向量，在初始化时为全1矩阵。即将作为W矩阵的初始化。结合该步后，在形式上W矩阵融合了distance编码中的距离信息与相对信息。

值得注意的是，原作者是设计distance编码时专门设计了一个不考虑位置信息衰减的头。即该头的W矩阵是一个全一的下三角矩阵。

**Time-shit** 在介绍TimeMix之前，要先介绍一下RWKV所使用的Time-shit技巧。

原文：[Time-shift: 一行代码，免费提高 Transformer 性能（无参数，无耗时） - 知乎][Time-shift_ _ Transformer _ -]

Time-shiit是原作者提出的一种几乎零成本提升模型效果的trick，实现代码如下所示。

Torch实现
    C=x.shape[-1]
    self.time_shift = nn.ZeroPad2d((0,0,1,0))
    x = torch.cat([self.time_shift(x)[:, :-1, :C//2], x[:, :, C//2:]], dim = -1)
    Keras实现
    d=K.shape(x)[-1]
    x=K.concatenate(K.temporal_padding(x,(1,0))[:,:-1,:d//2],x[:,:,d//2:])

可以看出不论哪个框架也就两行就能实现了，为了便于读者理解。假设存在一个3x4的矩阵。

![b6d903ac15f206e64d762bbc3815c2e5.png][]

在经过time-shift后变为

![b598a87eb43dfcc44b0385f3a6afb225.png][]

其实就相当于插入一个小的RNN，实验表明简单的trick能让模型的更快更好地收敛。

**TimeMix** TimeMix是RWKV中用于代替self-attention的部分，基于AFT的基础上做出改进兼具了线性的速度和较好的性能。在进行该步前，需要对输入的x进行time-shift。

同self-attention中的QKV矩阵一样，RWKV中也有对应的RKV矩阵。对与输出矩阵中第i个头的第j个token而言计算步骤如下所示。

![876bb22f5be751d871e84e20a0f9445c.png][]

这其中是一个\[hiden\_size,hiden\_size\]大小的方阵，与常规attention一样用于最后的输出。而是一个\[seq\_len,hiden\_size\]大小的矩阵，其作用笔者猜测应该是类似于bias的作用。

**ChannelMix** ChannelMix 是RWKV中用于替代FFN的部分。类似于tiny attention之于attention。ChannelMix本质上来说是一个tiny TimeMix。

在进行该步计算前，和TimeMix一样要先进行一次time-shift。随后依然要计算出RKV矩阵和W权重。不过有所不同的是在这一步中假设输入x的维度是embed\_size，则R的维度应和X相同。KV的维度是用户所自定义的hidden\_size,W的形状为\[hidden\_size,embed\_size\].

通过设置较小的hidden\_size可以实现一个tiny版TimeMix，能在对性能影响较小的情况下实现提速。当hidden\_size==embed\_size时，可以看作一个不考虑位置信息和归一化的TimeMix或者看作点乘式的FFN。

具体计算公式如下所示

![76757ddd9ce41f4b72475936a57f58f1.png][]

**总结** 本文介绍了一种鱼和熊掌兼得的模型。既能和AFT一样兼具通用性和高效，distance位置编码的设计使得模型也具备面对超长文本的能力。

实际实验效果可以去看原文的内容，本文只对其结构进行介绍。但总体而言，笔者测试过基于GPT的ai写小说和基于RWKV的ai写小说。相比较而言，RWKV的写出来的文章会更流畅，并且在训练时收敛速度页更快。

参考文献

\[1\] Are Pre-trained Convolutions Better than Pre-trained Transformers [https://arxiv.org/pdf/2105.03322.pdf][https_arxiv.org_pdf_2105.03322.pdf]

\[2\] MLP-Mixer: An all-MLP Architecture for Vision [https://arxiv.org/pdf/2105.01601.pdf][https_arxiv.org_pdf_2105.01601.pdf]

\[3\] Pay Attention to MLPs [https://arxiv.org/pdf/2105.08050.pdf][https_arxiv.org_pdf_2105.08050.pdf]

\[4\] Synthesizer: Rethinking Self-Attention in Transformer Models [https://arxiv.org/abs/2005.00743][https_arxiv.org_abs_2005.00743]

\[5\] Rethinking Attention with Performers [https://arxiv.org/abs/2009.14794][https_arxiv.org_abs_2009.14794]

\[6\] Reformer: The Efficient Transformer [https://arxiv.org/abs/2001.04451][https_arxiv.org_abs_2001.04451]

\[7\] Linformer: Self-Attention with Linear Complexity [https://arxiv.org/abs/2006.04768][https_arxiv.org_abs_2006.04768]

\[8\] 线性Attention的探索：Attention必须有个Softmax吗？ [https://spaces.ac.cn/archives/7546][https_spaces.ac.cn_archives_7546]

\[9\] Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention [https://arxiv.org/abs/2102.03902][https_arxiv.org_abs_2102.03902]

\[10\] Adaptive Multi-Resolution Attention with Linear Complexity [https://arxiv.org/abs/2108.04962][https_arxiv.org_abs_2108.04962]

\[11\] An Attention Free Transformer [https://arxiv.org/abs/2105.14103][https_arxiv.org_abs_2105.14103]

\[12\] Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation

[https://arxiv.org/abs/2108.12409][https_arxiv.org_abs_2108.12409]

--------------------

[RWKV:一种鱼和熊掌兼得的线性transformer模型 - 知乎][RWKV_transformer_ -]

[https://www.youtube.com/watch?v=oaP8\_fUFVWw][https_www.youtube.com_watch_v_oaP8_fUFVWw]

[389581d4382ca7d3bac9e67d37e0b55c.png]: https://img-blog.csdnimg.cn/img_convert/389581d4382ca7d3bac9e67d37e0b55c.png
[207478c9b1cbfe4720453d4923eca057.png]: https://img-blog.csdnimg.cn/img_convert/207478c9b1cbfe4720453d4923eca057.png
[RWKV is all you need_ Transformer -]: https://zhuanlan.zhihu.com/p/397985790
[4e025381e9095daf16e3db08ae92d4d3.png]: https://img-blog.csdnimg.cn/img_convert/4e025381e9095daf16e3db08ae92d4d3.png
[626d77aa7011009cc7104f9e644a5c3e.png]: https://img-blog.csdnimg.cn/img_convert/626d77aa7011009cc7104f9e644a5c3e.png
[d25dd40ff68d40f6cebf4a6275a58632.png]: https://img-blog.csdnimg.cn/img_convert/d25dd40ff68d40f6cebf4a6275a58632.png
[8fd06942db3772d3fa4b62f94967214b.png]: https://img-blog.csdnimg.cn/img_convert/8fd06942db3772d3fa4b62f94967214b.png
[Time-shift_ _ Transformer _ -]: https://zhuanlan.zhihu.com/p/399480671
[b6d903ac15f206e64d762bbc3815c2e5.png]: https://img-blog.csdnimg.cn/img_convert/b6d903ac15f206e64d762bbc3815c2e5.png
[b598a87eb43dfcc44b0385f3a6afb225.png]: https://img-blog.csdnimg.cn/img_convert/b598a87eb43dfcc44b0385f3a6afb225.png
[876bb22f5be751d871e84e20a0f9445c.png]: https://img-blog.csdnimg.cn/img_convert/876bb22f5be751d871e84e20a0f9445c.png
[76757ddd9ce41f4b72475936a57f58f1.png]: https://img-blog.csdnimg.cn/img_convert/76757ddd9ce41f4b72475936a57f58f1.png
[https_arxiv.org_pdf_2105.03322.pdf]: https://link.zhihu.com/?target=https%3A//arxiv.org/pdf/2105.03322.pdf
[https_arxiv.org_pdf_2105.01601.pdf]: https://link.zhihu.com/?target=https%3A//arxiv.org/pdf/2105.01601.pdf
[https_arxiv.org_pdf_2105.08050.pdf]: https://link.zhihu.com/?target=https%3A//arxiv.org/pdf/2105.08050.pdf
[https_arxiv.org_abs_2005.00743]: https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2005.00743
[https_arxiv.org_abs_2009.14794]: https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2009.14794
[https_arxiv.org_abs_2001.04451]: https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2001.04451
[https_arxiv.org_abs_2006.04768]: https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2006.04768
[https_spaces.ac.cn_archives_7546]: https://link.zhihu.com/?target=https%3A//spaces.ac.cn/archives/7546
[https_arxiv.org_abs_2102.03902]: https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2102.03902
[https_arxiv.org_abs_2108.04962]: https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2108.04962
[https_arxiv.org_abs_2105.14103]: https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2105.14103
[https_arxiv.org_abs_2108.12409]: https://link.zhihu.com/?target=https%3A//arxiv.org/abs/2108.12409
[RWKV_transformer_ -]: https://zhuanlan.zhihu.com/p/437714049
[https_www.youtube.com_watch_v_oaP8_fUFVWw]: https://www.youtube.com/watch?v=oaP8_fUFVWw