CI tru S : Chunked Instruction-aware State Eviction for Long Sequence Modeling (2024)

Abstract: Long sequence modeling has gained broad interest as large language models (LLMs) continue to advance. Recent research has identified that a large portion of hidden states within the key-value caches of Transformer models can be discarded (also termed evicted) withoutaffecting the perplexity performance in generating long sequences. However, we show that these methods, despite preserving perplexity performance, often drop information that is important for solving downstream tasks, a problem which we call information neglect. To address this issue, we introduce Chunked Instruction-aware State Eviction (CItruS), a novel modeling technique that integrates the attention preferences useful for a downstream task into the eviction process of hidden states. In addition, we design a method for chunked sequence processing to further improve efficiency. Our training-free method exhibits superior performance on long sequence comprehension and retrieval tasks over several strong baselines under the same memory budget, while preserving language modeling perplexity. The code and data have been released at https://github.com/ybai-nlp/CItruS.

## introduction recent advances in large language models (llms) have raised interest in long sequence modeling @xcite . several studies have found that information relevant to the next token prediction task often accumulates in the hidden representations of just a few tokens, and the attention distributions tend to focus sparsely on these tokens @xcite @xcite . this observation has resulted in methods that model longer sequences by evicting unnecessary key-value caches during the language modeling process @xcite , mostly based on the attention weights each token receives from the following context. however, these methods achieve limited performance on downstream tasks that require specific information from long documents (e.g., question answering), suggesting that they struggle to retain the detailed information necessary for such tasks. we refer to this condition as the information neglect problem. this issue arises because the cache acquired through state eviction is based only on the local document context. there is no explicit signal for the model to ensure that it is useful for solving downstream tasks. consider figure 1 , which shows two attention distributions-one from a document context and one from an instruction prompt-when applying the mistral 7b instruct model to a sample from the qasper dataset. note that the two differ substantially in their weighting of positions, suggesting that the document context-derived attention weights may not capture well the task specified by the instructions. 1 in this paper, we propose to address this information neglect issue by incorporating the instruction text into the state eviction process. our method, chunked instruction-aware state eviction (citrus), decomposes long sequence processing into two different subprocesses: language modeling and task solving. for the language modeling process, we propose chunked state eviction, splitting the long sequence input into large chunks while maintaining a cache that only stores the most important key-value states, which we show allows the model to efficiently and effectively encode long documents. as for the task-solving process, we propose an instruction-aware cache, either independent of or shared with the language modeling cache, which maintains the specific detailed information required to generate responses in downstream settings. the instruction-aware cache is then used to generate the final response for solving the task. our approach can be seen as analogous to ideas from cognitive science that language and thought can be disentangled in human language processing @xcite , we evaluate citrus on three tasks: long document reading comprehension, knowledge retrieval, and language modeling. our approach improves downstream task performance over several strong baselines by large margins and enables the retrieval of desired information hidden within a long document of up to one million tokens. furthermore, the model maintains high language modeling performance with a low perplexity. notably, citrus is applicable to all the transformer-based decoderonly model without any further training, improving the model's ability to conduct downstream tasks for input sequences with arbitrary lengths. overall, our contributions are summarized as follows: 1) we define and demonstrate the information neglect problem in state-eviction methods. ## the information neglect problem in this section, we demonstrate the information neglect problem of existing state eviction methods. state eviction methods have two basic elements: a key-value cache c that maintains the most important hidden states for language modeling and a strategy s to evict unnecessary states from the keyvalue cache, thereby making room for new states. by iteratively evicting the most unnecessary tokens from the cache, the model achieves the capability to model long sequences of arbitrary lengths. s is usually based on the attention weight a cache state receives from tokens later in the sequence. the information neglect problem stems from the observation that the preserved states useful for language modeling are not necessarily the ones for a downstream task (e.g., answering a specific question). we demonstrate this by measuring the difference between the top-k states selected by a document context compared to those selected by a specific instruction text (figure 2 ). specifically, we select one context and encode it to acquire a cache that could be evicted (i.e, context 1 in figure 2 ). then, we use another piece of context (i.e., context 2 in figure 2 ) and the instruction text, both with the same length, to evict the cache separately, retaining the top-k most important hidden states. by computing the overlap of the differently evicted caches, we draw conclusions about the information neglect scenarios during the eviction-based language modeling process. more experimental setup for these experiments is shown in appendix a. we use the same setting to acquire the results in figure 1 . we conduct this experiment on the full test set of the qasper dataset @xcite . we use the averaged attention score of all the tokens from one piece of text to the cache to select the most important states, which is further described in section 4.1. as shown in figure 3 , the hidden states focused on by the document context and the downstream instruction text are remarkably different, reflected by an intersection ratio lower than 0.2 in the middle layers. supported by the above experiments, we claim that if only the attention distribution of the context chunk is used to select the key-value states relevant to language modeling, some information specifically related to the final instruction text will be discarded during the encoding of the document, possibly decrease the task performance. a similar line of work that models long sequence with sliding-window-based methods @xcite also suffers from information neglect problems, where we provide detailed description in appendix d. additionally, we conduct another set of experiments that demonstrate the performance degradation of the standard state eviction models compared to the standard models that use the full text as input. results supporting the presence of information neglect problem are presented in appendix e. ## methods to address the problem of information neglect, we propose to decompose the inference procedure of large language models into two different subprocesses: the language modeling process and the task solving process, shown in figure 4 . for the language modeling process, we propose to use chunked state eviction methods to make the modeling process more efficient. for the task solving process, we propose instruction-aware state eviction, using the hidden states of the final instruction prompt as an additional instruction-aware query to extract and preserve the task-related information in a keyvalue cache. then, we utilize this key-value cache to generate a task-specific response. for downstream tasks with a long document input d and a final instruction i (a piece of text that prompt the model to conduct the downstream tasks), our proposed method generates a corresponding response r according to i. ## conclusion we have proposed citrus, an inference-time state eviction method for large language models (llms) that improves their performance on long sequence downstream tasks. it features a large chunked sequence processing procedure and an instructionaware cache that helps with solving downstream tasks. experiments on long document reading comprehension, knowledge retrieval, and language modeling show the utility of our method compared to strong baselines. our work demonstrates the possibility of generalizing standard llms trained on text constrained to certain lengths to processing longer sequences without any parameter adjustments. our evaluation mainly focuses on retrieving task-related information from a long document. future work may consider extending more high-level abilities (e.g., multi-hop and compositional reasoning) to the long sequence regime. moreover, trainable components can be further introduced to facilitate this process.