G raph R eader: Building Graph-based Agent to Enhance Long-Context Abilities of Large Language Models (2024)

Abstract: Long-context capabilities are essential for large language models (LLMs) to tackle complex and long-input tasks. Despite numerous efforts made to optimize LLMs for long contexts, challenges persist in robustly processing long inputs. In this paper, we introduce GraphReader, a graph-based agent system designed to handle long texts by structuring them into a graph and employing an agent to explore this graph autonomously. Upon receiving a question, the agent first undertakes a step-by-step analysis and devises a rational plan. It then invokes a set of predefined functions to read node content and neighbors, facilitating a coarse-to-fine exploration of the graph. Throughout the exploration, the agent continuously records new insights and reflects on current circumstances to optimize the process until it has gathered sufficient information to generate an answer. Experimental results on the LV-Eval dataset reveal that GraphReader using a 4k context window, consistently outperforms GPT-4-128k across context lengths from 16k to 256k by a large margin. Additionally, our approach demonstrates superior performance on four challenging single-hop and multi-hop benchmarks.

## introduction large language models (llms) have made great progress on natural language understanding and generation @xcite @xcite @xcite @xcite . however, transformer-based llms still struggle in handling long contexts due to the limitation of context window and memory usage. current techniques for solving the long-context tasks of llms can be divided into two perspectives: 1) model-level, which includes finetuning with modified positional embeddings @xcite * first four authors contributed equally. † corresponding author. ‡ project leader. graphreader outperforms existing open-sourced and closed-source models while demonstrating a scalable performance in very long contexts. in contrast, other models exhibit a significant decrease in performance as context length increases. @xcite ; @xcite @xcite , and applying transformer variants with modified attention mechanisms @xcite @xcite ; 2) agent-level, i.e., employing retrieval-augmented llm or agent to process long contexts with a limited context window llm @xcite . however, model-level methods typically train llms with target length texts, posing challenges in constructing training datasets and incurring high training costs @xcite . additionally, longcontext llms optimized with these methods tend to overlook crucial details in long contexts, known as "lost in the middle" @xcite , limiting their ability to address complex tasks, such as multi-hop questions. agent-level approaches transform input text into a tree @xcite or paginated pages @xcite , failing to capture multi-hop and long-range dependencies, thus limiting their effectiveness on very long contexts, as shown in figure 1 . to address these issues, we propose a graphbased agent named graphreader. as illustrated in figure 2 , graphreader first segments long texts into discrete chunks, extracts essential information, and compresses these into key elements and atomic facts. these key elements and facts are then used to construct a graph with nodes representing key elements and their associated atomic facts. this graph structure effectively captures long-range dependencies and multi-hop relationships within long text. subsequently, graphreader autonomously explores this graph using predefined functions, guided by a step-by-step rational plan. based on a given question, the agent progressively accesses information from coarse key elements and atomic facts to detailed original text chunks, taking notes and reflecting until it gathers sufficient information to generate an answer. in summary, our main contributions are threefold: ## related work long-context llms recent efforts @xcite @xcite have focused on positional interpolation (pi) to enhance long-context capabilities. however, these methods require training on full-length texts, leading to significant increases in data and training costs @xcite @xcite . thus, pose @xcite and skipalign @xcite investigate data skip strategy, but tend to neglect detailed information in long texts @xcite @xcite . furthermore, despite how extensively the context window is expanded, it remains constrained by a predefined fixed length. to address these limitations, transformer variants with modified attention mechanisms have been proposed @xcite @xcite . however, these models are prone to losing earlier information. retrieval retrieval augmented generation (rag) leverages an extensive database of documents to extract task-related information that aids in response generation. many efforts investigate various levels of retrieval granularity, including tokens @xcite , entities @xcite , and chunks @xcite . other approaches have explored diverse retrieval methods, such as bm25 @xcite and learning-based strategies @xcite @xcite . despite its capabilities, rag faces challenges in addressing complex questions due to difficulties in developing robust decision-making mechanisms. in contrast, we employ agents that use planning and reflection to gather essential information, effectively tackling complex problems. agent for retrieval recent work has increasingly leveraged llms as agents to tackle complex problems, utilizing their strong planning and reflection abilities @xcite . these abilities have been applied to complex tasks such as function call @xcite and kgqa @xcite . agents are also capable of retrieving unstructured information. for example, @xcite simulates human actions to search on internet for specific answers. additionally, memwalker @xcite and pearl @xcite organize documents into a tree structure, while readagent @xcite condenses documents into a gist memory directory. however, these approaches often struggle with multi-hop questions. kgp @xcite organizes documents into graphs, but it primarily uses the agent to generate queries, thereby not fully exploiting the agent's capabilities for planning and reflection. ## conclusion this paper introduces graphreader, a graph-based agent designed to enhance the long-context capabilities of large language models. graphreader organizes long texts into graph structures and employs an autonomous agent to explore the graph, successfully establishing long-range dependencies within a relatively small 4k context window. experiments demonstrate that graphreader outperforms gpt-4 with a 128k input length across various long-context single-hop and multi-hop questionanswering benchmarks.