CItruS: Chunked Instruction-aware State Eviction for Long Sequence Modeling

TL;DR

This work introduces 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, and designs a method for chunked sequence processing to further improve efficiency.

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) without affecting 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.

Figures (12)

• Figure 1: One sample from attention distributions in the 16th layer of the Mistral 7B Instruct model applied to the Qasper dataset. The attention distributions are calculated from a document context and an instruction text to the key-value cache. The x-axis represents different positions within the key-value cache, while the y-axis represents the attention weights. The positions are reordered by descending attention weights from the context, and positions with low attention weights are omitted for clarity.
• Figure 2: The illustration of our experiments that apply intersection calculation to explore the information neglect problem in state eviction models.
• Figure 3: The difference between the top-$k$ hidden states selected by the instruction text and the document context with the $k$ set as $20$, conducted with Mistral 7B Instruct. Context-instruction intersection represents the overlap between the top-$k$ hidden states selected by the attention distribution from one piece of the context in the long document and the instruction text to a key-value cache.
• Figure 4: The illustration of our proposed different subprocesses for task-specific long sequence modeling. Each process serves as different roles.
• Figure 5: The illustration of different cache designs for our proposed Standard CSE and CItruS.