Neurocache: Efficient Vector Retrieval for Long-range Language Modeling

TL;DR

Neurocache is an approach to extend the effective context size of large language models (LLMs) using an external vector cache to store its past states using an efficient k-nearest-neighbor algorithm to retrieve relevant past states and incorporate them into the attention process.

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

Figures (3)

• Figure 1: Performance and Scalability of Neurocache vs. Memorizing Transformers wu-etal-2022-memorizingtrans 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.
• Figure 2: Documents are segmented into sequences of $n$ tokens and processed sequentially through a Transformer decoder stack. For each text segment, mid-layer hidden states $H \in \mathbb{R}^{n \times h}$ are projected into a compact representation $C \in \mathbb{R}^{n \times d}$ using a learned weight matrix $W_p \in \mathbb{R}^{d \times h}$. This projection enhances the efficiency of $k$NN retrieval of the most relevant past states $C_{ret} \in \mathbb{R}^{n \times k \times d}$ from the cache $C_{cache}$. These states $C_{ret}$ are used by cache-augmented layers to generate keys/values for cache attention. The output of cache attention is added to the self-attention output before being fed to the feed-forward network (FFN). Finally, the cache $C_{cache}$ is updated to include $C$ while maintaining a constant size of $m$ entries.
• Figure 3: This implementation showcases the cache-attention computation in the model. It calculates the attention weights through a dot product between the queries and keys, applies a softmax to obtain probabilities, and then computes the weighted sum of the extended values to generate the final attention output.