RWKV：为 Transformer 时代重塑 RNN
计算与语言
2023-12-12 v2 人工智能
摘要
Transformers 几乎革新了所有自然语言处理（NLP）任务，但存在随序列长度二次增长的内存与计算复杂度。相比之下，循环神经网络（RNNs）在内存与计算需求上呈线性增长，但因并行化与可扩展性方面的局限而难以匹敌 Transformers 的同等性能。我们提出一种新颖模型架构——Receptance Weighted Key Value（RWKV），它结合了 Transformers 的高效可并行训练与 RNNs 的高效推理。我们的方法利用线性注意力机制，并允许将模型表述为 Transformer 或 RNN，从而在训练时并行计算，并在推理时保持恒定的计算与内存复杂度。我们将模型扩展至多达 140 亿参数，这是迄今训练过的最大的稠密 RNN，并发现 RWKV 性能与同等规模的 Transformers 相当，表明未来工作可利用该架构创建更高效模型。本工作向调和序列处理任务中计算效率与模型性能间的权衡迈出了重要一步。
引用
@article{arxiv.2305.13048,
title = {RWKV: Reinventing RNNs for the Transformer Era},
author = {Bo Peng and Eric Alcaide and Quentin Anthony and Alon Albalak and Samuel Arcadinho and Stella Biderman and Huanqi Cao and Xin Cheng and Michael Chung and Matteo Grella and Kranthi Kiran GV and Xuzheng He and Haowen Hou and Jiaju Lin and Przemyslaw Kazienko and Jan Kocon and Jiaming Kong and Bartlomiej Koptyra and Hayden Lau and Krishna Sri Ipsit Mantri and Ferdinand Mom and Atsushi Saito and Guangyu Song and Xiangru Tang and Bolun Wang and Johan S. Wind and Stanislaw Wozniak and Ruichong Zhang and Zhenyuan Zhang and Qihang Zhao and Peng Zhou and Qinghua Zhou and Jian Zhu and Rui-Jie Zhu},
journal= {arXiv preprint arXiv:2305.13048},
year = {2023}
}