Moyen clé-valeur

Résumé

Nous présentons Key-Value Means (« KVM »), une nouvelle récurrence par blocs pour l'attention qui peut prendre en charge soit un état de taille fixe, soit un état croissant. Équiper un transformer de base robuste de couches d'attention KVM de taille fixe produit un RNN segmenté fort en O(N), tout en n'ajoutant qu'un nombre insignifiant de nouveaux paramètres. Nous entraînons un transformer avec un cache KVM extensible et montrons qu'il obtient des performances compétitives sur des tests de long contexte avec seulement un temps de pré-remplissage sous-quadratique et une croissance d'état sous-linéaire. KVM est implémentable avec des opérations standard et sans noyaux personnalisés, et supporte un entraînement et un pré-remplissage parallélisables par segments. Il offre de nombreux avantages à la fois des transformers traditionnels (mémoire de contexte extensible, entraînement et pré-remplissage parallélisables par segments) et des RNN linéaires dans un seul package unifié. Il peut être utilisé sur chaque couche, économisant la mémoire du cache KV, et permettant une gamme continue de choix de complexité temporelle de pré-remplissage entre O(N) et O(N^2). Il peut également être implémenté dans une solution hybride en tandem avec des couches LRNN en lieu et place de l'attention traditionnelle, pour compléter le LRNN avec une utilisation améliorée de la longueur de contexte à croissance mémoire sous-linéaire et un décodage de long contexte. Nous publions notre code à l'adresse https://github.com/recursal/KVM-paper et les modèles entraînés à l'adresse https://huggingface.co/collections/recursal/key-value-means sous la licence Apache 2.0.

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.