17 天前
RWKV：面向Transformer时代的RNN革新
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, Rui-Jie Zhu
摘要
Transformer模型已彻底革新了几乎所有自然语言处理（NLP）任务，但其在内存和计算复杂度方面存在与序列长度呈二次方增长的瓶颈。相比之下，循环神经网络（RNN）在内存和计算需求上表现出线性增长，具备更高的计算效率，但由于并行化能力受限，其可扩展性较差，难以达到Transformer模型的性能水平。为此，我们提出一种新型模型架构——可接受度加权键值网络（Receptance Weighted Key Value, RWKV），该架构融合了Transformer的高效并行训练能力与RNN的高效推理特性。我们的方法基于线性注意力机制，使模型在训练阶段可被建模为Transformer，从而实现并行计算；在推理阶段则表现为RNN结构，保持恒定的计算与内存复杂度。我们成功训练了参数规模高达140亿的RWKV模型，成为迄今为止参数量最大的稠密RNN架构。实验结果表明，RWKV在性能上可与同规模的Transformer模型相媲美，表明该架构有望在未来推动更高效序列建模模型的发展。本研究为解决序列处理任务中计算效率与模型性能之间的权衡问题迈出了重要一步。