Knowledge Graph Enhanced Memory-Augmented Retrieval for Long Context Modeling

Technische Samenvatting: KGERMAR

Probleemstelling
Taalmodellering met lange contexten staat voor een fundamentele uitdaging die verder gaat dan het simpelweg uitbreiden van contextvensters: het behouden van een coherent begrip van entiteitstoestanden en relaties over duizenden tokens heen. Standaardbenaderingen lijden aan twee primaire beperkingen. Ten eerste lijden uitgebreide-attention modellen (bijv. LongLoRA, YARN) aan "lost-in-the-middle" degradatie, waarbij informatie in het midden van een lange context systematisch minder aandacht krijgt. Ten tweede vertrouwen geheugen-geaugmenteerde retrieval-systemen (bijv. MemLong, ERMAR) uitsluitend op semantische gelijkenis. Hoewel ze effectief zijn voor het ophalen van passages met overeenkomende termen, slagen ze er niet in om onderscheid te maken tussen passages die louter gerelateerde termen toevallig noemen en passages die expliciete causale of structurele relaties beschrijven tussen dezelfde entiteiten. Deze beperking is cruciaal in scenario's die vereisen dat er structureel wordt geredeneerd, zoals het volgen van de progressie van een fouttoestand van een apparaat tijdens een langdurige probleemoplossende sessie, waarbij de relevante context geen lexicale overlap kan delen met de query, maar wel diepe structurele verbindingen deelt.

Methodologie
Het artikel stelt KGERMAR (Knowledge Graph Enhanced Memory-Augmented Retrieval) voor, een raamwerk dat tijdens de inferentie dynamische, contextspecifieke kennisgrafen construeert om grafiek-structurele embeddings te integreren met tekstuele semantiek. Het systeem werkt via vier geïntegreerde componenten:

1. Contextuele Kennisgrafiek Constructie:

   • Extractie: Het systeem voert real-time Named Entity Recognition (NER) uit met een BERT-large model (gefine-tuned op CoNLL-2003) en Relation Extraction met een BERT-base model (gefine-tuned op TACRED).
   • Grafiekopbouw: Het extraheert entiteiten en relaties (bijv. CAUSEDBY, PARTOF, USEDFOR) binnen een glijdend venster. Om de kwaliteit te waarborgen, maakt het gebruik van entiteitsconsolidatie (het samenvoegen van varianten via string- en embedding-gelijkenis), relatie-betrouwbaarheidsscores (combinatie van classifier-zekerheid en frequentie) en grafiekfiltering.
   • Dynamisch Karakter: In tegenstelling tot statische benaderingen (bijv. ConceptNet), wordt deze grafiek gebouwd vanuit de invoertekst zelf, waardoor domeinspecifieke en sessiespecifieke entiteiten worden vastgelegd (bijv. een specifiek printermodel en de unieke foutcodes daarvan) die niet bestaan in vooraf geconstrueerde kennisbases.

2. Grafiek-verrijkte Embedding Generatie:

   • Codering: Geëxtraheerde entiteiten worden gecodeerd met behulp van Relational Graph Convolutional Networks (R-GCN) om zowel individuele eigenschappen als structurele rollen te vangen via multi-hop message passing.
   • Fusie: Deze grafiek-structurele embeddings worden gefuseerd met tekstuele entiteit-embeddings (afgeleid van het basis taalmodel) via cross-modale attention. Dit creëert een hybride representatie die zowel oppervlakkige semantiek als diepere relationele patronen vastlegt.

3. Multi-component Geheugenarchitectuur:

   • KGERMAR onderhoudt drie gespecialiseerde geheugenbanken:
     • Contextueel Geheugen: Slaat token-niveau key-value paren op van de taalmodellagen (volgens ERMAR).
     • Semantisch Geheugen: Slaat dense tekst-embeddings op van niet-overlappende chunks (met behulp van BGE-M3).
     • Structureel Geheugen: Slaat de gefuseerde tekst-grafiek entiteit-embeddings op.
   • Alle banken worden incrementeel bijgewerkt en beheerd via LRU-evictie wanneer de capaciteit wordt bereikt.

4. Hybride Retrieval Mechanisme:

   • Voor een gegeven query voert het systeem parallelle top-K retrieval uit vanuit alle drie de banken.
   • Retrieval-signalen worden gecombineerd met behulp van geleerde fusiegewichten ($\lambda_1, \lambda_2, \lambda_3$) om een definitieve score te produceren.
   • De opgehaalde representaties worden in het taalmodel geïnjecteerd via retrieval causal attention in aangewezen hogere lagen.

Trainingsdoelstellingen
Het model maakt gebruik van multi-task learning om vier doelstellingen gelijktijdig te optimaliseren:

• Taalmodellering Verlies ($L_{LM}$): Next-token voorspelling met opgehaalde context.
• Kennisgrafiek Embedding Verlies ($L_{KG}$): Link-predictie via binary cross-entropy om de R-GCN te trainen.
• Retrieval Ranking Verlies ($L_{retrieval}$): Margin-gebaseerde ranking om ervoor te zorgen dat relevante contexten hoger worden gescoord dan irrelevante contexten.
• Cross-Modale Alignment Verlies ($L_{alignment}$): Zorgt voor consistentie tussen tekstuele en structurele representaties.

Belangrijkste Bijdragen

1. Dynamische Kennisgrafiek Constructie: Een pijplijn die tijdens de inferentie real-time domeinspecifieke entiteiten en relaties uit de invoertekst extraheert, wat adaptatie mogelijk maakt aan gespecialiseerde inhoud die afwezig is in statische kennisbases.
2. Hybride Multi-component Geheugen: Een architectuur die contextuele, semantische en structurele representaties integreert. De structurele geheugenbank legt specifiek entiteitsrelaties vast die complementair zijn aan, en vaak discriminatiever dan, pure semantische gelijkenis.
3. Empirische Validatie: Demonstratie dat het integreren van structurele kennis substantiële verbeteringen oplevert in long-context modellering, waarbij tot 8,5% lagere perplexiteit en 2–2,5x betere geheugenefficiëntie wordt behaald vergeleken met sterke baselines.

Experimentele Resultaten
Het raamwerk werd geëvalueerd op SlimPajama (84,7K voorbeelden voor fine-tuning) en drie benchmark datasets: WikiText-103, PG-19, en Proof-pile.

• Perplexity: KGERMAR behaalde de laagste perplexiteit onder alle geheugen-geaugmenteerde 3B modellen over contextlengtes van 1K tot 32K tokens. Op WikiText-103 presteerde het significant beter dan ERMAR en MemLong (bijv. 7,14 vs. 7,80 perplexiteit bij 16K tokens).
• In-Context Learning (ICL): Op vijf NLU-taken (SST-2, MR, Subj, SST-5, MPQA) behaalde KGERMAR de hoogste gemiddelde nauwkeurigheid in zowel 4-shot (78,6%) als 20-shot (81,2%) settings, waarbij het ERMAR en MemLong overtrof. De winst was bijzonder merkbaar in taken die entiteits-niveau discriminatie vereisen (bijv. Subj, MPQA).
• Geheugenefficiëntie: KGERMAR demonstreerde superieure geheugenefficiëntie, waarbij het 2–2,5x minder geheugengebruik behield dan de baselines. Bij 32K tokens gebruikte het 36% minder piekgeheugen dan ERMAR en 41% minder dan MemLong. Het genormaliseerde geheugen-per-token verbeterde naarmate de contextlengte toenam, wat de bekwaamheid valideert om de kwadratische geheugengroei van standaard attention-mechanismen te mitigeren.
• Scaling Gedrag: Hoewel KGERMAR een iets hogere per-token latentie vertoont door real-time grafiekconstructie, vertoont het een bijna lineaire latentiegroei met de sequentielengte, wat het voorspelbaarder maakt voor long-context deployment dan ERMAR.

Significantie
Het artikel beweert dat KGERMAR de geheugen-geaugmenteerde taalmodellering vooruit helpt door de verschuiving van de afhankelijkheid van vaste, statische kennisbases naar domein-adaptieve kennisrepresentatie. Door dynamisch kennisgrafieken te construeren uit de invoercontext, kan het systeem redeneren over entiteitstoestanden en causale progressies die alleen door semantische gelijkenis niet kunnen worden gevangen. Deze aanpak biedt een praktische oplossing voor long-context modellering op standaard hardware, waarbij een balans wordt gevonden tussen hoge modelleringskwaliteit en een aanzienlijk verminderde geheugenoverhead, wat cruciaal is voor toepassingen in technische ondersteuning, wetenschelijke literatuur en juridische documentanalyse waar structureel redeneren essentieel is.