GoldFinch: Híbrido de Alto Desempenho RWKV/Transformer com Pré-Preenchimento Linear e Compressão Extrema de Cache 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

Resumo

Apresentamos o GoldFinch, um modelo de sequência híbrido Linear Attention/Transformer que utiliza uma nova técnica para gerar de forma eficiente um KV-Cache altamente comprimido e reutilizável em tempo e espaço lineares em relação ao comprimento da sequência. O GoldFinch empilha o nosso novo transformador GOLD em cima de uma versão aprimorada da arquitetura Finch (RWKV-6). Treinamos modelos de classe de até 1.5B de parâmetros das arquiteturas Finch, Llama e GoldFinch, e observamos uma melhoria significativa no desempenho de modelagem em comparação com Finch e Llama. Nossas economias de tamanho de cache aumentam linearmente com a contagem de camadas do modelo, variando de 756 a 2550 vezes menores do que o cache tradicional do transformer para tamanhos comuns, possibilitando a inferência de comprimentos de contexto extremamente grandes mesmo em hardware limitado. Embora a geração autoregressiva tenha complexidade temporal O(n) por token devido à atenção, o cálculo de pré-preenchimento do estado inicial completo do cache para um contexto submetido custa apenas O(1) por token devido ao uso de uma rede neural recorrente (RNN) para gerar esse cache. Disponibilizamos nossos pesos treinados e código de treinamento sob a licença Apache 2.0 para uso comunitário.

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.