HarmonyOS 鸿蒙Next：结合RNN与Transformer双重优点，深度解析大语言模型RWKV

HarmonyOS 鸿蒙Next：结合RNN与Transformer双重优点，深度解析大语言模型RWKV 摘要：RWKV在多语言处理、小说写作、长期记忆保持等方面表现出色，可以主要应用于自然语言处理任务，例如文本分类、命名实体识别、情感分析等。

本文分享自华为云社区《【云驻共创】昇思MindSpore技术公开课 RWKV 模型架构深度解析》，作者：Freedom123。

一、前言

Transformer模型作为一种革命性的神经网络架构，于2017年由Vaswani等人提出，并在诸多任务中取得了显著的成功。Transformer的核心思想是自注意力机制，通过全局建模和并行计算，极大地提高了模型对长距离依赖关系的建模能力。但是Transformer在处理长序列时面临内存和计算复杂度的问题，因为其复杂度与序列长度呈二次关系一直未业内人员所诟病。今天我们学习的RWKV，它作为对Transformers模型的替代，正在引起越来越多的开发人员的关注。RWKV模型以简单、高效、可解释性强等特点，成为自然语言处理领域的新宠。下面让我们一起来学习RWKV模型。

二、RWKV简介

RWKV（Receptance Weighted Key Value）是一个结合了RNN与Transformer双重优点的模型架构，由香港大学物理系毕业的彭博首次提出。其名称源于其 Time-mix 和 Channel-mix 层中使用的四个主要模型元素：R（Receptance）：用于接收以往信息；W（Weight）：是位置权重衰减向量，是可训练的模型参数；K（Key）：是类似于传统注意力中 K 的向量；V（Value）：是类似于传统注意力中 V 的向量。

RWKV模型作为一种革新性的大型语言模型，结合了RNN的线性复杂度和Transformer的并行处理优势，引入了Token shift和Channel Mix机制来优化位置编码和多头注意力机制，解决了传统Transformer模型在处理长序列时的计算复杂度问题。RWKV在多语言处理、小说写作、长期记忆保持等方面表现出色，可以主要应用于自然语言处理任务，例如文本分类、命名实体识别、情感分析等。

三、RWKV模型的演进

RWKV模型之所以发展到今天的结构经历了五个阶段，从RNN结构到LSTM结构，到GRU结构，到GNMT结构，到Transformers结构，最后到RMKV结构，下面我们一一来学习每种模型结构，务必做到对模型结构都有一个清晰的认识。

四、RWKV模型

RWKV是一个结合了RNN与Transformer双重优点的模型架构，是一个RNN架构的模型，但是可以像transformer一样高效训练。RWKV 模型通过 Time-mix 和 Channel-mix 层的组合，以及 distance encoding 的使用，实现了更高效的 Transformer 结构，并且增强了模型的表达能力和泛化能力。Time-mix 层与 AFT（Attention Free Transformer）层相似，采用了一种注意力归一化的方法，以消除传统 Transformer 模型中存在的计算浪费问题。Channel-mix 层则与 GeLU（Gated Linear Unit）层相似，使用了一个 gating mechanism 来控制每条通道的输入和输出。另外，RWKV 模型采用了类似于 AliBi 编码的位置编码方式，将每个位置的信息添加到模型的输入中，以增强模型的时序信息处理能力。这种位置编码方式称为 distance encoding，它考虑了不同位置之间的距离衰减特性，RWKV结构如下图所示：

五、RWKV模型代码阅读

1. RWKV模型推理代码

代码解释：

1. 引入库
2. 对输出进行校验
3. 加载RWKV/rwkv-4-169m-pile模型，并且输入提示词
4. 运行模型，解码生成内容
5. 期望输出与真实输出内容进行校验

2. Channel Mixing模块代码

x通道混合层接受与此标记对应的输入，以及x与前一个标记对应的输入，我们称之为last_x。last_x存储在这个 RWKV 层的state。其余输入是学习RWKV 的 parameters。首先，我们使用学习的权重对x和进行线性插值last_x。我们将此插值x作为输入运行到具有平方 relu 激活的 2 层前馈网络，最后与另一个前馈网络的 sigmoid 激活相乘（在经典 RNN 术语中，这称为门控）。

3. Time mixing模块代码

时间混合的开始类似于通道混合，通过将此标记的插入x到最后一个标记的x。然后我们应用学到的矩阵以获得“key”, “value” and “receptance”向量。

六、与其他模型的比较

1. 复杂度对比

从和Transformer，Reformer，Performer，Linear Transformers，AFT-full，AFT-local，MEGA等模型的复杂度比较中可以看的出来，RWKV模型的时间复杂度和空间负责度都是最低的，分别为O(Td)和O(d)，其中T 表示序列长度，d 表示特征维度，c 表示 MEGA 的二次注意力块大小。

2. 精度对比

RWKV 似乎可以像 SOTA transformer一样缩放。至少多达140亿个参数。在同等规模参数中，RWKV-4系列与Pythia和GPT-J比都是很有优势的。

3. 推理速度和内存占用

RWKV网络与不同类型的Transformer性能的实验结果对比如下图所示。RWKV时间消耗随序列长度是线性增加，且时间消耗远小于各种类型的Transformer。

七、小结

本节我们学习了RWKV模型，我们掌握了RWKV模型结构的整个演进过程，从最初的RNN结构，到LSTM结构，到GRU结构，到GNTM模型，到Transformers模型，最后到RWKV模型，我们学习了每种模型结构出现的原因，以及其对应的优势和不足。接下来，我们学习了RWKV模型，Time Mixing模块和Channel Mixing模块。我们通过学习RWKV模型的python代码，对RWKV模型从复杂度，精度，推理速度，内存占用等四个维度和其他模型进行了对比。

通过本节学习，我们对RWKV模型有了一个全面的认识，RWKV模型正在作为一颗在大模型领域的新星正在受到越来越多社区开发者的关注，希望RWKV模型在接下来的版本迭代过程中能给大家带来更多的惊喜。

点击关注，第一时间了解华为云新鲜技术