Mémorisation de contexte pour une génération efficace de longs contextes

Résumé

Les applications modernes des grands modèles de langage (LLM) reposent de plus en plus sur de longs préfixes de conditionnement pour contrôler le comportement du modèle au moment de l'inférence. Bien que l'inférence augmentée par préfixe soit efficace, elle souffre de deux limitations structurelles : i) l'influence du préfixe s'estompe à mesure que la génération progresse, et ii) le calcul d'attention sur le préfixe évolue linéairement avec sa longueur. Les approches existantes conservent soit le préfixe dans l'attention tout en le compressant, soit l'internalisent dans les paramètres du modèle via un entraînement basé sur les gradients. La première méthode continue de traiter le préfixe lors de l'inférence, tandis que la seconde est coûteuse en entraînement et mal adaptée aux mises à jour du préfixe. Pour résoudre ces problèmes, nous proposons la mémoire d'état d'attention, une approche sans entraînement qui externalise le préfixe dans une mémoire légère basée sur une recherche d'états d'attention précalculés entre les jetons du préfixe et ceux de la requête. Sur ManyICLBench avec LLaMA-3.1-8B, notre méthode améliore la précision par rapport à l'apprentissage en contexte pour des budgets mémoire de 1K à 8K, tout en réduisant la latence d'attention d'un facteur 1,36 à 8K, et dépasse les performances RAG à attention complète sur le benchmark NBA en utilisant seulement 20 % de son empreinte mémoire.

English

Modern large language model (LLM) applications increasingly rely on long conditioning prefixes to control model behavior at inference time. While prefix-augmented inference is effective, it incurs two structural limitations: i) the prefix's influence fades as generation proceeds, and ii) attention computation over the prefix scales linearly with its length. Existing approaches either keep the prefix in attention while compressing it, or internalize it into model parameters through gradient-based training. The former still attends to the prefix at inference, while the latter is training-intensive and ill-suited to prefix updates. To address these issues, we propose attention-state memory, a training-free approach that externalizes the prefix into a lightweight, lookup-based memory of precomputed attention states between prefix and query tokens. On ManyICLBench with LLaMA-3.1-8B, our method improves accuracy over in-context learning at 1K-8K memory budgets while reducing attention latency by 1.36x at 8K, and surpasses full-attention RAG performance on NBA benchmark using only 20% of its memory footprint.