RWKV 介绍——一个兼具 Transformer 优点的 RNN

ChatGPT 和聊天机器人驱动的应用在自然语言处理 (NLP) 领域引起了广泛关注。社区一直在寻求强大、可靠的开源模型用于其应用和用例。这些强大模型的兴起源于 Vaswani 等人于 2017 年首次引入的基于 Transformer 的模型的普及和广泛采用。这些模型显著优于此前基于循环神经网络 (RNN) 的最先进 NLP 模型，后者在该论文发表后被认为已“过时”。通过这篇博文，我们将介绍 RWKV 这一新架构的集成，它结合了 RNN 和 Transformer 的优点，并已于近期集成到 Hugging Face 的 Transformers 库中。

RWKV 项目概述

RWKV 项目由 BlinkDL (Bo Peng) 发起并主导，他积极参与并维护该项目。社区在官方 Discord 频道中组织起来，不断在性能（RWKV.cpp、量化等）、可扩展性（数据集处理和抓取）和研究（聊天微调、多模态微调等）等各个主题上增强项目成果。用于训练 RWKV 模型的 GPU 由 Stability AI 捐赠。

您可以通过加入 官方 Discord 频道 来参与，并通过以下两篇博文了解更多关于 RWKV 背后的通用思想：https://johanwind.github.io/2023/03/23/rwkv_overview.html / https://johanwind.github.io/2023/03/23/rwkv_details.html

Transformer 架构 vs RNN

RNN 架构是首批广泛用于处理数据序列的神经网络架构之一，与采用固定大小输入的经典架构不同。它将当前“token”（即数据流的当前数据点）和先前的“状态”作为输入，并计算预测的下一个 token 和预测的下一个状态。然后，新状态用于计算下一个 token 的预测，依此类推。RNN 还可以以不同的“模式”使用，从而使其能够应用于不同的场景，正如 Andrej Karpathy 的博文 所述，例如一对一（图像分类）、一对多（图像字幕）、多对一（序列分类）、多对多（序列生成）等。

由于 RNN 在每个步骤都使用相同的权重来计算预测，因此它们由于梯度消失问题而难以记忆长序列信息。为了解决这一限制，人们引入了 LSTM 或 GRU 等新架构。然而，Transformer 架构被证明是迄今为止解决此问题最有效的。

在 Transformer 架构中，输入 token 在自注意力模块中同时进行处理。token 首先使用 Query、Key 和 Value 权重线性投影到不同的空间。生成的矩阵直接用于计算注意力分数（通过 softmax，如下所示），然后乘以 Value 隐藏状态以获得最终的隐藏状态。这种设计使架构能够有效地缓解长序列问题，并且与 RNN 模型相比，其推理和训练速度更快。

在训练过程中，Transformer 架构相对于传统的 RNN 和 CNN 具有多项优势。其中最显著的优势是其学习上下文表示的能力。与一次处理一个词的 RNN 和 CNN 不同，Transformer 架构将输入序列作为一个整体进行处理。这使得它能够捕获序列中词语之间的长距离依赖关系，这对于语言翻译和问答等任务特别有用。

在推理过程中，RNN 在速度和内存效率方面具有一些优势。这些优势包括简单性（只需矩阵-向量操作）和内存效率（内存需求在推理过程中不会增加）。此外，由于计算只作用于当前 token 和状态，计算速度与上下文窗口长度保持不变。

RWKV 架构

RWKV 的灵感来源于 Apple 的 Attention Free Transformer。该架构经过精心简化和优化，使其能够转换为 RNN。此外，还添加了一些技巧，例如 TokenShift 和 SmallInitEmb (技巧列表详见 官方 GitHub 仓库的 README)，以提升其性能与 GPT 相当。没有这些技巧，模型的性能将大打折扣。目前，训练方面已有基础设施可将训练扩展至 14B 参数，并且 RWKV-4（截至目前的最新版本）中已迭代修复了一些问题，例如数值不稳定性。

RWKV：RNN 和 Transformer 的结合

如何将 Transformer 和 RNN 的优点结合起来？基于 Transformer 的模型的主要缺点是，当上下文窗口大于一定值时，运行模型可能变得具有挑战性，因为注意力分数是针对整个序列同时计算的。

RNN 天然支持非常长的上下文长度——仅受训练中看到的上下文长度限制，但通过精心编码，这可以扩展到数百万个 token。目前，RWKV 模型在 8192 (ctx8192) 的上下文长度上进行训练，它们的运行速度与 ctx1024 模型相同，并且需要相同的 RAM 量。

传统 RNN 模型的主要缺点以及 RWKV 的不同之处

1. 传统的 RNN 模型无法利用很长的上下文（LSTM 在用作 LM 时只能管理大约 100 个 token）。然而，RWKV 可以利用数千甚至更多 token，如下所示

2. 传统 RNN 模型在训练时无法并行化。RWKV 类似于“线性化 GPT”，它的训练速度比 GPT 快。

通过将这两种优势结合到单一架构中，我们希望 RWKV 能够成长为比其各部分之和更强大的模型。

RWKV 注意力机制的公式

该模型架构与经典基于 Transformer 的模型非常相似（即一个嵌入层、多个相同的层、层归一化以及一个因果语言建模头来预测下一个 token）。唯一的区别在于注意力层，它与传统的基于 Transformer 的模型完全不同。

为了更全面地理解注意力层，我们建议深入阅读 Johan Sokrates Wind 的博文中提供的详细解释。

现有检查点

纯语言模型：RWKV-4 模型

最受欢迎的 RWKV 模型参数量从约 1.7 亿到 14 亿不等。根据 RWKV 概述 博文，这些模型已在 Pile 数据集上进行训练，并在不同基准上与其他 SoTA 模型进行评估，它们的表现相当不错，结果非常可比。

指令微调/聊天版本：RWKV-4 Raven

Bo 还训练了一个 RWKV 架构的“聊天”版本，即 RWKV-4 Raven 模型。它是一个 RWKV-4 pile 模型（在 The Pile 数据集上预训练的 RWKV 模型），在 ALPACA、CodeAlpaca、Guanaco、GPT4All 和更多数据集上进行了微调。该模型有多个版本，训练语言不同（仅英语、英语+中文+日语、英语+日语等）和大小不同（1.5B 参数、7B 参数、14B 参数）。

所有 HF 转换后的模型都可以在 Hugging Face Hub 的 RWKV 组织中找到。

🤗 Transformers 集成

该架构已通过 此拉取请求 添加到 transformers 库中。截至本文撰写之时，您可以通过从源代码安装 transformers 或使用库的 main 分支来使用它。该架构与库紧密集成，您可以像使用任何其他架构一样使用它。

下面我们来看一些示例。

文本生成示例

要根据输入提示生成文本，您可以使用 pipeline 来生成文本。

from transformers import pipeline

model_id = "RWKV/rwkv-4-169m-pile"

prompt = "\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese."

pipe = pipeline("text-generation", model=model_id)
print(pipe(prompt, max_new_tokens=20))
>>> [{'generated_text': '\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese.\n\nThe researchers found that the dragons were able to communicate with each other, and that they were'}]

或者您可以从下面的代码片段开始运行

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile")
tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")

prompt = "\nIn a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese."

inputs = tokenizer(prompt, return_tensors="pt")
output = model.generate(inputs["input_ids"], max_new_tokens=20)
print(tokenizer.decode(output[0].tolist()))
>>> In a shocking finding, scientist discovered a herd of dragons living in a remote, previously unexplored valley, in Tibet. Even more surprising to the researchers was the fact that the dragons spoke perfect Chinese.\n\nThe researchers found that the dragons were able to communicate with each other, and that they were

使用 Raven 模型（聊天模型）

您可以像使用 Alpaca 风格一样提示聊天模型，下面是一个示例

from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "RWKV/rwkv-raven-1b5"

model = AutoModelForCausalLM.from_pretrained(model_id).to(0)
tokenizer = AutoTokenizer.from_pretrained(model_id)

question = "Tell me about ravens"
prompt = f"### Instruction: {question}\n### Response:"

inputs = tokenizer(prompt, return_tensors="pt").to(0)
output = model.generate(inputs["input_ids"], max_new_tokens=100)

print(tokenizer.decode(output[0].tolist(), skip_special_tokens=True))
>>> ### Instruction: Tell me about ravens
### Response: RAVENS are a type of bird that is native to the Middle East and North Africa. They are known for their intelligence, adaptability, and their ability to live in a variety of environments. RAVENS are known for their intelligence, adaptability, and their ability to live in a variety of environments. They are known for their intelligence, adaptability, and their ability to live in a variety of environments.

根据 Bo 的说法，更佳的指令技术详见 此 Discord 消息（点击前请务必加入频道）

| |

权重转换

任何用户都可以通过简单地运行 transformers 库中提供的转换脚本，将原始 RWKV 权重转换为 HF 格式。首先，将“原始”权重推送到 Hugging Face Hub（我们将其仓库命名为 RAW_HUB_REPO，原始文件命名为 RAW_FILE），然后运行转换脚本

python convert_rwkv_checkpoint_to_hf.py --repo_id RAW_HUB_REPO --checkpoint_file RAW_FILE --output_dir OUTPUT_DIR

如果您想将转换后的模型推送到 Hub（例如，在 dummy_user/converted-rwkv 下），请先不要登录 huggingface-cli login，然后再推送模型，然后运行

python convert_rwkv_checkpoint_to_hf.py --repo_id RAW_HUB_REPO --checkpoint_file RAW_FILE --output_dir OUTPUT_DIR --push_to_hub --model_name dummy_user/converted-rwkv

未来工作

多语言 RWKV

Bo 目前正在研究一个多语言语料库以训练 RWKV 模型。最近 发布了一个新的多语言分词器。

社区导向和研究项目

RWKV 社区非常活跃，并致力于多个后续方向的研究，一个酷炫的项目列表可以在 Discord 上的一个专用频道中找到（点击链接前请务必加入频道）。还有一个专门针对此架构研究的频道，欢迎加入并贡献！

模型压缩和加速

由于 RWKV 只需进行矩阵-向量运算，因此它是非标准和实验性计算硬件（例如光子处理器/加速器）的理想选择。

因此，该架构也可以自然地受益于经典的加速和压缩技术（例如 ONNX、4 位/8 位量化等），我们希望这能随着架构与 Transformer 的集成而普及给开发人员和从业者。

RWKV 在不久的将来也可以受益于 optimum 库中提出的加速技术。其中一些技术在 rwkv.cpp 仓库 或 rwkv-cpp-cuda 仓库中有所强调。

致谢

Hugging Face 团队衷心感谢 Bo 和 RWKV 社区抽出时间并回答我们关于该架构的问题。我们还要感谢他们的帮助和支持，并期待在 HF 生态系统中看到更多 RWKV 模型的应用。我们还要感谢 Johan Wind 在 RWKV 方面撰写的博文，这极大地帮助我们理解了该架构及其潜力。最后，我们还要特别感谢 ArEnSc 发起了最初的 transformers PR。此外，还要特别感谢 Merve Noyan、Maria Khalusova 和 Pedro Cuenca 慷慨地审阅了这篇博文，使其变得更好！

引用

如果您在工作中使用了 RWKV，请使用 以下 cff 引用。