RWKV-5/6（Eagle和Finch）：基于RNN的序列模型新突破
RWKV-5/6（Eagle和Finch）：基于RNN的序列模型新突破
RWKV-5（Eagle）和RWKV-6（Finch）是基于RWKV-4架构的序列模型改进版本。通过引入多头矩阵值状态和动态循环机制，这些模型在保持RNN推断效率的同时提高了表达能力。本文详细介绍了Eagle和Finch的架构设计、训练方法、数据集构建以及实验结果。
0. 摘要
我们介绍了Eagle（RWKV-5）和Finch（RWKV-6），这是在RWKV-4架构基础上改进的序列模型。我们的架构设计进步包括多头矩阵值状态和动态循环机制，这些改进提高了表达能力，同时保持了RNN的推断效率特性。我们引入了一个新的多语言语料库，拥有1120B个token，并基于贪婪匹配设计了一个快速的分词器，以增强多语言性能。我们训练了四个Eagle模型，参数范围从0.46到75亿，以及两个Finch模型，参数分别为16亿和31亿，并发现它们在各种基准测试中取得了竞争性能。我们在HuggingFace上以Apache 2.0许可证发布了我们的所有模型。
模型位于：https://huggingface.co/RWKV
训练代码位于：https://github.com/RWKV/RWKV-LM
推理代码位于：https://github.com/RWKV/ChatRWKV
时间并行训练代码位于：https://github.com/RWKV/RWKV-infctx-trainer
3. Eagle/Finch 架构
我们通过两个步骤对RWKV架构进行了优化，并观察到每一步都有显著的建模改进。与基线RWKV-4相比，Eagle增加了矩阵值注意力状态、注意力头上的LayerNorm、SiLU注意力门控和改进的初始化。它还取消了接受的Sigmoid激活。Finch进一步将数据依赖性应用于衰减调度和token移位。
核心架构与RWKV-4类似，由一系列堆叠的残差块组成，形状类似于传统的Transformer。根据（Tolstikhin等，2021年）的符号表示，每个块包含一个Pre-LayerNorm时间混合子层，后跟一个Pre-LayerNorm通道混合子层，如图1左侧所示。这些对应于Transformer的传统注意力和前馈网络子层。有关我们的训练实现及与RWKV-4的差异的更多细节，请参见附录B，有关速度和内存基准测试，请参见第9节。
4. 方法
在本节中，我们使用D来表示模型维度，除非明确说明，本节中出现的所有向量都是维度D/h，其中h表示注意力头的数量，属于R^(D/h)。为了简洁和简便起见，我们以每个头的方式显示计算，省略头索引。我们使用的约定是，除非明确转置，所有向量都是行向量，因此所有矩阵都在右侧操作。
4.1 Eagle
4.1.1 Eagle token 移位
我们采用了前一版RWKV中的token移位技术，类似于大小为2的一维因果卷积，如图1中下部所示。为了更好地介绍token移位技术，我们定义一些符号。在RWKV-4和Eagle token移位中使用的xt和x_(t−1)之间的线性插值（lerp）定义为：
其中每个μ□∈R^D是一个可学习的向量。token移位允许模型独立且独特地为每个头每个时间步的接受、键、值和门向量（r、k、v和g）的每个通道分配新旧信息的数量。这使得在单个层内可以形成归纳头，因为即使是单个头也可以直接将过去和当前的token数据分别累积到这些向量的不同子空间中。
4.1.2 Eagle 时间混合
Eagle时间混合的公式可以写成如下形式：
其中LayerNorm分别作用于h个头的每一个，这也等价于在h个组上进行的Group-Norm操作。值得注意的是，w是从w=exp(−exp(ω))获得的，其中ω∈R^(D/h)是实际的按头可训练参数。这确保w落在区间(0,1)内，保证了diag(w)是一个收缩矩阵。
wkvt注意力计算也可以用循环形式来写成：
RWKV的wkv项可以被视为线性注意力中归一化的k^T·v项的基于衰减的等效项。值得注意的是，对于给定的头j，循环状态s是k^T·v的总和，其中每个s通道在每个时间步都会由w的相应通道进行衰减。在应用接收向量r、门控g和输出权重w之前，会将每个通道学习的增益u与当前token的k^T·v相乘，并与状态求和，如图1右上方所示。这使得当前token相对于包含在衰减状态历史中的过去token的总和具有特殊的处理方式。接收向量乘以这个总和，就像在Linear Attention中的query项一样。
4.1.3 通道混合
在Eagle和Finch中，通道混合模块与之前的RWKV-4架构相同，只是隐藏维度从4D降低到了3.5D。这个降维考虑到了在Eagle时间混合中的新门控权重，以确保与先前模型在相同层数和嵌入维度上具有等参数关系。尽管在Finch中添加了一小部分新的LoRA权重参数，但我们并没有进一步降低隐藏维度。通道混合的公式与RWKV-4相同，但我们在此重新陈述它们以确保符号一致性，使用方程式3中的线性插值：
4.2 Finch
4.2.1 Finch token 移位
在Finch token移位中使用的xt和x_(t−1)之间的数据依赖线性插值（ddlerp）定义为：
其中μx和每个λ□引入了一个维度为D的可训练向量，每个A□∈R^(D×32)，B□∈R^(32×D)引入了新的可训练权重矩阵。对于下面所示的LoRA_ω的特殊情况，我们引入了双倍大小的可训练权重矩阵Aω∈R^(D×64)，Bω∈R^(64×D)。图1右下方可以找到示意图。请注意，预期未来的7B及更大的雀模型将进一步将这些权重矩阵的大小增加一倍或更多。这种增强了数据相关性的新token移位形式旨在将模型的能力扩展到超出RWKV-4/Eagle风格的token移位，以便每个通道分配的新旧数据量取决于当前和先前时间步的输入。
4.2.2 Finch 时间混合
wkv_t注意力计算也可以用循环的方式来写：
不像在Eagle中，wt在这里不是跨序列静态的（图1左边和右上方的虚线箭头）。这是Finch中衰减的核心变化，因为现在wt的每个通道可以以数据相关的方式随时间变化，而不是以前的固定学习向量。
上述新的LoRA机制用于获取学习向量（与Eagle中相同），并以低成本的方式用传入的输入确定的额外offset对其进行增强。请注意，LoRA过程本身使用的是Eagle风格的token移位值作为其输入，而不仅仅是最新的token。新的时间变化衰减wt更进一步，之后再次应用LoRA。直观地说，这是token移位的二阶变体，允许wt的每个通道根据当前和先前token的混合以及混合本身的方面而变化。
5. RWKV World Tokenizer
在语言建模中，标记化很重要，因为它会影响token之间的学习关系以及基于这些模式生成新文本。然而，构建单个语义块所需的token数量通常在非欧洲和其他代表性较低的语言中分布非常不均匀。基于字节对编码（Byte-pair-encoding，BPE）的tokenizer在训练时考虑到这种不平等会导致不仅在代表性较低的语言方面性能较低，而且还会产生不必要的经济成本，例如推理（Ahia等人，2023年）和持续的预训练与扩展的词汇表（Lin等人，2024年；Sasaki等人，2023年）。为了解决这些问题，我们从多个词汇文件中手动选择标记，以确保非欧洲语言得到充分代表。
为了构建标记器的词汇表，我们合并了以下tokenizer的词汇表，然后手动选择非欧洲语言的标记：
- GPT-NeoX-20B (Black等人，2022年): https://huggingface.co/EleutherAI/gpt-neox-20b
- GPT2 (Radford等人，2019年): https://huggingface.co/openai-community/gpt2
- tiktoken的cl100k_base: https://github.com/openai/tiktoken
- Llama2 (Touvron等人，2023年): https://huggingface.co/meta-llama/Llama-2-7b-hf
- Bloom (Workshop等人，2023年): https://huggingface.co/bigscience/bloom
该tokenizer的词汇表大小为V=65536，编号从0到65535，其中token按其字节长度排列。以下是简要概述：
- token 0：表示文本文档之间的边界，称为
或 。此token不编码任何特定内容，仅用于文档分隔。 - token 1-256：由字节编码组成（token k编码字节k-1），其中token 1-128对应于标准ASCII字符。
- token 257-65529：在UTF-8中具有至少2字节长度的token，包括单词、前缀和后缀、重音字母、汉字、韩文字母、平假名、片假名和表情符号。例如，汉字分配从token 10250到18493。
- token 65530-65535：保留用于未来使用的token。
这些指定旨在增强tokenizer在多语言语料库以及编程语言源代码上的效率。
该tokenizer是通过前缀树（Trie）实现的，以提高速度同时保持简单性。编码是通过从左到右将输入字符串与词汇表中的最长元素匹配来执行的。我们注意到，我们的tokenizer词汇表的构建是为了减轻对次要语言造成的不必要负担，而这种负担是朴素BPE和相关方法引起的。
6. RWKV World v2 数据集
我们在新的RWKV World v2数据集上训练我们的模型，这是一个新的多语言1120B token数据集，来自各种手动选择的公开可用数据源。该数据集旨在超越当前广泛用于训练LLM的许多数据集对英语的重视。我们这样做是为了支持全球绝大多数非英语母语人士的使用，以提高模型响应中的表现，并使迁移学习成为可能，使我们的模型可以跨文化和地区应用知识。我们非常注重事实知识和代码，但也注重文化作品，包括故事、书籍、字幕和对话。源数据大约有70%是英语，15%是多语言，15%是代码。我们在附录D中详细描述了我们数据集的组成部分。
7. 预训练模型
我们已经预训练并公开发布了六个受Apache 2.0许可的Eagle and Finch模型：Eagle 0.4B、Eagle 1.5B、Eagle 3B、Eagle 7B、Finch 1.6B和Finch 3B。所有模型都是在1120B token的RWKV World v2多语言语料库上进行训练的。有关详细的参数计数和FLOPs计算，请参见附录E。
8. 语言建模实验
9. 速度和存储 Benchmark
10. 多模态实验
11. 结论
在本工作中，我们介绍了Eagle（RWKV-5）和Finch（RWKV-6），通过整合多头矩阵值状态和动态数据驱动的递归机制，取得了基于RNN的语言模型的实质性进展。这些模型在MQAR和各种语言基准测试中表现出色，挑战了传统Transformer架构的主导地位，同时保留了关键的RNN优势。通过在Apache 2.0许可下公开提供模型并在广泛的多语言语料库上进行训练，我们的工作不仅推动了语言模型的能力发展，还强调了社区的可访问性和在各个领域的适用性。尽管我们意识到未来存在计算和伦理挑战，但我们希望Eagle和Finch的高效新架构和广泛可用性将有助于推动语言建模的边界，并为未来的创新铺平道路。
限制
Eagle和Finch模型在某些方面存在不足，可以在未来的工作中得到缓解和解决。我们尝试将Eagle用作大规模文本嵌入基准测试（theMassive Text Embedding Benchmark，MTEB）的嵌入模型，但未能获得强大的嵌入性能。我们认为它的状态是上下文的高质量嵌入，但需要适当的方法来聚合信息内容。我们将这留给未来的工作。
由于我们的训练语料库包含来自GPT-3.5和ChatGPT的一些合成数据，我们发布的模型表现出与ChatGPT类似的行为，并将模仿ChatGPT的对话风格和语调。例如，模型偶尔会声称它是由OpenAI训练的。然而，这不是RWKV架构的一般特性，而是数据和训练过程的特定结果。
未来工作
我们的1120B token的多语言训练语料库远远小于当今模型（如LLaMA2）的训练数据规模，扩展我们的训练语料库以使其更加多样化和广泛化是提高模型性能的关键优先事项。我们还计划训练和发布更大的Finch版本，如7B和14B参数，并通过专家混合（Mixture of Experts）进一步降低推理和训练成本来扩展其性能。