RWKV: Reinvenzione delle RNN per l'era dei Transformer

RWKV: Reinventing RNNs for the Transformer Era

May 22, 2023

Autori: Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Xiangru Tang, Bolun Wang, Johan S. Wind, Stansilaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Jian Zhu, Rui-Jie Zhu

cs.AI

Abstract

I Transformer hanno rivoluzionato quasi tutti i compiti di elaborazione del linguaggio naturale (NLP), ma soffrono di una complessità di memoria e computazionale che scala quadraticamente con la lunghezza della sequenza. Al contrario, le reti neurali ricorrenti (RNN) mostrano un scaling lineare nei requisiti di memoria e computazionali, ma faticano a raggiungere le stesse prestazioni dei Transformer a causa di limitazioni nella parallelizzazione e nella scalabilità. Proponiamo una nuova architettura di modello, Receptance Weighted Key Value (RWKV), che combina l'addestramento efficiente e parallelizzabile dei Transformer con l'inferenza efficiente delle RNN. Il nostro approccio sfrutta un meccanismo di attenzione lineare e ci permette di formulare il modello sia come Transformer che come RNN, parallelizzando i calcoli durante l'addestramento e mantenendo una complessità computazionale e di memoria costante durante l'inferenza, portando alla prima architettura non-Transformer scalata a decine di miliardi di parametri. I nostri esperimenti rivelano che RWKV performa alla pari con Transformer di dimensioni simili, suggerendo che lavori futuri possano sfruttare questa architettura per creare modelli più efficienti. Questo lavoro rappresenta un passo significativo verso la riconciliazione dei compromessi tra efficienza computazionale e prestazioni del modello nei compiti di elaborazione delle sequenze.

English

Transformers have revolutionized almost all natural language processing (NLP) tasks but suffer from memory and computational complexity that scales quadratically with sequence length. In contrast, recurrent neural networks (RNNs) exhibit linear scaling in memory and computational requirements but struggle to match the same performance as Transformers due to limitations in parallelization and scalability. We propose a novel model architecture, Receptance Weighted Key Value (RWKV), that combines the efficient parallelizable training of Transformers with the efficient inference of RNNs. Our approach leverages a linear attention mechanism and allows us to formulate the model as either a Transformer or an RNN, which parallelizes computations during training and maintains constant computational and memory complexity during inference, leading to the first non-transformer architecture to be scaled to tens of billions of parameters. Our experiments reveal that RWKV performs on par with similarly sized Transformers, suggesting that future work can leverage this architecture to create more efficient models. This work presents a significant step towards reconciling the trade-offs between computational efficiency and model performance in sequence processing tasks.