GoldFinch: Hoogwaardige RWKV/Transformer Hybride met Lineaire Pre-Fill en Extreme KV-Cache Compressie

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

Samenvatting

We introduceren GoldFinch, een hybride Lineaire Attention/Transformer-sequentiemodel dat een nieuwe techniek gebruikt om efficiënt een sterk gecomprimeerde en herbruikbare KV-Cache te genereren in lineaire tijd en ruimte ten opzichte van de sequentielengte. GoldFinch plaatst onze nieuwe GOLD-transformer bovenop een verbeterde versie van de Finch (RWKV-6)-architectuur. We trainen modellen van de Finch-, Llama- en GoldFinch-architecturen tot 1,5B parameters en vinden een aanzienlijk verbeterde modelprestatie ten opzichte van zowel Finch als Llama. Onze besparingen in cachegrootte nemen lineair toe met het aantal modellagen, variërend van 756 tot 2550 keer kleiner dan de traditionele transformercache voor gangbare groottes, wat inferentie van extreem grote contextlengtes mogelijk maakt, zelfs op beperkte hardware. Hoewel autoregressieve generatie een O(n)-tijdscomplexiteit per token heeft vanwege attention, kost de voorafgaande berekening van de volledige initiële cachestatus voor een ingediende context slechts O(1) tijd per token dankzij het gebruik van een recurrent neural network (RNN) om deze cache te genereren. We geven onze getrainde gewichten en trainingscode vrij onder de Apache 2.0-licentie voor gebruik door de gemeenschap.

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.