Transformer 模型在几乎所有自然语言处理（NLP）任务中都带来了革命，但其在序列长度上的内存和计算复杂性呈二次方增长。相比之下，循环神经网络（RNNs）在内存和计算需求上呈线性增长，但由于并行化和可扩展性的限制，很难达到与 Transformer 相同的性能水平。本文提出了一种新颖的模型架构，Receptance Weighted Key Value（RWKV），将 Transformer 的高效可并行训练与 RNN 的高效推理相结合。实验证明，RWKV 的性能与相同规模的 Transformer 相当。

深度学习技术在人工智能领域取得了重大进展，在各种科学和工业应用中发挥了关键作用。这些应用通常涉及复杂的序列数据处理任务，包括自然语言理解、对话式人工智能、时间序列分析等，其中用到的技术主要包括循环神经网络（RNNs）、卷积神经网络（CNNs）和 Transformer 等。

不过，这些方法各自存在不同的缺点，从而限制了它们在某些场景下的效率。循环神经网络（RNNs）面临着梯度消失的问题，使得它们难以对长序列进行训练。此外，在训练过程中无法在时间维度上并行化，进而限制了其可扩展性。另一方面，卷积神经网络（CNNs）只擅长捕捉局部模式，在处理长程依赖方面还很欠缺，而这对于许多序列处理任务至关重要。

Transformer 模型由于其处理局部和长程依赖关系的能力以及可并行化训练的特点而成为一个强大的替代方案，如 GPT-3、ChatGPT、GPT-4、LLaMA 和 Chinchilla 等都展示了这种架构的能力，推动了自然语言处理领域的前沿。尽管取得了这些重大进展，Transformer 中固有的自注意力机制带来了独特的挑战，主要是由于其二次复杂度造成的。这种复杂性使得该架构在涉及长输入序列或资源受限情况下计算成本高昂且占用内存。这也促使了大量研究的发布，旨在改善 Transformer 的扩展性，但往往以牺牲一些特性为代价。

为了应对这些挑战，一个由 27 所大学、研究机构组成的开源研究团队，联合发表论文《 RWKV: Reinventing RNNs for the Transformer Era 》，文中介绍了一种新型模型：RWKV（Receptance Weighted Key Value），这是一种新颖的架构，有效地结合了 RNN 和 Transformer 的优点，同时规避了两者的缺点。RWKV 设计精良，能够缓解 Transformer 所带来的内存瓶颈和二次方扩展问题，实现更有效的线性扩展，同时保留了使 Transformer 在这个领域占主导的一些性质。

• 论文地址：https://arxiv.org/pdf/2305.13048.pdf

• RWKV 模型下载：https://huggingface.co/BlinkDL/rwkv-4-raven

• Demo 地址：https://www.codewithgpu.com/i/app/BlinkDL/ChatRWKV/RWKV-4-Raven-7B

本文利用线性注意力机制，允许将模型定义为 Transformer 或 RNN，从而在训练期间并行化计算，并在推理过程中保持恒定的计算和内存复杂性，使其成为第一个可扩展到数百亿参数的非 Transformer 架构。

RWKV 其中的一个特征是它能够提供并行训练和强大的可扩展性，类似于 Transformer。此外，该研究对 RWKV 中的注意力机制进行了重新阐述，引入了线性注意力的一个变体，避开了传统点积（dot-product）token 交互，转而采用更有效的通道导向注意力（ channel directed attention ）。这种方法与传统的 Transformer 架构形成了鲜明的对比，其中特定的 token 交互主导了注意力。在 RWKV 中，线性注意力的实施是无需近似的，这在效率上提供了显著的改进，并增强了可扩展性，详见表 1。

该研究表示，开发 RWKV 的主要动机是弥补神经网络架构在计算效率和表达能力之间的差距。它为处理涉及数十亿参数的大规模模型的任务提供了一个有希望且可行的解决方案，以极低的计算成本展现出强有力的竞争性。

实验结果表明，RWKV 可以成为一个有价值的工具，用于解决各个领域扩展和部署人工智能模型的各种挑战，特别是那些涉及序列数据处理的领域。RWKV 为下一代更可持续、计算效率更高的序列处理任务的 AI 模型铺平了道路。

总结而言，本文的贡献如下：

• 引入了 RWKV 网络架构，该架构结合了 RNN 和 Transformer 的优点，同时减轻了它们已知的限制。

• 本文提出了一个新的注意力机制重构，进而提出线性注意力，避开了与标准 Transformer 模型相关的二次复杂性。

• 本文在基准数据集上进行了一系列全面的实验，展示了 RWKV 在处理涉及大规模模型和长距离依赖任务上的性能、效率和可扩展性。

• 发布了预训练模型，其大小从 1.69 亿到 140 亿的参数不等，这些模型是在 Pile 上训练的。

值得注意的是，论文参与机构之一的 EleutherAI 表示：这篇论文还不是最终版本，后续会不断完善。

RWKV 模型

RWKV 架构的名称来源于时间混合和通道混合块中使用的四个主要模型元素，分别如下：

• R：Receptance 向量，用于接收以往信息；

• W：权重（weight）是位置权重衰减向量，是可训练的模型参数；

• K：键（Key）是类似于传统注意力中 K 的向量；

• V：值（Value）是类似于传统注意力中 V 的向量。

每一时间步的主要元素之间的交互是相乘增加的，具体如下图 2 所示。

架构细节

专栏文章内容及配图由作者撰写发布，仅供工程师学习之用，如有侵权或者其他违规问题，请联系本站处理。 联系我们