Pre-computed memory or on-the-fly encoding? A hybrid approach to retrieval augmentation makes the most of your compute

Michiel De Jong, Yury Zemlyanskiy, Nicholas Fitzgerald, Joshua Ainslie, Sumit Sanghai, Fei Sha, William W. Cohen
Proceedings of the 40th International Conference on Machine Learning, PMLR 202:7329-7342, 2023.

Abstract

Retrieval-augmented language models such as Fusion-in-Decoder are powerful, setting the state of the art on a variety of knowledge-intensive tasks. However, they are also expensive, due to the need to encode a large number of retrieved passages. Some work avoids this cost by pre-encoding a text corpus into a memory and retrieving dense representations directly. However, pre-encoding memory incurs a severe quality penalty as the memory representations are not conditioned on the current input. We propose LUMEN, a hybrid between these two extremes, pre-computing the majority of the retrieval representation and completing the encoding on the fly using a live encoder that is conditioned on the question and fine-tuned for the task. We show that LUMEN significantly outperforms pure memory on multiple question-answering tasks while being much cheaper than FiD, and outperforms both for any given compute budget. Moreover, the advantage of LUMEN over FiD increases with model size.

Cite this Paper

BibTeX
@InProceedings{pmlr-v202-de-jong23a, title = {Pre-computed memory or on-the-fly encoding? {A} hybrid approach to retrieval augmentation makes the most of your compute}, author = {De Jong, Michiel and Zemlyanskiy, Yury and Fitzgerald, Nicholas and Ainslie, Joshua and Sanghai, Sumit and Sha, Fei and Cohen, William W.}, booktitle = {Proceedings of the 40th International Conference on Machine Learning}, pages = {7329--7342}, year = {2023}, editor = {Krause, Andreas and Brunskill, Emma and Cho, Kyunghyun and Engelhardt, Barbara and Sabato, Sivan and Scarlett, Jonathan}, volume = {202}, series = {Proceedings of Machine Learning Research}, month = {23--29 Jul}, publisher = {PMLR}, pdf = {https://proceedings.mlr.press/v202/de-jong23a/de-jong23a.pdf}, url = {https://proceedings.mlr.press/v202/de-jong23a.html}, abstract = {Retrieval-augmented language models such as Fusion-in-Decoder are powerful, setting the state of the art on a variety of knowledge-intensive tasks. However, they are also expensive, due to the need to encode a large number of retrieved passages. Some work avoids this cost by pre-encoding a text corpus into a memory and retrieving dense representations directly. However, pre-encoding memory incurs a severe quality penalty as the memory representations are not conditioned on the current input. We propose LUMEN, a hybrid between these two extremes, pre-computing the majority of the retrieval representation and completing the encoding on the fly using a live encoder that is conditioned on the question and fine-tuned for the task. We show that LUMEN significantly outperforms pure memory on multiple question-answering tasks while being much cheaper than FiD, and outperforms both for any given compute budget. Moreover, the advantage of LUMEN over FiD increases with model size.} }
Endnote
%0 Conference Paper %T Pre-computed memory or on-the-fly encoding? A hybrid approach to retrieval augmentation makes the most of your compute %A Michiel De Jong %A Yury Zemlyanskiy %A Nicholas Fitzgerald %A Joshua Ainslie %A Sumit Sanghai %A Fei Sha %A William W. Cohen %B Proceedings of the 40th International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2023 %E Andreas Krause %E Emma Brunskill %E Kyunghyun Cho %E Barbara Engelhardt %E Sivan Sabato %E Jonathan Scarlett %F pmlr-v202-de-jong23a %I PMLR %P 7329--7342 %U https://proceedings.mlr.press/v202/de-jong23a.html %V 202 %X Retrieval-augmented language models such as Fusion-in-Decoder are powerful, setting the state of the art on a variety of knowledge-intensive tasks. However, they are also expensive, due to the need to encode a large number of retrieved passages. Some work avoids this cost by pre-encoding a text corpus into a memory and retrieving dense representations directly. However, pre-encoding memory incurs a severe quality penalty as the memory representations are not conditioned on the current input. We propose LUMEN, a hybrid between these two extremes, pre-computing the majority of the retrieval representation and completing the encoding on the fly using a live encoder that is conditioned on the question and fine-tuned for the task. We show that LUMEN significantly outperforms pure memory on multiple question-answering tasks while being much cheaper than FiD, and outperforms both for any given compute budget. Moreover, the advantage of LUMEN over FiD increases with model size.
APA
De Jong, M., Zemlyanskiy, Y., Fitzgerald, N., Ainslie, J., Sanghai, S., Sha, F. & Cohen, W.W.. (2023). Pre-computed memory or on-the-fly encoding? A hybrid approach to retrieval augmentation makes the most of your compute. Proceedings of the 40th International Conference on Machine Learning, in Proceedings of Machine Learning Research 202:7329-7342 Available from https://proceedings.mlr.press/v202/de-jong23a.html.

Related Material