MEMORY-VQ: Compression for Tractable Internet-Scale Memory
Yury Zemlyanskiy, Michiel de Jong, Luke Vilnis, Santiago Ontanon, William Cohen, Sumit Sanghai, Joshua Ainslie
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.
BibTeX
@inproceedings{zemlyanskiy-etal-2024-memory,
title = "{MEMORY}-{VQ}: Compression for Tractable {I}nternet-Scale Memory",
author = "Zemlyanskiy, Yury and
de Jong, Michiel and
Vilnis, Luke and
Ontanon, Santiago and
Cohen, William and
Sanghai, Sumit and
Ainslie, Joshua",
editor = "Duh, Kevin and
Gomez, Helena and
Bethard, Steven",
booktitle = "Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 2: Short Papers)",
month = jun,
year = "2024",
address = "Mexico City, Mexico",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2024.naacl-short.64/",
doi = "10.18653/v1/2024.naacl-short.64",
pages = "737--744"
}