Context Memorization for Efficient Long Context Generation

Résumé Technique : Mémoire d'État d'Attention pour une Génération de Contexte Long Efficace

Énoncé du Problème
Les applications modernes de Modèles de Langage à Grande Échelle (LLM) reposent de plus en plus sur des préfixes de conditionnement longs (par exemple, des exemples en contexte, des documents récupérés ou des instructions d'agents) pour guider l'inférence. Cependant, l'inférence augmentée par préfixe fait face à deux limitations structurelles principales :

1. Déclin du Préfixe : À mesure que la génération progresse, l'attention du modèle se distribue sur les tokens nouvellement générés, ce qui fait que l'influence du préfixe s'estompe, en particulier dans les scénarios de contexte long.
2. Inefficacité de l'Inférence : Le calcul de l'attention sur le préfixe évolue linéairement avec sa longueur ($O(L)$), imposant une latence et une surcharge mémoire significatives lors de la phase de préremplissage et à chaque étape de décodage. Bien que la mise en cache du préfixe amortisse les coûts de préremplissage, elle entraîne toujours une consommation mémoire substantielle.

Les solutions existantes se divisent en deux catégories, toutes deux présentant des inconvénients :

• Internalisation du Préfixe : Méthodes qui ajustent les paramètres du modèle ou utilisent des hyperréseaux pour internaliser le comportement du préfixe. Elles nécessitent un entraînement basé sur le gradient, ce qui les rend gourmandes en ressources, lentes à mettre à jour et inadaptées aux changements dynamiques de préfixe.
• Compression du Préfixe : Méthodes qui maintiennent le préfixe dans le mécanisme d'attention mais le compressent (par exemple, élagage de tokens ou éviction du cache KV). Elles nécessitent toujours de porter attention au préfixe compressé lors de l'inférence, ce qui signifie que la latence évolue avec la longueur compressée et que l'influence du préfixe reste sujette au déclin.

Méthodologie : Mémoire d'État d'Attention (ASM)
Les auteurs proposent la Mémoire d'État d'Attention, une approche sans entraînement qui externalise le préfixe vers une mémoire légère basée sur la recherche d'états d'attention précalculés. Cette méthode exploite l'identité softmax en ligne pour décomposer l'attention sur un bloc de clés concaténé en une attention sur des sous-blocs.

L'idée centrale est que, pour une requête donnée, l'état d'attention $(a, Z)$ — composé de la sortie d'attention et du terme de normalisation — suffit à reconstruire la contribution du préfixe à l'attention sans se ré-attacher aux clés et valeurs originales.

Le cadre fonctionne en deux phases :

1. Phase de Calibration Hors Ligne (Construction de la Mémoire) :

   • Collecte : Le modèle effectue des passages avant sur un ensemble de traces de requête-réponse représentatives concaténées avec le préfixe. Pour chaque token de requête dans la réponse, le système enregistre le vecteur de requête et l'état d'attention correspondant $(a, Z)$ sur le préfixe.
   • Regroupement : Les vecteurs de requête collectés sont regroupés (par exemple, via K-moyennes) en $K$ centroïdes.
   • Agrégation : Pour chaque groupe, une entrée représentative unique est créée. Le vecteur de requête du centroïde est la moyenne des requêtes du groupe. L'état d'attention $(\bar{a}, \bar{Z})$ est calculé via une agrégation consciente de l'attention, fusionnant les états de tous les membres du groupe en utilisant la règle de mise à jour softmax en ligne. Cela garantit que le centroïde représente une moyenne sans perte de l'attention du préfixe du groupe.
   • Efficacité : Pour gérer des préfixes très longs (par exemple, 16K+ tokens) sans dépasser la mémoire GPU, le préfixe peut être partitionné en blocs, encodés indépendamment, puis fusionnés dans la mémoire en utilisant la propriété de compositionnalité de l'identité softmax en ligne.

2. Phase d'Inférence En Ligne :

   • Récupération : Pour une requête entrante, le système identifie le centroïde de mémoire le plus proche en utilisant la similarité cosinus (optionnellement avec des transformations de blanchiment et la gestion de RoPE).
   • Fusion : L'état d'attention précalculé récupéré $(\bar{a}, \bar{Z})$ est fusionné avec la sortie d'auto-attention de la requête (calculée sur les tokens non préfixes) en utilisant l'opérateur de fusion défini par l'identité softmax en ligne. Cela permet de retrouver la sortie d'attention exacte qui aurait été produite si le préfixe avait été présent, sans porter attention aux tokens du préfixe.
   • Complexité : Le coût de la recherche évolue logarithmiquement ($O(\log K)$) avec la taille de la mémoire via un index hiérarchique, indépendamment de la longueur originale du préfixe.

Contributions Clés

1. Externalisation Sans Entraînement : L'article propose une méthode pour externaliser les longs préfixes vers un dictionnaire de mémoire compact sans nécessiter d'entraînement basé sur le gradient ni de modèles auxiliaires.
2. Extension du Softmax en Ligne : Les auteurs étendent l'identité softmax en ligne du calcul efficace de l'attention au domaine de la réutilisation de préfixes entre requêtes, permettant une reconstruction sans perte de l'attention du préfixe via la récupération.
3. Performance et Efficacité : Les expériences démontrent que l'ASM égale ou dépasse les performances de l'attention complète tout en réduisant considérablement les coûts de calcul.

Résultats Expérimentaux
La méthode a été évaluée sur ManyICLBench (Apprentissage en Contexte) et RuleArena/NBA (Génération Augmentée par Récupération) en utilisant LLaMA-3.1-8B.

• Apprentissage en Contexte (ICL) : Sur ManyICLBench, l'ASM a amélioré la précision par rapport à l'apprentissage en contexte standard sur des budgets mémoire de 1K à 8K entrées. Avec un budget mémoire de 8K, elle a réduit la latence d'attention de 1,36× par rapport à l'attention complète. L'ASM a surpassé les bases de compression de cache KV (KVzip), qui ont souvent fait s'effondrer des informations critiques comme les tokens d'étiquette.
• Génération Augmentée par Récupération (RAG) : Sur le benchmark NBA (impliquant environ 20K tokens de règles), l'ASM a surpassé les performances de RAG à attention complète en utilisant seulement 20 % de l'empreinte mémoire. Les performances optimales ont été observées à une taille de codebook intermédiaire (4K entrées), suggérant que le nombre d'entrées est un hyperparamètre spécifique à la tâche plutôt qu'un paramètre qui devrait simplement être maximisé.
• Latence : Bien que la recherche hiérarchique entraîne une surcharge plus élevée avec des budgets mémoire très faibles, la croissance de la latence est sous-linéaire. L'ASM devient plus rapide que l'attention complète à environ 4K entrées mémoire et atteint une accélération de 1,8× à 16K entrées.

Signification et Revendications
L'article revendique que la Mémoire d'État d'Attention adresse les coûts structurels de la génération de contexte long en découplant la connaissance du préfixe du calcul d'attention par requête. En externalisant le préfixe vers une mémoire précalculée basée sur la recherche, la méthode élimine l'échelle linéaire de la latence d'inférence avec la longueur du préfixe et supprime le problème de déclin du préfixe associé aux mécanismes d'attention standard.

Les auteurs positionnent ce travail comme une étape vers une direction de recherche plus large : l'externalisation des connaissances des LLM vers des représentations compactes et réutilisables qui existent au-delà des invites textuelles et des paramètres du modèle. Ils soulignent que leur approche offre un compromis favorable entre précision et efficacité, fournissant une alternative sans entraînement à l'internalisation des paramètres et une alternative plus efficace à la compression de préfixe.