GoldFinch: Hochleistungs-RWKV/Transformer-Hybrid mit linearem Vorabfüllen und extrem hoher KV-Cache-Komprimierung

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

July 16, 2024

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

cs.AI

Zusammenfassung

Wir stellen GoldFinch vor, ein hybrides Linear Attention/Transformer-Sequenzmodell, das eine neue Technik verwendet, um effizient einen hochkomprimierten und wiederverwendbaren KV-Cache in linearer Zeit und Raum in Bezug auf die Sequenzlänge zu generieren. GoldFinch stapelt unseren neuen GOLD-Transformer auf eine verbesserte Version der Finch (RWKV-6)-Architektur. Wir trainieren bis zu 1,5 Milliarden Parameterklassenmodelle der Finch-, Llama- und GoldFinch-Architekturen und stellen eine dramatisch verbesserte Modellierungsleistung im Vergleich zu Finch und Llama fest. Unsere Cache-Größeneinsparungen steigen linear mit der Anzahl der Modellschichten und sind für gängige Größen zwischen 756 und 2550 Mal kleiner als der traditionelle Transformer-Cache, was die Inferenz extrem großer Kontextlängen auch auf begrenzter Hardware ermöglicht. Obwohl die autoregressive Generierung aufgrund der Aufmerksamkeit eine Zeitkomplexität von O(n) pro Token hat, kostet die Vorab-Berechnung des gesamten anfänglichen Cache-Zustands für einen übermittelten Kontext nur O(1) Zeit pro Token aufgrund der Verwendung eines rekurrenten neuronalen Netzwerks (RNN) zur Generierung dieses Caches. Wir veröffentlichen unsere trainierten Gewichte und Trainingscode unter der Apache 2.0-Lizenz zur gemeinschaftlichen Nutzung.

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.