Sleutel-waarde gemiddelden

Samenvatting

We presenteren Key-Value Means ("KVM"), een nieuwe blok-recurrentie voor aandacht die zowel een vaste grootte als een groeiende toestand kan accommoderen. Het uitrusten van een sterke transformer-baseline met vaste-grootte KVM-aandachtslagen levert een sterke O(N) chunked RNN op, terwijl slechts een verwaarloosbaar aantal nieuwe parameters wordt toegevoegd. We trainen een transformer met een uitbreidbare KVM-cache en tonen aan dat deze concurrerend presteert op lange-contexttesten met slechts subkwadratische voorvultijd en sublineaire toestandsgroei. KVM is implementeerbaar met standaardbewerkingen en zonder aangepaste kernels, en ondersteunt chunkgewijs paralleliseerbare training en voorvulling. Het biedt veel van de voordelen van zowel traditionele transformers (uitbreidbaar contextgeheugen, chunkgewijs paralleliseerbare training en voorvulling) als lineaire RNN's in één uniform pakket. Het kan op elke laag worden gebruikt, waardoor KV-cachegeheugen wordt bespaard en een continu bereik aan keuzes voor de voorvultijdcomplexiteit tussen O(N) en O(N^2) mogelijk wordt. Het kan ook worden geïmplementeerd in een hybride oplossing samen met LRNN-lagen in plaats van traditionele aandacht, om de LRNN aan te vullen met verbeterd sublineair geheugengroeicontextlengtegebruik en lange-contextdecodering. We publiceren onze code op https://github.com/recursal/KVM-paper en getrainde modellen op https://huggingface.co/collections/recursal/key-value-means onder de Apache 2.0-licentie.

English

We present Key-Value Means ("KVM"), a novel block-recurrence for attention that can accommodate either fixed-size or growing state. Equipping a strong transformer baseline with fixed-size KVM attention layers yields a strong O(N) chunked RNN, while adding only an insignificant number of new parameters. We train a transformer with a growable KVM cache and show it performs competitively on long-context tests with only subquadratic prefill time and sublinear state growth. KVM is implementable with standard operations and without custom kernels, and supports chunk-wise parallelizable training and prefill. It provides many of the benefits of both traditional transformers (expandable context memory, chunk-wise parallelizable training and prefill) and linear RNNs in a single unified package. It can be used on every layer, saving KV-cache memory, and allowing a continuous range of choices of prefill time complexity between O(N) and O(N^2). It can also be implemented in a hybrid solution in tandem with LRNN layers in place of traditional attention, to supplement the LRNN with improved sublinear memory growth context length usage and long context decoding. We release our code at https://github.com/recursal/KVM-paper and trained models at https://huggingface.co/collections/recursal/key-value-means under the Apache 2.0 license.