Red de Retentiva: Un Sucesor del Transformer para Modelos de Lenguaje a Gran Escala

Retentive Network: A Successor to Transformer for Large Language Models

July 17, 2023

Autores: Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, Furu Wei

cs.AI

Resumen

En este trabajo, proponemos Retentive Network (RetNet) como una arquitectura base para modelos de lenguaje a gran escala, logrando simultáneamente paralelismo en el entrenamiento, inferencia de bajo costo y un buen rendimiento. Teóricamente, derivamos la conexión entre recurrencia y atención. Luego, proponemos el mecanismo de retención para el modelado de secuencias, que admite tres paradigmas de cálculo: paralelo, recurrente y recurrente por fragmentos. Específicamente, la representación paralela permite el paralelismo en el entrenamiento. La representación recurrente habilita una inferencia de bajo costo O(1), lo que mejora el rendimiento en términos de rendimiento de decodificación, latencia y memoria de GPU sin sacrificar el rendimiento. La representación recurrente por fragmentos facilita el modelado eficiente de secuencias largas con complejidad lineal, donde cada fragmento se codifica en paralelo mientras se resume recurrentemente los fragmentos. Los resultados experimentales en modelado de lenguaje muestran que RetNet logra resultados favorables en escalabilidad, entrenamiento paralelo, despliegue de bajo costo e inferencia eficiente. Estas propiedades intrigantes convierten a RetNet en un fuerte sucesor del Transformer para modelos de lenguaje a gran escala. El código estará disponible en https://aka.ms/retnet.

English

In this work, we propose Retentive Network (RetNet) as a foundation architecture for large language models, simultaneously achieving training parallelism, low-cost inference, and good performance. We theoretically derive the connection between recurrence and attention. Then we propose the retention mechanism for sequence modeling, which supports three computation paradigms, i.e., parallel, recurrent, and chunkwise recurrent. Specifically, the parallel representation allows for training parallelism. The recurrent representation enables low-cost O(1) inference, which improves decoding throughput, latency, and GPU memory without sacrificing performance. The chunkwise recurrent representation facilitates efficient long-sequence modeling with linear complexity, where each chunk is encoded parallelly while recurrently summarizing the chunks. Experimental results on language modeling show that RetNet achieves favorable scaling results, parallel training, low-cost deployment, and efficient inference. The intriguing properties make RetNet a strong successor to Transformer for large language models. Code will be available at https://aka.ms/retnet.