MEMORY - VQ : Compression for Tractable I nternet-Scale Memory (2024)

Abstract: Retrieval augmentation is a powerful but expensive method to make language models more knowledgeable about the world. Memory-based methods like LUMEN (de Jong et al., 2023a) pre-compute token representations for retrieved passages to drastically speed up inference. However, memory also leads to much greater storage requirements from storing pre-computed representations. We propose MEMORY-VQ, a new method to reduce storage requirements of memory-augmented models without sacrificing performance. Our method uses a vector quantization variational autoencoder (VQ-VAE) to compress token representations. We apply MEMORY-VQ to the LUMEN model to obtain LUMEN-VQ, a memory model that achieves a 16x compression rate with comparable performance on the KILT benchmark. LUMEN-VQ enables practical retrieval augmentation even for extremely large retrieval corpora.

## introduction retrieval augmentation is a common method to improve the factual knowledge of language models @xcite @xcite @xcite . retrieval provides a model with additional context in the form of text passages relevant to an input query. however, retrieval augmentation comes at an increased computational cost, as the model must process the retrieved passages on-the-fly. a recent line of work @xcite @xcite @xcite representations instead of raw text, which avoids the cost of reading retrieved passages from scratch. one such model, lumen, stands out for its strong performance, achieving 3x faster inference than standard fusion-in-decoder (izacard and grave, 2021) (fid) with minimal loss in quality. however, these pre-encoding memory models use much more storage than traditional retrievalaugmented models -lumen saves an embedding for each token in the corpus, which takes up much more space than token ids. our contribution is the first paper on compressing pre-encoded token memory representations. this compression makes memory methods such as lumen practical even for extremely large retrieval corpora. previous works (e.g., @xcite @xcite ) have focused on token representation compression for late-interaction reranking models. in contrast, our approach compresses the intermediate representations of a language model. these compressed representations are used as inputs into an llm, and the compression layers' parameters are trained alongside the rest of the model. ## background we aim to match fid and lumen performance in quality while reducing lumen storage requirements. we first describe fid and lumen, methods on which memory-vq is built, and their storage requirements. for an in-depth analysis, please see de @xcite . we follow up with background on vector quantization, including product quantization and vq-vae used for memory-vq. ## memory-vq we propose memory-vq, an efficient method for reducing storage requirements for memory-based models. the high-level idea is to compress memories using vector quantization techniques and store integer codes instead of the original memory vectors. codes are decompressed into vectors on the fly. applying the method to lumen yields the following lumen-vq model. to perform compressvq and decompressvq we apply product quantization, splitting each vector into subspaces and independently quantizing each subspace using vq-vae. codes are an exponential moving average of memory vectors assigned to the code in each batch. appendix a in van den oord et al. (2017) contains a detailed description. for training the compression layer jointly with the model, we follow the vq-vae recipe, but we avoid using the commitment loss in our experiments as it led to model divergence. to initialize the codebooks, we use a procedure similar to kmeans++ initialization @xcite . additionally, we perform code-book reset @xcite using the same procedure to re-initialize infrequently used codes. we divide memories into g subspaces, and if needed, pad memories with zeros to ensure divisibility. each subspace has c codes. therefore the storage requirement for each quantized vector is the number of subspaces multiplied by the number of bits required to represent each id, which is the logarithm of the number of codes. ## related work memory models retrieval augmentation can be computationally expensive due to the additional context that language models need to process. to mitigate this, memory models like lumen @xcite , @xcite , and others @xcite @xcite @xcite @xcite @xcite store pre-computed representations in memory. memory-vq focuses on improving the storage requirements for memory-based models. while our experiments involve the lumen (de jong et al., 2023a) model due to its strong performance, the method applies to a broader range of models. compression for late-interaction reranking memory-vq focuses on compression for lateinteraction memory models, while other works have explored compression for late-interaction reranking. for instance, sdr @xcite employs an autoencoder to reduce token representation dimensionality, followed by product quantization. becr @xcite utilizes localitysensitive hashing for token representation compression. cq @xcite learns vector quantization parameters by treating codes as learnable weights and uses gumbel-softmax for differentiable nearest code determination. finally, col-bertv2 @xcite proposes a custom compression scheme combining pq and integer quantization to handle reconstruction residuals. ## conclusion we introduced memory-vq, a novel approach for reducing the storage requirements of memoryaugmented language models without compromising performance. by employing vq-vae to compress token representations, we obtain a lumen model with 16x compression, denoted as lumen-vq. remarkably, lumen-vq maintains performance close to lumen and fid and benefits from lumen inference speed-ups with sharply reduced storage cost. using memory-vq, memory aug-mentation is a practical solution for drastic inference speedups with extensive retrieval corpora.