GoldFinch: Высокопроизводительный гибрид RWKV/Transformer с линейной предварительной заполненностью и экстремальным сжатием кэша KV.

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

July 16, 2024

Авторы: Daniel Goldstein, Fares Obeid, Eric Alcaide, Guangyu Song, Eugene Cheah

cs.AI

Аннотация

Мы представляем GoldFinch, гибридную модель последовательности с использованием линейного внимания/трансформера, которая использует новую технику для эффективной генерации высококомпрессионного и многократно используемого KV-кэша за линейное время и пространство относительно длины последовательности. GoldFinch ставит наш новый трансформер GOLD поверх улучшенной версии архитектуры Finch (RWKV-6). Мы обучаем до 1.5 млрд классовых моделей параметров архитектур Finch, Llama и GoldFinch и обнаруживаем драматически улучшенную производительность модели по сравнению с Finch и Llama. Наши экономии размера кэша увеличиваются линейно с количеством слоев модели, варьируя от 756 до 2550 раз меньше, чем у традиционного кэша трансформера для обычных размеров, обеспечивая вывод крайне больших длин контекста даже на ограниченном оборудовании. Хотя авторегрессионная генерация имеет временную сложность O(n) на токен из-за внимания, предварительное вычисление всего начального состояния кэша для представленного контекста стоит только O(1) времени на токен из-за использования рекуррентной нейронной сети (RNN) для генерации этого кэша. Мы выпускаем наши обученные веса и код обучения под лицензией Apache 2.0 для использования сообществом.

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.