CoChat
Sign Up
PreprintCited by 2Topic ModelingInformation Retrieval and Search BehaviorNatural Language Processing Techniques

Accelerating Retrieval-Augmented Language Model Serving with Speculation

Open access
Zhang, Zhihao, Zhu, Alan, Yang, Lijie, Xu, Yihua, Li, Lanting, Phothilimthana, Phitchaya Mangpo, Jia, Zhihao·2024·arXiv
DOI 10.48550/arxiv.2401.14021arXiv 2401.14021
Open this paper in CoChat — free accountOpen PDF

Abstract

Retrieval-augmented language models (RaLM) have demonstrated the potential to solve knowledge-intensive natural language processing (NLP) tasks by combining a non-parametric knowledge base with a parametric language model. Instead of fine-tuning a fully parametric model, RaLM excels at its low-cost adaptation to the latest data and better source attribution mechanisms. Among various RaLM approaches, iterative RaLM delivers a better generation quality due to a more frequent interaction between the retriever and the language model. Despite the benefits, iterative RaLM usually encounters high overheads due to the frequent retrieval step. To this end, we propose RaLMSpec, a speculation-inspired framework that provides generic speed-up over iterative RaLM while preserving the same model outputs through speculative retrieval and batched verification. By further incorporating prefetching, optimal speculation stride scheduler, and asynchronous verification, RaLMSpec can automatically exploit the acceleration potential to the fullest. For naive iterative RaLM serving, extensive evaluations over three language models on four downstream QA datasets demonstrate that RaLMSpec can achieve a speed-up ratio of 1.75-2.39x, 1.04-1.39x, and 1.31-1.77x when the retriever is an exact dense retriever, approximate dense retriever, and sparse retriever respectively compared with the baseline. For KNN-LM serving, RaLMSpec can achieve a speed-up ratio up to 7.59x and 2.45x when the retriever is an exact dense retriever and approximate dense retriever, respectively, compared with the baseline.

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2401.14021

Cites

0 citation errors · 0 not checkable

Statistics

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Cites

The publisher hasn’t deposited this paper’s reference list.

Research works printed but not linked: 46

  1. Uri Alon, Frank Xu, Junxian He, Sudipta Sengupta, Dan Roth, and Graham Neubig. Neuro-symbolic language modeling with automaton-augmented retrieval. In International Conference on Machine Learning, pp. 468-485. PMLR, 2022.
  2. Akari Asai, Sewon Min, Zexuan Zhong, and Danqi Chen. Acl 2023 tutorial: Retrieval-based lan- guage models and applications. ACL 2023, 2023.
  3. Jonathan Berant, Andrew Chou, Roy Frostig, and Percy Liang. Semantic parsing on Freebase from question-answer pairs. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp. 1533-1544, Seattle, Washington, USA, October 2013. Association for Computational Linguistics. URL https://www.aclweb.org/anthology/D13-1160.
  4. Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Milli- can, George Bm Van Den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. In International conference on machine learning, pp. 2206-2240. PMLR, 2022.
  5. Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
  6. F Warren Burton. Speculative computation, parallelism, and functional programming. IEEE Trans- actions on Computers, 100(12):1190-1193, 1985.
  7. Charlie Chen, Sebastian Borgeaud, Geoffrey Irving, Jean-Baptiste Lespiau, Laurent Sifre, and John Jumper. Accelerating large language model decoding with speculative sampling. arXiv preprint arXiv:2302.01318, 2023. · arXiv
  8. Danqi Chen, Adam Fisch, Jason Weston, and Antoine Bordes. Reading wikipedia to answer open- domain questions. In 55th Annual Meeting of the Association for Computational Linguistics, ACL 2017, pp. 1870-1879. Association for Computational Linguistics (ACL), 2017.
  9. Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022. · arXiv
  10. Andrew Drozdov, Shufan Wang, Razieh Rahimi, Andrew McCallum, Hamed Zamani, and Mohit Iyyer. You can't pick your neighbors, or can you? when and how to rely on retrieval in the k nn-lm. arXiv preprint arXiv:2210.15859, 2022. · arXiv
  11. Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Mingwei Chang. Retrieval augmented language model pre-training. In International conference on machine learning, pp. 3929-3938.
  12. PMLR, 2020.
  13. Junxian He, Graham Neubig, and Taylor Berg-Kirkpatrick. Efficient nearest neighbor language models. arXiv preprint arXiv:2109.04212, 2021. · arXiv
  14. Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning. arXiv preprint arXiv:2112.09118, 2021. · arXiv
  15. Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi-Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave. Few-shot learning with re- trieval augmented language models. arXiv preprint arXiv:2208.03299, 2022. · arXiv
  16. Zhengbao Jiang, Frank F Xu, Luyu Gao, Zhiqing Sun, Qian Liu, Jane Dwivedi-Yu, Yiming Yang, Jamie Callan, and Graham Neubig. Active retrieval augmented generation. arXiv preprint arXiv:2305.06983, 2023. · arXiv
  17. Joao Gante. Assisted generation: a new direction toward low-latency text generation, 2023. URL https://huggingface.co/blog/assisted-generation. Preprint.
  18. Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535-547, 2019.
  19. Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. triviaqa: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension. arXiv e-prints, art. arXiv:1705.03551, 2017. · arXiv
  20. Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. Dense passage retrieval for open-domain question answering. arXiv preprint arXiv:2004.04906, 2020. · arXiv
  21. Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. arXiv preprint arXiv:1911.00172, 2019. · arXiv
  22. Omar Khattab and Matei Zaharia. Colbert: Efficient and effective passage search via contextualized late interaction over bert. In Proceedings of the 43rd International ACM SIGIR conference on research and development in Information Retrieval, pp. 39-48, 2020.
  23. Omar Khattab, Keshav Santhanam, Xiang Lisa Li, David Hall, Percy Liang, Christopher Potts, and Matei Zaharia. Demonstrate-search-predict: Composing retrieval and language models for knowledge-intensive nlp. arXiv preprint arXiv:2212.14024, 2022. · arXiv
  24. Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, et al. Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:453-466, 2019.
  25. Yaniv Leviathan, Matan Kalman, and Yossi Matias. Fast inference from transformers via speculative decoding. arXiv preprint arXiv:2211.17192, 2022. · arXiv
  26. Patrick Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, et al. Retrieval-augmented genera- tion for knowledge-intensive nlp tasks. Advances in Neural Information Processing Systems, 33: 9459-9474, 2020.
  27. Jimmy Lin, Xueguang Ma, Sheng-Chieh Lin, Jheng-Hong Yang, Ronak Pradeep, and Rodrigo Nogueira. Pyserini: A Python toolkit for reproducible information retrieval research with sparse and dense representations. In Proceedings of the 44th Annual International ACM SIGIR Con- ference on Research and Development in Information Retrieval (SIGIR 2021), pp. 2356-2362, 2021.
  28. Yu A Malkov and Dmitry A Yashunin. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions on pattern analysis and machine intelligence, 42(4):824-836, 2018.
  29. Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Rae Ying Yee Wong, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative llm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 2023. · arXiv
  30. Joon Sung Park, Joseph C O'Brien, Carrie J Cai, Meredith Ringel Morris, Percy Liang, and Michael S Bernstein. Generative agents: Interactive simulacra of human behavior. arXiv preprint arXiv:2304.03442, 2023. · arXiv
  31. Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
  32. Ori Ram, Yoav Levine, Itay Dalmedigos, Dor Muhlgay, Amnon Shashua, Kevin Leyton- Brown, and Yoav Shoham. In-context retrieval-augmented language models. arXiv preprint arXiv:2302.00083, 2023. · arXiv
  33. Juan Ramos et al. Using tf-idf to determine word relevance in document queries. In Proceedings of the first instructional conference on machine learning, volume 242, pp. 29-48. Citeseer, 2003.
  34. Stephen Robertson, Hugo Zaragoza, et al. The probabilistic relevance framework: Bm25 and be- yond. Foundations and Trends® in Information Retrieval, 3(4):333-389, 2009.
  35. Ohad Rubin and Jonathan Berant. Long-range language modeling with self-retrieval. arXiv preprint arXiv:2306.13421, 2023. · arXiv
  36. Keshav Santhanam, Omar Khattab, Jon Saad-Falcon, Christopher Potts, and Matei Zaharia. Colbertv2: Effective and efficient retrieval via lightweight late interaction. arXiv preprint arXiv:2112.01488, 2021. · arXiv
  37. Weijia Shi, Sewon Min, Michihiro Yasunaga, Minjoon Seo, Rich James, Mike Lewis, Luke Zettle- moyer, and Wen-tau Yih. Replug: Retrieval-augmented black-box language models. arXiv preprint arXiv:2301.12652, 2023. · arXiv
  38. Mitchell Stern, Noam Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autore- gressive models. Advances in Neural Information Processing Systems, 31, 2018.
  39. Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.09288, 2023. · arXiv
  40. Guanzhi Wang, Yuqi Xie, Yunfan Jiang, Ajay Mandlekar, Chaowei Xiao, Yuke Zhu, Linxi Fan, and Anima Anandkumar. Voyager: An open-ended embodied agent with large language models. arXiv preprint arXiv:2305.16291, 2023a. · arXiv
  41. Weizhi Wang, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, and Furu Wei. Aug- menting language models with long-term memory. arXiv preprint arXiv:2306.07174, 2023b. · arXiv
  42. Lee Xiong, Chenyan Xiong, Ye Li, Kwok-Fung Tang, Jialin Liu, Paul Bennett, Junaid Ahmed, and Arnold Overwijk. Approximate nearest neighbor negative contrastive learning for dense text retrieval. arXiv preprint arXiv:2007.00808, 2020. · arXiv
  43. Nan Yang, Tao Ge, Liang Wang, Binxing Jiao, Daxin Jiang, Linjun Yang, Rangan Majumder, and Furu Wei. Inference with reference: Lossless acceleration of large language models. arXiv preprint arXiv:2304.04487, 2023. · arXiv
  44. Zexuan Zhong, Tao Lei, and Danqi Chen. Training language models with memory augmentation. arXiv preprint arXiv:2205.12674, 2022. · arXiv
  45. Wangchunshu Zhou, Yuchen Eleanor Jiang, Peng Cui, Tiannan Wang, Zhenxin Xiao, Yifan Hou, Ryan Cotterell, and Mrinmaya Sachan. Recurrentgpt: Interactive generation of (arbitrarily) long text. arXiv preprint arXiv:2305.13304, 2023. · arXiv
  46. Xizhou Zhu, Yuntao Chen, Hao Tian, Chenxin Tao, Weijie Su, Chenyu Yang, Gao Huang, Bin Li, Lewei Lu, Xiaogang Wang, et al. Ghost in the minecraft: Generally capable agents for open-world enviroments via large language models with text-based knowledge and memory. arXiv preprint arXiv:2305.17144, 2023. Preprint. Table 8: Averaged latency measured in seconds over GPT-2. OPT and LLaMA-2 with the sparse retriever. Model Method Wiki QA WQ NQ Trivia QA GPT2 Baseline 7.41 ± 1.34 7.03 ± 1.15 7.23 ± 0.11 6.80 ± 0.09 RaLMSpec 5.18 ± 0.13 5.30 ± 0.95 5.34 ± 0.11 5.40 ± 0.03 RaLMSpec+P(20) 5.23 ± 0.23 4.58 ± 0.01 5.17 ± 0.05 5.50 ± 0.04 RaLMSpec+P(256) 6.88 ± 0.66 7.16 ± 1.34 6.76 ± 0.27 6.91 ± 0.13 RaLMSpec+S 5.62 ± 0.96 5.03 ± 0.68 5.24 ± 0.13 5.61 ± 0.11 RaLMSpec+A 5.34 ± 0.89 4.99 ± 0.86 4.76 ± 0.14 5.04 ± 0.12 RaLMSpec+P(20)SA 4.49 ± 0.09 4.57 ± 0.81 4.54 ± 0.02 4.99 ± 0.01 RaLMSpec+P(256)SA 6.66 ± 1.25 6.50 ± 1.39 5.54 ± 0.02 5.91 ± 0.03 OPT Baseline 7.68 ± 0.01 1.83 ± 0.01 2.62 ± 0.01 4.71 ± 0.02 RaLMSpec 5.63 ± 0.01 1.93 ± 0.01 2.60 ± 0.01 4.07 ± 0.02 RaLMSpec+P(20) 3.00 ± 0.01 2.06 ± 0.01 2.50 ± 0.02 4.27 ± 0.01 RaLMSpec+P(256) 7.13 ± 0.01 2.45 ± 0.01 3.22 ± 0.01 5.24 ± 0.02 RaLMSpec+S 2.79 ± 0.01 1.69 ± 0.01 2.32 ± 0.01 4.27 ± 0.01 RaLMSpec+A 5.27 ± 0.02 1.86 ± 0.01 2.28 ± 0.01 3.82 ± 0.02 RaLMSpec+P(20)SA 2.46 ± 0.01 1.57 ± 0.01 1.88 ± 0.01 3.59 ± 0.01 RaLMSpec+P(256)SA 6.37 ± 0.01 1.92 ± 0.04 2.44 ± 0.02 4.53 ± 0.01 LLaMA Baseline 10.75 ± 0.32 10.55 ± 0.07 10.72 ± 0.10 11.06 ± 0.25 RaLMSpec 11.47 ± 0.17 11.02 ± 0.31 11.10 ± 0.27 10.79 ± 0.20 RaLMSpec+P(20) 11.27 ± 0.14 11.04 ± 0.22 10.69 ± 0.15 10.66 ± 0.23 RaLMSpec+P(256) 12.83 ± 0.21 13.35 ± 0.81 12.48 ± 0.22 12.60 ± 0.36 RaLMSpec+S 10.38 ± 0.28 10.19 ± 0.19 9.95 ± 0.04 10.18 ± 0.08 RaLMSpec+A 10.88 ± 0.26 10.66 ± 0.12 10.56 ± 0.20 10.16 ± 0.18 RaLMSpec+P(20)SA 8.28 ± 0.18 8.18 ± 0.10 8.26 ± 0.14 8.09 ± 0.18 RaLMSpec+P(256)SA 9.46 ± 0.17 10.42 ± 0.74 9.64 ± 0.08 9.36 ± 0.16 · arXiv

Cited by

Sorted by citations

Related

Authors
Zhang, Z., Zhu, A., Yang, L., Xu, Y., Li, L., Phothilimthana, P. M., & Jia, Z.
Affiliations
  • Carnegie Mellon University (US)
  • University of California, Berkeley (US)
  • Google DeepMind (United Kingdom) (GB)
ORCID
Venue
arXiv
Publisher
arXiv
Published
2024
ISSN
2331-8422 (ISSN-L)
DOI
10.48550/arxiv.2401.14021
arXiv
2401.14021
OpenAlex
W4391272538
Cited by
1
References
0
Citations by year
2024: 1
Type
posted content
Language
en
Open access
Licence
Locations
  • arXiv (Cornell University) (submittedVersion, cc-by-sa)
  • arXiv (Cornell University)
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Topic Modeling
  • Physical Sciences › Computer Science › Information Systems › Information Retrieval and Search Behavior
  • Physical Sciences › Computer Science › Artificial Intelligence › Natural Language Processing Techniques
Keywords
retrieval-augmented language models, prefetching, sparse retrieval, question answering

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review