Neurocache: Efficient Vector Retrieval for Long-range Language Modeling (2024)

Abstract: This paper introduces Neurocache, an approach to extend the effective context size of large language models (LLMs) using an external vector cache to store its past states. Like recent vector retrieval approaches, Neurocache uses an efficient k-nearest-neighbor (kNN) algorithm to retrieve relevant past states and incorporate them into the attention process. Neurocache improves upon previous methods by (1) storing compressed states, which reduces cache size; (2) performing a single retrieval operation per token which increases inference speed; and (3) extending the retrieval window to neighboring states, which improves both language modeling and downstream task accuracy. Our experiments show the effectiveness of Neurocache both for models trained from scratch and for pre-trained models such as Llama2-7B and Mistral-7B when enhanced with the cache mechanism. We also compare Neurocache with text retrieval methods and show improvements in single-document question-answering and few-shot learning tasks. We made the source code available under: https://github.com/alisafaya/neurocache

## introduction recent advancements in natural language processing have been significantly driven by the development of large language models (llms) such as gpt-3, gpt-4, llama, and llama2 @xcite @xcite . while demonstrating impressive capabilities, these models are constrained by limited context window sizes. this limitation becomes apparent in tasks that require understanding long documents, such as document summarization and academic literature review, where processing hundreds of thousands of tokens is necessary. various methods, including sparse attention @xcite @xcite , have been explored to address this @xcite on pg-19: the graph illustrates neurocache's consistently lower token perplexity and faster inference times across various cache sizes on the project gutenberg-19 dataset, demonstrating its efficiency and scalability. limitation. however, these approaches often struggle to utilize their extended contexts @xcite fully. recent research by @xcite shows that retrieval-augmented models with shorter contexts (4k tokens) can match the performance of models with longer contexts (16k/32k tokens), maintaining efficiency during inference. this emphasizes the potential of retrieval-augmented strategies in llms. in response to these challenges, we introduce neurocache. neurocache employs an efficient knearest-neighbor (knn) strategy for retrieving relevant past states from a compressed external vector cache. this approach is designed to optimize hidden state caching and retrieval, thereby enhancing language modeling quality and increasing inference speed. neurocache advances over exiting methods by reducing the cache size through the storage of compressed states, performing a single retrieval operation per token to boost inference speed, and extending the retrieval window to include neighboring states for improved language modeling and downstream accuracy. figure 1 illustrates the advantages of neurocache in terms of inference speed and language modeling accuracy over methods like memorizing transformers @xcite . our evaluation of neurocache encompasses both models trained from scratch and established pretrained models such as llama2-7b and mistral-7b @xcite , demonstrating its effectiveness in enhancing language models for downstream tasks. specifically, we highlight neurocache's improvements in singledocument question-answering and few-shot learning tasks when compared to traditional text retrieval methods. moreover, neurocache's integration extends the maximum context length of these models to 128k tokens, indicating its significant impact on long-document processing. in summary, neurocache represents a substantial step forward in addressing the challenges of processing long documents in llms, offering a blend of efficiency, adaptability, and enhanced performance. our comprehensive experiments and analysis showcase neurocache's potential in revolutionizing the understanding of long documents in natural language processing. ## related work transformers have made significant advancements in natural language processing but face challenges in processing long contexts. various methods have been developed to extend the context window while maintaining computational efficiency @xcite . recent methods include the continued training or fine-tuning of short-context language models @xcite , positional interpolation @xcite , alibi @xcite , and sparse and efficient attention designs @xcite @xcite . these approaches reflect the evolving landscape of solutions for managing extended attention windows in large language models (llms). however, language models still encounter difficulties in processing longer contexts @xcite . studies have indicated that retrievalaugmented models with shorter contexts (4k) can surpass models with longer contexts (16k/32k) in performance @xcite . prominent strategies in this area are text retrieval and vector retrieval. text retrieval involves identifying and processing the most relevant segments of long documents. vector retrieval, on the other hand, integrates relevant hidden representations of the input, like hidden states or key-value pairs, into the model. ## language modeling we assess neurocache's effectiveness via two experimental approaches: pre-training language mod-1 we assume an algorithm with quicksort-style partitioning is used. ## downstream evaluation we assess the performance of models augmented with neurocache, particularly llama2-7b and mistral-7b adapted on longpile, using seven distinct downstream tasks from the longbench suite @xcite . these tasks cover a range of scenarios, including single-document questionanswering (qa), multi-document qa, and fewshot learning. we utilize a zero-shot evaluation approach for the single-document and multidocument qa tasks. conversely, in the few-shot learning tasks, a small set of examples is provided to the models, serving as part of the extended context. ## conclusion this paper introduced neurocache, an approach designed to improve long document processing in language models. neurocache employs a k-nearestneighbor (knn) strategy for integrating relevant past states from compressed hidden representations, thus extending the context window of transformer decoders. notably, neurocache enhances the maximum context length of models like llama2-7b and mistral-7b to 128k tokens. our findings indicate that neurocache offers improvements in inference speed and language modeling accuracy. it demonstrates proficiency in singledocument question-answering and few-shot learning, though it faces challenges in multi-document scenarios. neurocache's competitive performance and adaptability highlight its potential utility in various applications. in summary, neurocache contributes to the field by enabling more efficient handling of extended contexts in existing language models. future work may explore further optimizations for multidocument tasks and the extension of neurocache to different model architectures and domains.