GoldFinch : Hybride RWKV/Transformer Haute Performance avec Pré-Remplissage Linéaire et Compression Extrême du Cache KV

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

July 16, 2024

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

cs.AI

Résumé

Nous présentons GoldFinch, un modèle de séquence hybride Attention Linéaire/Transformer qui utilise une nouvelle technique pour générer efficacement un KV-Cache hautement compressé et réutilisable en temps et espace linéaires par rapport à la longueur de la séquence. GoldFinch superpose notre nouveau transformateur GOLD au-dessus d'une version améliorée de l'architecture Finch (RWKV-6). Nous entraînons des modèles de la classe des 1,5 milliards de paramètres pour les architectures Finch, Llama et GoldFinch, et observons une amélioration spectaculaire des performances de modélisation par rapport à Finch et Llama. Nos économies de taille de cache augmentent linéairement avec le nombre de couches du modèle, allant de 756 à 2550 fois plus petit que le cache traditionnel des transformateurs pour les tailles courantes, permettant ainsi l'inférence de contextes extrêmement longs même sur du matériel limité. Bien que la génération autoregressive ait une complexité temporelle de O(n) par token en raison de l'attention, le calcul de pré-remplissage de l'état initial complet du cache pour un contexte soumis ne coûte que O(1) temps par token grâce à l'utilisation d'un réseau de neurones récurrent (RNN) pour générer ce cache. Nous publions nos poids entraînés et notre code d'entraînement sous licence Apache 2.0 pour une utilisation communautaire.

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.