Knowledge Graph Enhanced Memory-Augmented Retrieval for Long Context Modeling

Résumé Technique : KGERMAR

Énoncé du Problème
La modélisation de langage à contexte long fait face à un défi fondamental qui dépasse la simple extension des fenêtres de contexte : maintenir une compréhension cohérente des états et des relations entre les entités à travers des milliers de jetons. Les approches standards souffrent de deux limitations primaires. Premièrement, les modèles à attention étendue (ex. : LongLoRA, YARN) souffrent d'une dégradation de type « perdu au milieu » (lost-in-the-middle), où l'information située au milieu d'un contexte long est systématiquement sous-attendue. Deuxièmement, les systèmes de récupération augmentés par la mémoire (ex. : MemLong, ERMAR) reposent exclusivement sur la similitude sémantique. Bien qu'efficaces pour récupérer des passages contenant des termes correspondants, ils ne parviennent pas à distinguer les passages qui mentionnent simplement des termes liés de manière fortuite de ceux qui décrivent des relations causales ou structurelles explicites entre les mêmes entités. Cette limitation est critique dans les scénarios exigeant un raisonnement structurel, comme le suivi de la progression de l'état d'erreur d'un appareil au cours d'une longue session de dépannage, où le contexte pertinent peut ne partager aucune superposition lexicale avec la requête, mais partage des connexions structurelles profondes.

Méthodologie
Le document propose KGERMAR (Knowledge Graph Enhanced Memory-Augmented Retrieval), un cadre qui construit des graphes de connaissances dynamiques et spécifiques au contexte pendant l'inférence pour intégrer les plongements structurels de graphes à la sémantique textuelle. Le système fonctionne via quatre composantes intégrées :

1. Construction de Graphe de Connaissances Contextuel :

   • Extraction : Le système effectue une reconnaissance d'entités nommées (NER) en temps réel à l'aide d'un modèle BERT-large (affiné sur CoNLL-2003) et une extraction de relations à l'aide d'un modèle BERT-base (affiné sur TACRED).
   • Construction du Graphe : Il extrait les entités et les relations (ex. : CAUSEDBY, PARTOF, USEDFOR) au sein d'une fenêtre glissante. Pour garantir la qualité, il emploie la consolidation d'entités (fusion des variantes via la similitude de chaînes de caractères et de plongements), un score de confiance des relations (combinant la certitude du classificateur et la fréquence), et un filtrage de graphe.
   • Nature Dynamique : Contrairement aux approches statiques (ex. : ConceptNet), ce graphe est construit à partir du texte d'entrée lui-même, capturant des entités spécifiques au domaine et à la session (ex. : un modèle d'imprimante spécifique et ses codes d'erreur uniques) qui n'existent pas dans les bases de connaissances pré-construites.

2. Génération de Plongements Enrichis par le Graphe :

   • Encodage : Les entités extraites sont encodées à l'aide de Réseaux de Convolution de Graphes Relationnels (R-GCN) pour capturer à la fois les propriétés individuelles et les rôles structurels via un passage de messages multi-sauts.
   • Fusion : Ces plongements structurels de graphe sont fusionnés avec les plongements d'entités textuelles (dérivés du modèle de langue de base) via une attention cross-modale. Cela crée une représentation hybride qui capture à la fois la sémantique de surface et les motifs relationnels plus profonds.

3. Architecture de Mémoire Multi-Composante :

   • KGERMAR maintient trois banques de mémoire spécialisées :
     • Mémoire Contextuelle : Stocke les paires clé-valeur au niveau du jeton des couches du modèle de langue (suivant ERMAR).
     • Mémoire Sémantique : Stocke les plongements textuels denses de segments non chevauchants (utilisant BGE-M3).
     • Mémoire Structurelle : Stocke les plongements fusionnés texte-entité du graphe.
   • Toutes les banques sont mises à jour de manière incrémentale et gérées via une éviction LRU lorsque la capacité est atteinte.

4. Mécanisme de Récupération Hybride :

   • Pour une requête donnée, le système effectue une récupération parallèle top-K à partir des trois banques.
   • Les signaux de récupération sont combinés à l'aide de poids de fusion appris ($\lambda_1, \lambda_2, \lambda_3$) pour produire un score final.
   • Les représentations récupérées sont injectées dans le modèle de langue via une attention causale de récupération dans les couches supérieures désignées.

Objectifs d'Entraînement
Le modèle utilise l'apprentissage multi-tâche pour optimiser simultanément quatre objectifs :

• Perte de Modélisation de Langage ($L_{LM}$) : Prédiction du jeton suivant en utilisant le contexte récupéré.
• Perte de Plongement de Graphe de Connaissances ($L_{KG}$) : Prédiction de liens via l'entropie croisée binaire pour entraîner le R-GCN.
• Perte de Classement de Récupération ($L_{retrieval}$) : Classement basé sur la marge pour garantir que les contextes pertinents soient mieux scorés que les contextes non pertinents.
• Perte d'Alignement Cross-Modale ($L_{alignment}$) : Assure la cohérence entre les représentations textuelles et structurelles.

Contributions Clés

1. Construction de Graphe de Connaissances Dynamique : Un pipeline qui extrait des entités et des relations spécifiques au domaine à partir du texte d'entrée en temps réel lors de l'inférence, permettant une adaptation à des contenus spécialisés absents des bases de connaissances statiques.
2. Mémoire Multi-Composante Hybride : Une architecture intégrant des représentations contextuelles, sémantiques et structurelles. La banque de mémoire structurelle capture spécifiquement les relations entre entités qui sont complémentaires et souvent plus discriminantes que la pure similitude sémantique.
3. Validation Empirique : Démonstration que l'incorporation de connaissances structurelles produit des améliorations substantielles de la modélisation de contexte long, atteignant jusqu'à 8,5 % de perplexité en moins et une efficacité de mémoire 2 à 2,5 fois supérieure par rapport aux bases solides.

Résultats Expérimentaux
Le cadre a été évalué sur SlimPajama (84,7K exemples pour l'affinage) et trois jeux de données de référence : WikiText-103, PG-19 et Proof-pile.

• Perplexité : KGERMAR a atteint la perplexité la plus basse parmi tous les modèles augmentés par la mémoire de 3B sur des longueurs de contexte allant de 1K à 32K jetons. Sur WikiText-103, il a surpassé significativement ERMAR et MemLong (ex. : 7,14 contre 7,80 de perplexité à 16K jetons).
• Apprentissage en Contexte (ICL) : Sur cinq tâches de NLU (SST-2, MR, Subj, SST-5, MPQA), KGERMAR a obtenu la précision moyenne la plus élevée dans les configurations à 4-shots (78,6 %) et 20-shots (81,2 %), surpassant ERMAR et MemLong. Les gains étaient particulièrement notables dans les tâches nécessitant une discrimination au niveau de l'entité (ex. : Subj, MPQA).
• Efficacité de la Mémoire : KGERMAR a démontré une efficacité de mémoire supérieure, maintenant une utilisation de mémoire 2 à 2,5 fois plus faible que les bases de référence. À 32K jetons, il a utilisé 36 % de mémoire de pointe en moins qu'ERMAR et 41 % de moins que MemLong. La mémoire normalisée par jeton s'est améliorée à mesure que la longueur du contexte augmentait, validant sa capacité à atténuer la croissance quadratique de la mémoire des mécanismes d'attention standards.
• Comportement de Mise à l'Échelle : Bien que KGERMAR présente une latence par jeton légèrement plus élevée due à la construction du graphe en temps réel, il démontre une croissance de latence quasi linéaire avec la longueur de la séquence, ce qui le rend plus prévisible pour le déploiement en contexte long que ERMAR.

Signification
Le document affirme que KGERMAR fait progresser la modélisation de langage augmentée par la mémoire en passant d'une dépendance aux bases de connaissances fixes et statiques à une représentation de connaissance adaptative au domaine. En construisant dynamiquement des graphes de connaissances à partir du contexte d'entrée, le système peut raisonner sur les états des entités et les progressions causales que la similitude sémantique seule ne peut capturer. Cette approche offre une solution pratique pour la modélisation de contexte long sur du matériel courant, offrant un équilibre entre une haute qualité de modélisation et une réduction significative de la surcharge de mémoire, ce qui est crucial pour les applications impliquant le support technique, la littérature scientifique et l'analyse de documents juridiques où le raisonnement structurel est primordial.