Résumé technique : Récupération augmentée par mémoire classée améliorée (ERMAR)

1. Énoncé du problème

Les modèles de langage de grande taille (LLM) rencontrent des limitations fondamentales dans le traitement de contextes étendus en raison de la complexité quadratique des mécanismes d'attention et de la demande croissante en mémoire durant la génération. Bien que des approches existantes comme MemLong (Liu et al., 2024) tentent de résoudre ce problème en stockant un contexte historique précalculé dans une banque de mémoire et en récupérant des segments pertinents via une similarité de produit scalaire, elles souffrent d'inefficacités critiques :

• Traitement uniforme : MemLong traite tous les paires clé-valeur (K-V) avec un poids égal, indépendamment de leur pertinence contextuelle, entraînant une surcharge d'informations.
• Précision réduite : L'absence d'un mécanisme de classement nuancé se traduit par une précision de récupération plus faible, en particulier dans les scénarios nécessitant un dialogue soutenu ou une analyse de documents.
• Gestion statique : Les méthodes actuelles reposent souvent sur des structures de mémoire fixes qui ne s'adaptent ni à l'importance du contenu ni aux schémas d'utilisation.

2. Méthodologie : Le cadre ERMAR

Les auteurs proposent la Récupération augmentée par mémoire classée améliorée (ERMAR), un cadre conçu pour classer dynamiquement les entrées de mémoire en fonction de leur pertinence. L'architecture s'appuie sur la base de référence MemLong mais introduit un nouveau mécanisme de score de pertinence et un modèle de re-classement point par point inspiré des techniques d'apprentissage pour le classement dans la récupération d'information.

Architecture de base

Le système ERMAR fonctionne en quatre étapes séquentielles :

1. Récupération de mémoire longue : L'historique d'entrée est encodé en paires clé-valeur (K-V) au niveau des segments et stocké dans une banque de mémoire aux côtés d'embeddings denses capturant les représentations sémantiques.
2. Recherche : Lors de la réception d'une nouvelle requête, son embedding est utilisé pour rechercher dans la banque de mémoire des paires K-V candidates via une correspondance de similarité.
3. Re-classement (Nouveauté centrale) : Les candidats récupérés passent par un module de re-classement. Contrairement à MemLong, ERMAR emploie un mécanisme de score de pertinence basé sur la multiplication pour calculer les scores des paires requête-document. Ce module réorganise dynamiquement les embeddings K-V en fonction de leurs scores de pertinence, en priorisant les informations les plus pertinentes.
4. Fusion de mémoire et génération : Seules les paires K-V les mieux classées sont injectées dans le bloc supérieur entraînables du transformateur via l'Attention causale de récupération, conditionnant la sortie finale sur le contexte historique le plus pertinent.

Composants techniques clés

• Score de pertinence : Le score de pertinence $\alpha(q, K)$ est défini à l'aide d'une fonction softmax sur le produit scalaire de l'embedding de requête $q$ et des embeddings de clés $K$, normalisé par $\sqrt{d_{ret}}$. Cela imite les opérations d'attention pour attribuer une importance relative aux entrées de mémoire.
• Score de pertinence avec récupération adaptative (RSAR) : Ce mécanisme classe dynamiquement les entrées de mémoire $(K_j, V_j, s_j)$ où $s_j$ est le score de pertinence. Il emploie une stratégie d'élagage pour rejeter les entrées en dessous d'un seuil prédéfini, optimisant l'utilisation de la mémoire en ne conservant que les informations contextuelles critiques.
• Gestion hybride de la mémoire : ERMAR utilise une stratégie d'élagage dynamique qui conserve les 10 % les plus récents de la mémoire, priorise les 80 % du milieu en fonction de la fréquence de récupération, et rejette les 10 % les plus anciens comme potentiellement obsolètes.
• Stratégie d'entraînement : Le modèle utilise un bloc inférieur de transformateur figé (produisant des paires K-V) et un bloc supérieur entraînables avec une attention augmentée par récupération. Il est affiné à l'aide de LoRA sur le jeu de données SlimPajama.

3. Contributions principales

L'article identifie trois améliorations principales introduites par ERMAR :

1. Métrique de similarité sémantique : Une nouvelle métrique pour mesurer l'alignement contextuel entre les embeddings de requête et les paires K-V, dépassant les simples recherches de similarité.
2. Fonction de score pondérée : Une fonction qui prend en compte à la fois la similarité de contenu et la pertinence contextuelle pour calculer les scores de pertinence.
3. Intégration de l'historique d'utilisation : L'intégration des schémas d'utilisation historiques pour affiner l'évaluation de la pertinence et la priorisation de la mémoire.

4. Résultats expérimentaux

ERMAR a été évalué par rapport à des modèles de l'état de l'art (incluant MemLong, LongLoRA, YARN et Phi-3) sur des échelles de 3 milliards et 7 milliards de paramètres, sur des benchmarks tels que WikiText-103, PG-19 et Proof-Pile.

• Modélisation du langage en contexte long :
  • ERMAR a obtenu des résultats de l'état de l'art parmi les modèles de 3 milliards de paramètres, surpassant systématiquement MemLong dans 10 configurations sur 12 (83,3 %).
  • Il a démontré une stabilité supérieure lors de la mise à l'échelle vers des contextes plus longs. Sur PG-19, la perplexité n'a augmenté que de 0,31 % lors du passage de 4k à 16k tokens, comparé à une dégradation significative des modèles de référence.
  • À 32k tokens, ERMAR a maintenu des performances compétitives, montrant une amélioration de 0,90 % de la perplexité par rapport à MemLong sur PG-19.
• Apprentissage en contexte (ICL) :
  • Dans des configurations à 4 coups et 20 coups sur cinq tâches de compréhension du langage naturel (SST-2, MR, Subj, SST-5, MPQA), ERMAR a surpassé les références.
  • Des gains notables ont été observés sur MPQA (+21 % par rapport à MemLong en 20 coups) et Subj (+26 %).
• Efficacité mémoire :
  • ERMAR a réduit la mémoire réservée jusqu'à 30 % par rapport à MemLong pour des contextes de 16k tokens (16,61 Go contre 23,77 Go).
  • Il a obtenu une amélioration de huit fois de l'efficacité mémoire par token lors de la mise à l'échelle de 1k à 16k tokens.
• Latence et débit :
  • Bien qu'ERMAR ait présenté une latence légèrement plus élevée (augmentation relative de 9 à 35 % sur des contextes plus courts), il a démontré une variance de latence significativement plus faible (±45–59 ms contre ±154–214 ms pour MemLong), indiquant une plus grande stabilité d'exécution.
  • Le débit est resté compétitif à environ 2,4k tokens/seconde pour des contextes plus longs.

5. Importance et affirmations

Les auteurs affirment qu'ERMAR fournit une approche plus rigoureuse de la gestion de mémoire en contexte long en s'attaquant aux limites du traitement statique et uniforme de la mémoire.

• Évolutivité : Le cadre atténue efficacement la dilution de l'information dans les séquences étendues, permettant aux modèles de maintenir une cohérence sur des contextes de 32k tokens.
• Précision de récupération : En intégrant un paradigme à deux étapes de récupération puis de re-classement, ERMAR garantit que les informations historiques les plus sémantiquement pertinentes sont priorisées, conduisant à des réponses plus précises et cohérentes dans les scénarios de dialogue soutenu.
• Efficacité : Les stratégies de gestion et d'élagage adaptatifs de la mémoire permettent des réductions significatives de la surcharge mémoire sans sacrifier les performances, en faisant une solution viable pour les applications en contexte long à ressources limitées.

L'article conclut que, bien qu'ERMAR introduise une surcharge de calcul par rapport aux LLM standards, le compromis est justifié par des gains substantiels en précision de récupération, en rétention de contexte et en efficacité mémoire. Un travail futur est proposé pour optimiser le cadre pour des jeux de données spécialisés et des tâches de raisonnement complexe.