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と同等の性能を達成するのが難しい。本研究では、Transformerの効率的な並列学習とRNNの効率的な推論を統合した新しいモデルアーキテクチャ、Receptance Weighted Key Value（RWKV）を提案する。本手法は線形アテンション機構を活用し、モデルをTransformerとしてもRNNとしても定式化可能であるため、学習時には計算を並列化しつつ、推論時には計算およびメモリの複雑性を一定に保つことが可能となる。我々は、これまでにない規模の140億パラメータにまでモデルをスケーリングし、これまでに訓練された最大の密結合型RNN（dense RNN）を実現した。その結果、RWKVは同規模のTransformerと同等の性能を発揮することが確認され、今後の研究においてこのアーキテクチャを活用することで、より効率的なモデルの構築が可能であることが示唆された。本研究は、シーケンス処理タスクにおける計算効率とモデル性能のトレードオフを是正する上で、重要な一歩を踏み出したものである。