GoldFinch: Alto Rendimiento Híbrido RWKV/Transformador con Precarga Lineal y Compresión Extrema de Caché KV.

GoldFinch: High Performance RWKV/Transformer Hybrid with Linear Pre-Fill and Extreme KV-Cache Compression

July 16, 2024

Autores: Daniel Goldstein, Fares Obeid, Eric Alcaide, Guangyu Song, Eugene Cheah

cs.AI

Resumen

Presentamos GoldFinch, un modelo de secuencia híbrido de Atención Lineal/Transformer que utiliza una nueva técnica para generar de manera eficiente un KV-Cache altamente comprimido y reutilizable en tiempo y espacio lineal con respecto a la longitud de la secuencia. GoldFinch apila nuestro nuevo transformador GOLD sobre una versión mejorada de la arquitectura Finch (RWKV-6). Entrenamos modelos de clase de hasta 1.5 mil millones de parámetros de las arquitecturas Finch, Llama y GoldFinch, y encontramos una mejora drástica en el rendimiento de modelado en comparación con Finch y Llama. Nuestros ahorros en el tamaño de la caché aumentan linealmente con el número de capas del modelo, siendo de 756-2550 veces más pequeños que la caché tradicional del transformer para tamaños comunes, lo que permite inferir longitudes de contexto extremadamente grandes incluso en hardware limitado. Aunque la generación autoregresiva tiene una complejidad temporal de O(n) por token debido a la atención, el cálculo de precarga del estado inicial completo de la caché para un contexto enviado cuesta solo O(1) tiempo por token debido al uso de una red neuronal recurrente (RNN) para generar esta caché. Publicamos nuestros pesos entrenados y código de entrenamiento bajo la licencia Apache 2.0 para uso comunitario.

English

We introduce GoldFinch, a hybrid Linear Attention/Transformer sequence model that uses a new technique to efficiently generate a highly compressed and reusable KV-Cache in linear time and space with respect to sequence length. GoldFinch stacks our new GOLD transformer on top of an enhanced version of the Finch (RWKV-6) architecture. We train up to 1.5B parameter class models of the Finch, Llama, and GoldFinch architectures, and find dramatically improved modeling performance relative to both Finch and Llama. Our cache size savings increase linearly with model layer count, ranging from 756-2550 times smaller than the traditional transformer cache for common sizes, enabling inference of extremely large context lengths even on limited hardware. Although autoregressive generation has O(n) time complexity per token because of attention, pre-fill computation of the entire initial cache state for a submitted context costs only O(1) time per token due to the use of a recurrent neural network (RNN) to generate this cache. We release our trained weights and training code under the Apache 2.0 license for community use.