HGMem: Hypergraph-based Working Memory to Improve Multi-step RAG for Long-Context Complex Relational Modeling
Abstract
Multi-step retrieval-augmented generation (RAG) has become a widely adopted strategy for enhancing large language models (LLMs) on tasks that demand global comprehension and intensive reasoning. Although many RAG systems incorporate a working memory to consolidate information, existing designs primarily function as a passive storage for isolated facts. This static nature overlooks crucial high-order correlations among primitive facts, thereby limiting models' capacity for multi-step reasoning and resulting in fragmented reasoning and weak global sense-making within extended contexts. We introduce HGMem, a hypergraph-based working memory system, extending the concept of memory beyond simple storage into a dynamic, expressive structure for complex reasoning and global understanding. In our approach, memory is represented as a hypergraph where hyperedges correspond to distinct memory units, enabling the progressive formation of high-order interactions within memory. This mechanism connects facts and thoughts around the focal problem, evolving the memory into an integrated and situated knowledge structure that provides strong propositions for deeper reasoning. We evaluate HGMem on several challenging global sense-making benchmarks. Extensive experiments and in-depth analyses demonstrate that our method consistently improves multi-step RAG and substantially outperforms strong baseline systems across diverse datasets.
Lay Summary
(1) While multi-step Retrieval-Augmented Generation (RAG) helps AI models answer complex queries, current systems rely on working memory mechanisms that treat information as a static list of primitive facts. This unstructured or rigidly schema-bound storage fails at global sense-making tasks, leaving the model with a fragmented view of complex, long-context relationships. (2) To address this, we propose HGMem, a working memory system that allows memorized content to actively evolve. HGMem structures memory as a hypergraph, which naturally encodes intricate, high-order correlations among multiple entities simultaneously rather than just simple pairs. As the AI interacts with the text, it progressively merges and updates these memory points to dynamically build a global perspective. (3) Extensive evaluations show that HGMem consistently outperforms strong multi-step RAG baselines on complex reasoning benchmarks. Notably, when powered by a medium-sized open-source model (Qwen2.5-32B), HGMem matches or exceeds the performance of GPT-4o, demonstrating a highly efficient path toward advanced contextual awareness.