SEGMENT +: Long Text Processing with Short-Context Language Models (2024)

Abstract: There is a growing interest in expanding the input capacity of language models (LMs) across various domains. However, simply increasing the context window does not guarantee robust performance across diverse long-input processing tasks, such as understanding extensive documents and extracting detailed information from lengthy and noisy data. In response, we introduce Segment+, a general framework that enables LMs to handle extended inputs within limited context windows efficiently. Segment+ utilizes structured notes and a filtering module to manage information flow, resulting in a system that is both controllable and interpretable. Our extensive experiments across various model sizes, focusing on long-document question-answering and Needle-in-a-Haystack tasks, demonstrate the effectiveness of Segment+ in improving performance.

## introduction language models (lms) have shown exceptional performance in a wide range of nlp tasks @xcite . due to the relatively short context window of most lms, they face unique challenges in contexts such as longdocument question answering, long-term memory maintenance, and processing lengthy, noisy contexts @xcite @xcite @xcite . efficiently processing long texts across various tasks remains a core challenge in this community. to reduce input length for handling long text, traditional retrieval is a simple and fast method but struggles with tasks requiring multiple pieces of information, often missing details and introducing noise @xcite . enhancements like query rewriters @xcite and figure 1 : this picture illustrates the use of short-context models to tackle long document question answering tasks in segment + . the process begins by gathering relevant context from the document for a specific question. only notes labeled 'keep' are used as the context to derive the final answer, avoiding noise. to access the entire text and perform multiple intermediate reasoning steps to enhance performance on complex tasks, memory management involves repeatedly using short-context lms combined with document pre/post-processing for managing long texts @xcite . for such systems, the key challenge lies in designing a mechanism to manage information flow across dif-ferent invocation times. however, past research often relies heavily on the model's inherent capabilities for planning and spontaneous decision-making, resulting in an uncontrolled reasoning process with noisy, free-form text expressions. to address the above issues and challenges, we propose the segment + , a robust and controllable framework that helps language models process long texts with a limited context window. the key insight of the segment + is capturing the characteristics of queries and designing two specific components to gather and merge different types of information from long inputs. as shown in figure 1 , the segment + agent has two stages. in stage 1, the model collects structured notes in parallel from all segments, with each note containing an evidence and a reasoning part. after filtering out unhelpful notes labeled as 'remove', the remaining notes proceed to stage 2. in stage 2, the notes are divided into batches, maintaining the same order as the input, with a maximum token limit. each batch of notes is then merged into one updated structured note. this process is iterated until the remaining notes can fit into the context window as the final context for answering the question. the challenge of processing long texts can be effectively tackled by using two types of components for information flow control. we notice that some questions require specific detailed information, while others need further reasoning across different parts of the content. therefore, we create an evidence component for gathering original sentences from the input, focusing on precision, and a reasoning component to help the model compress context into high-level semantic information, focusing on recall. this division, using both evidence and reasoning, makes the process both controllable and interpretable. significantly, retrieval and long-context language models can both benefit our method by either moderately increasing the processing context window or narrowing the input range that needs to be processed, which does not require high accuracy. in short, our contributions include: 1) we introduce a versatile framework for long context processing, applicable across language models of varying sizes and multiple text domains. 2) our method, leveraging a robust reasoning schema, outperforms other agent-based baselines and advanced long context models in long text processing tasks. 3) we conduct a thorough analysis of segment + , high-lighting the importance of structured information control. long context lms language models perform well in a variety of applications but struggle with large texts due to limited context windows @xcite @xcite . through techniques such as position interpolation and continuous pretraining, researchers have attempted to expand the context windows, thereby improving performance for both long and short document tasks @xcite . however, these approaches are limited by data quality and feasible window size constraints @xcite . besides, the models' inability to handle queries when key information is scattered across a large text is also a notable challenge @xcite . @xcite discovers that this issue arises from attention failure and can be alleviated by training models with a specially designed task. thus, the performance and robustness of models processing very long texts in a longer context remain uncertain, and limited resources restrict our ability to indefinitely expand the model's context window. our method addresses this challenge by dynamically optimizing the use of the available context window, thus allowing for the effective extraction of crucial information scattered throughout the text. this approach shifts the focus from merely widening the window to better utilizing the model's current context capabilities. memory management rather than processing long text in one go, we can use language models as agents to handle the long input task step by step. in such systems, memory is pivotal not only for storing out-of-window information but also as a foundation for lifelong learning through historical data analysis @xcite . using llm agents for long text tasks has its advantages. firstly, efficiently organizing and utilizing memory can improve performance in both documents qa and dialogue tasks @xcite . secondly, capturing the document structure and employing agent navigators across document segments is advantageous for document qa tasks @xcite . lastly, leveraging the task decomposition and plan-and-solve abilities of agents also benefits these tasks @xcite . our approach stands out by showcasing how it's effective across variously sized models, unlike others that require high language model capabilities. ## method the core challenge of processing long inputs with short-context models is how to control the information flow within different segments. in other words, how to retain the most useful information while using the fewest tokens. we first construct this process into a two-stage pipeline and then introduce each step in detail. ## experiments in this section, we apply the segment + to two distinct long text reasoning tasks: long document question answering and needle-in-a-haystack tasks. for the long document qa task, we analyze segment + 's ability to compress reading contexts and merge information efficiently in parallel ( §4.1). for the needle-in-a-haystack task, we evaluate segment + 's resistance to noise while accurately gathering essential sentences and do reasoning for final answer.( §4.2). understanding long documents has long been a common research issue in the nlp field, posing challenges due to the increasing length of texts and the complexities involved in comprehensive reasoning. ## conclusion in this paper, we introduce segment + , a simple yet effective plug-and-play methodology designed to augment the processing of long inputs within limited context windows, leveraging structured information flow control motivated by data characteristics and filtering mechanisms. our extensive experiments and analyses substantiate that segment + significantly enhances performance in long document question answering and noisy text processing, thereby illustrating its broad applicability across diverse domains in this field. when compared to agent-based methods, segment + not only achieves superior performance but also exhibits greater stability. furthermore, this information control schema holds potential for broader applications in scenarios requiring long input processing, such as in agent memory management and video information processing.