XC-Cache : Utilisation de l'attention croisée pour accéder au contexte mis en cache afin d'améliorer l'efficacité de l'inférence des grands modèles de langage

XC-Cache: Cross-Attending to Cached Context for Efficient LLM Inference

April 23, 2024

Auteurs: João Monteiro, Étienne Marcotte, Pierre-André Noël, Valentina Zantedeschi, David Vázquez, Nicolas Chapados, Christopher Pal, Perouz Taslakian

cs.AI

Résumé

Les approches d'apprentissage en contexte (ICL) exploitent généralement des techniques de prompting pour conditionner la génération des modèles de langage à décodeur unique sur des informations de référence. Le traitement en temps réel d'un contexte est inefficace en raison du coût quadratique des opérations d'auto-attention, ce qui rend la mise en cache souhaitable. Cependant, la mise en cache des états des transformateurs peut facilement nécessiter presque autant d'espace que les paramètres du modèle. Lorsque le contexte approprié n'est pas connu à l'avance, la mise en cache de l'ICL peut s'avérer difficile. Ce travail aborde ces limitations en introduisant des modèles qui, inspirés par l'architecture encodeur-décodeur, utilisent l'attention croisée pour conditionner la génération sur un texte de référence sans recourir au prompting. Plus précisément, nous exploitons des modèles pré-entraînés à décodeur unique et n'entraînons qu'un petit nombre de couches supplémentaires. Nous utilisons la tâche de Question-Réponse (QA) comme banc d'essai pour évaluer la capacité de nos modèles à effectuer une génération conditionnelle et observons qu'ils surpassent l'ICL, sont comparables aux grands modèles de langage (LLM) avec prompting finement ajustés, et réduisent considérablement l'empreinte mémoire par rapport à la mise en cache standard des clés-valeurs (KV) d'un facteur de deux ordres de grandeur.

English

In-context learning (ICL) approaches typically leverage prompting to condition decoder-only language model generation on reference information. Just-in-time processing of a context is inefficient due to the quadratic cost of self-attention operations, and caching is desirable. However, caching transformer states can easily require almost as much space as the model parameters. When the right context isn't known in advance, caching ICL can be challenging. This work addresses these limitations by introducing models that, inspired by the encoder-decoder architecture, use cross-attention to condition generation on reference text without the prompt. More precisely, we leverage pre-trained decoder-only models and only train a small number of added layers. We use Question-Answering (QA) as a testbed to evaluate the ability of our models to perform conditional generation and observe that they outperform ICL, are comparable to fine-tuned prompted LLMs, and drastically reduce the space footprint relative to standard KV caching by two orders of magnitude.