I nfini P ot: Infinite Context Processing on Memory-Constrained LLM s (2024)

Abstract: Handling long input contexts remains a significant challenge for Large Language Models (LLMs), particularly in resource-constrained environments such as mobile devices. Our work aims to address this limitation by introducing InfiniPot, a novel KV cache control framework designed to enable pre-trained LLMs to manage extensive sequences within fixed memory constraints efficiently, without requiring additional training. InfiniPot leverages Continual Context Distillation (CCD), an iterative process that compresses and retains essential information through novel importance metrics, effectively maintaining critical data even without access to future context. Our comprehensive evaluations indicate that InfiniPot significantly outperforms models trained for long contexts in various NLP tasks, establishing its efficacy and versatility. This work represents a substantial advancement toward making LLMs applicable to a broader range of real-world scenarios.

## introduction large language models (llms) have revolutionized the field of natural language processing (nlp) by achieving unprecedented performance across diverse tasks. however, their capacity to handle long input contexts remains notably limited, primarily due to the predefined maximum context length set during pretraining @xcite . moreover, even within these trained context lengths, llms often struggle to effectively process and maintain coherence over extended sequences @xcite @xcite . the challenge becomes further pronounced in resource-constrained environments such as mobile devices, where memory and computational power are significantly limited. in long context scenario, key-value (kv) cache size operates as a hard constraint on the available memory, making llm inference even more challenging. efficiently processing long sequences is essential for enabling applications such as document summarization and complex question answering directly on edge devices, yet current llms fall short of this capability. traditional approaches to addressing long context inputs in llms include either increasing the memory capacity to match the long context requirements or using streaming inputs to process information incrementally; however, both methods have substantial drawbacks. specifically, increasing memory to accommodate long contexts is often impractical, especially in on-device scenarios due to hardware constraints. moreover, current streaming-based solutions typically leverage only recent context for memory saving, thus failing to fully capitalize on the entire input context @xcite @xcite . several attempts have been made to manage long contexts in llms, exploring techniques like parameterefficient fine-tuning (peft) and recurrent attention mechanisms @xcite @xcite @xcite . these approaches, however, typically require significant memory increments or additional training steps, leading to inefficiencies and becoming impractical. other methods have tried modifying positional encoding to prevent out-ofdistribution issues but do not address the increased computational memory demands @xcite @xcite . additionally, there have been efforts to utilize key-value (kv)-cache compression. existing kv-cache compression methods primarily focus on compressing prior contexts during the generation stage, rather than effectively managing long input contexts with strict memory constraints @xcite @xcite . in transformer-based auto-regressive llms, the attention mechanism assigns scores to each token based on future context. however, in situations where full context is unavailable, identifying critical input parts becomes challenging, often resulting in information loss and suboptimal performance in memoryconstrained settings. to address these challenges, we propose infinipot, a novel kv-cache control framework that allows pre-trained llms to handle infinitely long contexts within fixed memory constraints. our model works akin to a virtual 'pot' that processes incoming ingredients. when the pot nears overflow, it distills unnecessary parts and retains only the essentials. similarly, infinipot serves as a 'memory pot' that processes incoming token sequence. when the number of kv-cache entries approaches its limit, it compresses the current kv-cache context into a smaller set of entries. this consume-and-compress cycle continues, preventing kv-cache overflow. to achieve this, we introduce continual context distillation (ccd), an iterative process that effectively obtains and retains essential information through novel importance metrics. the catalyst prompt (cap) consists of carefully designed volatile prompts injected right before the kv-cache reaches its limit, providing strong guidance in generating attention scores and approximating the future importance of tokens within a finite context. meanwhile, novelty under compression (nuc) score prioritizes new information that the pretrained model or the previous context has not encountered by assigning higher importance to such novel content. by combining cap and nuc, our approach can distinguish representative and novel tokens from less critical ones. this enables the management of long contexts efficiently through multiple ccd cycles. the proposed methodology promises to empower llms to handle extended sequences effectively, making them applicable to a broader range of nlp tasks. our thorough evaluations demonstrate that infinipot enables pre-trained models to achieve performance comparable to, or even surpassing, models explicitly trained to handle long contexts on various long-context nlp tasks, all without requiring additional training. in summary, our key contributions are as follows: ## infinipot: infinite context processing on memory-constrained llms in this section, we explain infinipot, the proposed kv-cache control framework that enables memoryconstrained llm to handle extremely long contexts. we first start from the problem definition (sec. 3.1). then, we propose a novel context compression method, continual context distillation (ccd) (sec. 3.2). in particular, the method consists of two unique components, namely catalyst prompt (cap) (sec. 3.3.1) and novelty under compression (nuc) (sec. 3.3.2). finally, we introduce an effective rope position assigning policy that prevents out-of-distribution position problem after kv-cache compression (sec. 3.5). ## conclusion in this paper, we addressed the challenge of enabling large language models (llms) to handle long input contexts efficiently in memory-constrained environments. we proposed infinipot, a novel kv-cache control framework that uses continual context distillation (ccd) to iteratively compress and distill essential information, even without access to future context. our evaluations demonstrated that infinipot-equipped llms can manage extended sequences effectively, achieving performance comparable to or surpassing models explicitly trained for long-context tasks. this work significantly extends the capabilities of pre-trained llms, making them more versatile and applicable to a broader range of nlp tasks without requiring additional training.