RWKV:一种鱼和熊掌兼得的线性transformer模型
众所周知，在NLP和计算机视觉领域中，transformer及其变种已经取得了显著进展。然而，其中最核心的自注意力机制因其计算复杂度为二次方（即O(N²)）的问题而受到了批评。
在不更改transformer block这一整体架构的前提下
另一种思路将自注意力机制替换为其他低复杂度组件。例如研究表明，在计算机视觉领域中使用膨胀卷积替代自注意力机制也取得了不错的效果[1]。值得注意的是具有里程碑意义的是MLP-Mixer[2]它不仅展现了在计算机视觉领域的卓越性能而且兼顾了语言处理能力的gMLP与aMLP等创新架构[3]。然而这些方法均存在一定程度的不足例如当前基于语言处理领域的改进版本Synthesizer与gMLP相较于传统自注意力机制仍有差距尽管aMLP虽然在某些方面表现有所提升但其本质上仍依赖于自注意力机制从而未能完全实现加速目标。值得注意的是今年暑假期间苹果提出了AFT模型[11]声称其是目前为止最高效的transformer架构。
以σ函数为激活函数的标准AFT框架中,QKV机制对应于自注意力机制,w参数则是一个经过训练得到的有效参数矩阵.从关注机理来看,AFT模型确实采用了点积运算实现关注,在执行自回归任务时只需对权重矩阵W施加掩码约束即可.值得注意的是,W矩阵本身带有位置编码能力,这一特点不仅解决了部分线性注意力无法有效执行自回归的问题,同时也巧妙地消除了传统transformer架构中的位置编码需求.因此可以说,AFT模型实现了"一石三鸟"的最佳关注效果.然而,成也萧何败也萧何,AFT模型的核心优势同时也是其显著缺陷所在,W矩阵通常情况下应具有[max_len,max_len]维度.换言之,AFT架构所能处理的有效文本长度受到W参数量级限制,当处理超过几千个字符的内容时,其计算规模已经接近于BERT级别的要求.针对这一局限性,本文作者提出了一种创新性的解决方案...
RWKV
该模型采用了transformer block架构（见图）。相较于传统transformer架构（...），RWKV模型替换了位置编码和时间混合机制（Position Encoding & TimeMix），同时替换了通道混合机制（ChannelMix）。其他组件则保持与传统transformer一致。
Position Matrix RWKV采用了位置编码机制，其结构与现有的AliBi编码方案具有相似性[12]。为了便于讨论和引用,本文将该位置编码命名为distance编码。本文主要基于距离衰减特性对该位置编码进行了研究和探讨。对于第i个头中的第j个token而言,其位置编码由以下公式确定:其中n_head代表总头数,max_len则表示允许的最大序列长度。
在学术界主流的观点认为RNN架构具有内建的时间序列特性，并无需依赖transformer模型中的位置编码机制。从计算流程来看，在RNN架构中信息仅传递至当前token及其之前的输入单元。当时间间隔增大时，较早输入的信息影响力逐步减弱。基于此，在设计位置编码机制时也充分考虑了这一特点。
在RWKV模型中，不会直接作用于输入序列X执行上述运算步骤；相反地，在AFT架构的基础上引入了类似于时间门限混合机制（Time-Mix）的过程，并通过权重矩阵W实现这一功能。具体而言，在权重矩阵W中包含了三维结构[heads×序列长度×序列长度]的信息参数；这些参数的具体数值则由以下所述公式确定。
从这一角度来看，在RWKV架构中,AFT模型中的权重矩阵并非通过训练获得而是基于预定义公式推导得出,这使得AFT模型能够有效缓解长文本难以处理的问题以及因参数数量激增而导致的计算效率低下
在处理过程中，在任务文本长度受限的情况下。例如机器翻译任务或基于RWKV模型的应用场景。在这些场景中，并不需要面对长文本的问题。因此可以通过调整W矩阵来增加位置信息的数量。参考公式如下
在初始化阶段，我们采用两个具有不同维度形状[n_head,seq_len,1]和[n_head,1,seq_len]的向量来构建初始权重矩阵。这些向量将被用于W矩阵的初始化过程。在此基础上，通过这一步骤处理，在形式上W矩阵融合了distance编码中的距离信息与相对位置信息。
值得注意的是，在distance编码的设计过程中专门开发了一个不考虑位置信息衰减影响的头部模型。即该头部模型对应的W矩阵是一个全1的下三角矩阵。
时间延迟（time-delay）应被用于在引入混合模型（introduce）时进行调整。
Time-shift: 利用时间移位技术, 只需一行代码即可优化 Transformer 模型, 无需任何参数或计算开销 - 知乎
Time shift是一种近乎无成本提升模型效果的技巧，由提出者提出，并且其代码实现相对简便。
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的矩阵在这种情况下相关的算法运算将更加直观易懂
在经过time-shift后变为
从本质上来看,这相当于嵌入一个较简单的RNN网络;实验结果表明这种简洁的技术能够显著提升模型收敛速度.
TimeMix TimeMix作为RWKV中的一个替代机制，在自注意力机制的基础上进行了优化与改进。其主要优势在于结合了AFT框架的同时实现了线性计算特性，并且展现出较高的性能水平。为了确保后续计算的连续性，在执行该操作之前应先完成时间轴上的位移处理。
同self-attention机制中的Q、K、V矩阵相仿，在RWKV模型中也存在对应的R、K、V矩阵结构。对于输出矩阵中第i个子层的第j个输入单元而言，请具体计算流程如下所述。
其中构成一个[hiden_size,hiden_size]大小的方阵，并与常规attention机制类似地用于最后的输出。而是一个[seq_len,hiden_size]大小的矩阵，并推测其作用类似于bias的作用。
在RWKV架构中，ChannelMix 作为替代FFN层的关键组件发挥着重要作用。它与 tiny attention 在 attention 模块中的作用相仿，并从本质上讲，它借鉴了 tiny TimeMix 的设计理念与机制。
在该步骤开始之前，请先执行一次时间偏移操作（类似于TimeMix的方法）。随后仍需计算RKV矩阵以及W权重矩阵。然而，在这一特定步骤中假设输入x的空间大小为embed_size，则R矩阵的空间大小应与X一致；同时KV矩阵的空间大小由用户自行设定为hidden_size，并且权重矩阵W的形状被设定为[hidden_size, embed_size]。
通过配置较小规模的hidden_size有助于构建一个简化的TimeMix模型，在不影响性能的前提下仍能有效提升运行速度。当隐藏层大小等于嵌入维度时（即hidden_size=embed_size），该模型可被视为空间注意力机制而不进行位置编码，并等价于多头全连接层（即类似于FFN）。
具体计算公式如下所示
总结
总结
实证结果具体可见于原文内容中, 本文着重介绍其架构设计. 相比而言, RWKV表现出更高的流畅性, 同样在训练收敛速度上也更为迅速.
参考文献
[1] Which is better between pre-trained convolutions and pre-trained transformers? https://arxiv.org/pdf/2105.03322.pdf
MLP-Mixer: Such a completely based on all MLPs architecture designed for vision tasks, as introduced in the paper https://arxiv.org/pdf/2105.01601.pdf
Give Due Attention To MLPs https://arxiv.org/pdf/2105.08050.pdf
[4] Synthesizer: Rethinking Self-Attention in Transformer Models https://arxiv.org/abs/2005.00743
Revisiting the Attention Mechanism through Performers https://arxiv.org/abs/2009.14794
该方法：该高效的Transformer模型：...
该方法基于自注意力机制与线性复杂度算法的结合使用，在理论上显著地提高了计算效率，并实现了对大规模数据处理的能力提升
对线性注意力机制的研究：是否有必要在Attention中强制引入Softmax？ [https://spaces.ac.cn/archives/7546](https://link.zhihu.com/?target=https%3A//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
该方法提出了自适应多分辨率注意力机制及其线性复杂度
[编号] 不含注意力机制的Transformer模型 https://arxiv.org/abs/2105.14103
[12] Short Training Epochs and Long Testing Epochs: Self-attention Mechanisms Equipped with Linear Biases Facilitate Input Length Extrapolation.
访问该论文
由RWKV提出的一种兼具灵活与效率的线性变换器架构发布于知乎