2024年4月9日,RWKV Foundation 发布了 RWKV 新架构的最新论文——《Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence》,论文由 Bo PENG 和 RWKV 开源社区成员共同完成。

论文地址:arxiv.org/abs/2404.05…

论文内容摘要

论文提出了两种新的 RWKV 架构:Eagle (RWKV-5) 和 Finch (RWKV-6),这两种序列模型在 RWKV-4 架构的基础上进行了改进。

对比 RWKV-4 ,新 RWKV 架构的改进包括多头矩阵值状态(multi-headed matrix-valued states)和动态递归机制(dynamic recurrence mechanism),这些改进提高了 RWKV 模型的表达能力,同时保持 RNN 的推理效率特征。

新的 RWKV 架构还使用一个包含 1.12 T tokens 的新多语言语料库和一个基于 greedy matching 的快速分词器,以增强 RWKV 的多语言能力。

基于新的架构训练了四个Eagle (RWKV-5)模型(分别是0.4B、1.5B、3B、7B 参数),以及两个 Finch 模型(分别是 1.6B 、3B参数)。

新架构的RWKV模型在各种基准测试中都实现了具有竞争力的性能,所有 Eagle 和 Finch 模型都基于 Apache 2.0 开源许可证下在 HuggingFace 上发布:
huggingface.co/RWKV

RWKV 当前的多模态工作包含 RWKV Music(音乐方向) 和 VisualRWKV(图像方向)。

未来 RWKV 的工作将聚焦于扩大并改进训练数据集,以及训练和发布更大规模的 Finch (RWKV-6) 模型,例如 7B 和 14B 参数,并通过 Mixture of Experts 降低推理和训练成本。

新架构升级细节

Eagle (RWKV-5) 通过使用富有表现力的多头矩阵值状态(而不是向量值状态)、重新构造的接受态和额外的门控机制,改进了从 RWKV-4 中学习到的架构和学习衰减进度。

Finch (RWKV-6) 通过引入新的数据相关函数来进一步改进架构的表现能力和灵活性,用于时间混合和令牌移位模块,包括参数化线性插值。

此外,Finch (RWKV-6) 提出了对低秩自适应函数的新用法,以使可训练的权重矩阵能够以一种上下文相关的方式有效地增强学习到的数据衰减向量。

最后,RWKV新架构引入了一种新的分词器 ”RWKV World Tokenizer“,和一个新数据集“RWKVWorld v2”(1.12 T tokens),两者用于提高 RWKV 模型在多语言和代码数据上的性能。

  • 新的分词器 ”RWKV World Tokenizer“ 包含不常见语言的词汇,并且通过基于 Trie 的 greedy matching 进行快速分词。
  • 新数据集 “RWKV World v2” 是一个新的多语言 1.12 T tokens 数据集,取自各种手工选择的公开可用数据源,数据组成约 70% 是英语,15% 是多语言数据,15% 是代码。

新架构基准测试

RWKV对新架构和模型进行了各类语言建模实验和基准测试,以下为部分基准测试效果展示。

MQAR测试结果

MQAR (Multiple Query Associative Recall) 任务是一种用于评估语言模型的任务,旨在测试模型在多次查询情况下的联想记忆能力。在这类任务中,模型需要通过给定的多个查询来检索相关的信息。

MQAR任务的目标是衡量模型在多次查询下检索信息的能力,以及其对不同查询的适应性和准确性。

下图为RWKV-4、Eagle、 Finch 和其他非 Transformer 架构的 MQAR 任务测试结果。

RWKV-5/6(Eagle & Finch):基于矩阵值状态和动态递归的新模型架构

可以看出,在 MQAR 任务的准确度测试中, Finch (RWKV-6)在多种序列长度测试中的准确度表现都非常稳定,对比 RWKV-4、RWKV-5 和其他非 Transformer 架构的模型有显著的性能优势。

长上下文实验

在 PG19 测试集上测试了从 2048 tokens 开始的 RWKV-4、Eagle 和 Finch 的 loss 与序列位置。(所有模型均基于上下文长度 4096 进行预训练。)

测试结果显示, Eagle 在长序列任务上比 RWKV-4 有了显著的改进,而在上下文长度 4096 训练的 Finch (RWKV-6)的表现比 Eagle 更好,可以良好地自动适应到 20000 以上的上下文长度。

RWKV-5/6(Eagle & Finch):基于矩阵值状态和动态递归的新模型架构

速度和显存基准测试

速度和内存基准测试比较了 Finch、Mamba 和 Flash Attention 的类 Attention 内核的速度和显存利用率。

RWKV-5/6(Eagle & Finch):基于矩阵值状态和动态递归的新模型架构

RWKV-5/6(Eagle & Finch):基于矩阵值状态和动态递归的新模型架构

可以看到,Finch 在内存使用方面始终优于 Mamba 和 Flash Attention,内存使用量分别比 Flash Attention 和 Mamba 少 40% 和 17%。

未来 RWKV 团队会对 Finch 的 CUDA 实现做进一步优化(包括算法改进),带来速度的提升和更大的并行化。

其他测试,如语言模型评估测试基准、Bamboo 基准测试等,可以在论文中找到。

新架构的多语言任务表现

日语

RWKV-5/6(Eagle & Finch):基于矩阵值状态和动态递归的新模型架构

西班牙语

RWKV-5/6(Eagle & Finch):基于矩阵值状态和动态递归的新模型架构

阿拉伯语

RWKV-5/6(Eagle & Finch):基于矩阵值状态和动态递归的新模型架构

日语-英语

RWKV-5/6(Eagle & Finch):基于矩阵值状态和动态递归的新模型架构

基于新架构的模型

Eagle (RWKV-5)架构目前发布了四种参数的模型:

  • RWKV-5-Eagle 0.4B
  • RWKV-5-Eagle 1.5B
  • RWKV-5-Eagle 3B
  • RWKV-5-Eagle 7B

Eagle (RWKV-5)模型下载链接:huggingface.co/BlinkDL/rwk…

Finch (RWKV-6) 架构目前发布了两种参数的模型:

  • RWKV-6-Finch 1.6B
  • RWKV-6-Finch 3B

Finch (RWKV-6)模型下载链接:huggingface.co/BlinkDL/rwk…

未来工作

新 RWKV 架构的多语言训练语料库“RWKV World v2”只有 1.12T tokens ,这比 LLaMA2 等当代模型的训练数据量要小得多。因此,RWKV 的下一步重要工作是扩展 训练语料库,使其更加多样化,这是改进模型性能的关键事项。

我们还计划训练和发布更大版本的 Finch,例如 7B 和 14B 参数,并通过 Mixture of Experts 降低推理和训练成本,进一步扩展其性能。

附录

RWKV-6-Finch 最新架构图:

RWKV-5/6(Eagle & Finch):基于矩阵值状态和动态递归的新模型架构