GoldFinch: Ibrido RWKV/Transformer ad alte prestazioni con Pre-Fill lineare e compressione estrema della KV-Cache

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

July 16, 2024

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

cs.AI

Abstract

Presentiamo GoldFinch, un modello ibrido di sequenza Linear Attention/Transformer che utilizza una nuova tecnica per generare in modo efficiente una KV-Cache altamente compressa e riutilizzabile in tempo e spazio lineari rispetto alla lunghezza della sequenza. GoldFinch combina il nostro nuovo trasformatore GOLD con una versione potenziata dell'architettura Finch (RWKV-6). Addestriamo modelli di classe fino a 1,5 miliardi di parametri per le architetture Finch, Llama e GoldFinch, riscontrando un miglioramento significativo delle prestazioni di modellazione rispetto sia a Finch che a Llama. Il risparmio nella dimensione della cache aumenta linearmente con il numero di strati del modello, risultando da 756 a 2550 volte più piccolo rispetto alla cache tradizionale dei trasformatori per dimensioni comuni, consentendo l'inferenza di contesti estremamente lunghi anche su hardware limitato. Sebbene la generazione autoregressiva abbia una complessità temporale O(n) per token a causa dell'attenzione, il calcolo preliminare dell'intero stato iniziale della cache per un contesto inviato costa solo O(1) tempo per token grazie all'uso di una rete neurale ricorrente (RNN) per generare questa cache. Rilasciamo i nostri pesi addestrati e il codice di addestramento sotto licenza Apache 2.0 per l'uso della comunità.

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.