Retrieval augmentation is a powerful but expensive method to make language models more knowledgeable about the world. Memory-based methods like LUMEN 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.
@article{arxiv.2308.14903,
title = {MEMORY-VQ: Compression for Tractable Internet-Scale Memory},
author = {Yury Zemlyanskiy and Michiel de Jong and Luke Vilnis and Santiago Ontañón and William W. Cohen and Sumit Sanghai and Joshua Ainslie},
journal= {arXiv preprint arXiv:2308.14903},
year = {2023}
}