Knowledge Graph Enhanced Memory-Augmented Retrieval for Long Context Modeling

2026-06-11

AI Summary

This work addresses the challenge of maintaining consistent entity states and relational coherence across thousands of tokens in long-context language modeling, where semantic similarity alone proves insufficient for capturing deep structural dependencies. The authors propose a dynamic, context-aware knowledge graph construction method that extracts entities and relations in real time during inference, enabling domain-adaptive knowledge representation without reliance on a fixed external knowledge base. By integrating three complementary memory stores—contextual, semantic, and structural—and employing a retrieval mechanism driven by graph embeddings and learned attention weights, the model substantially enhances its capacity to model long-range dependencies. Evaluated across context lengths from 1K to 32K and five natural language understanding tasks, the approach achieves up to an 8.5% reduction in perplexity and a 2–2.5× improvement in memory efficiency.

Abstract

Long-context language modeling requires not only extending context windows but maintaining coherent understanding of entity states and relationships across thousands of tokens -- a challenge that semantic similarity alone cannot address. KGERMAR addresses this by constructing dynamic, context-specific knowledge graphs from input text during inference, enabling domain-adaptive retrieval that leverages both semantic similarity and explicit entity relationships. The framework performs real-time entity and relation extraction to build contextual knowledge graphs, then integrates graph-structural embeddings with textual semantics through a multi-component memory architecture. Three memory banks -- contextual, semantic, and structural -- are maintained with retrieval signals fused via learned weights to capture both surface-level semantics and deeper relational patterns. Evaluated on SlimPajama (84.7K training examples), WikiText-103 (4,358 examples), PG-19 (100 examples), and Proof-pile (46.3K examples), KGERMAR achieves up to 8.5\% lower perplexity and 2--2.5x better memory efficiency than memory-augmented baselines across context lengths from 1K to 32K tokens, with superior in-context learning performance across five NLU tasks. The dynamic knowledge graph construction approach advances memory-augmented language modeling by enabling domain-specific knowledge representation that adapts to input contexts rather than relying on fixed knowledge bases.

Problem

• long-context language modeling
• entity relationships
• knowledge graph
• memory-augmented retrieval
• context coherence

Innovation

• knowledge graph
• memory-augmented retrieval
• long-context modeling
• dynamic entity-relation extraction
• multi-component memory architecture