arXiv 2401.140212024-01-25cs.LGcs.CLcs.IR

用推测加速检索增强语言模型服务

Accelerating Retrieval-Augmented Language Model Serving with Speculation

Carnegie Mellon University（卡内基梅隆大学）；University of California, Berkeley（加州大学伯克利分校）；Google DeepMind（谷歌DeepMind）

Zhihao Zhang, Alan Zhu, Lijie Yang, Yihua Xu, Lanting Li, Phitchaya Mangpo Phothilimthana, Zhihao Jia

备注：: Preprint

阅读 PDF

arXiv 页面 ↗HTML 原文 ↗

AI总结：

针对迭代式RaLM频繁检索导致的高延迟问题，提出RaLMSpec框架，通过推测检索、批量验证、预取、调度和异步验证在保持输出不变的同时实现显著服务加速。

摘要：

检索增强语言模型（RaLM）通过将非参数化知识库与参数化语言模型相结合，已展现出解决知识密集型自然语言处理（NLP）任务的潜力。与微调完全参数化模型不同，RaLM擅长以低成本适配最新数据，并具备更好的来源归因机制。在各类RaLM方法中，迭代式RaLM由于检索器与语言模型之间更频繁的交互，可提供更好的生成质量。尽管有这些优势，迭代式RaLM通常会因频繁的检索步骤而产生高昂开销。为此，我们提出RaLMSpec，这是一个受推测启发的框架，通过推测检索和批量验证，在保持相同模型输出的同时，为迭代式RaLM提供通用加速。通过进一步引入预取、最优推测步长调度器和异步验证，RaLMSpec能够自动、最充分地挖掘加速潜力。对于朴素迭代式RaLM服务，在四个下游QA数据集上针对三种语言模型进行的大量评估表明，与基线相比，当检索器分别为精确稠密检索器、近似稠密检索器和稀疏检索器时，RaLMSpec可实现1.75–2.39倍、1.04–1.39倍和1.31–1.77倍的加速比。对于KNN-LM服务，与基线相比，当检索器为精确稠密检索器和近似稠密检索器时，RaLMSpec分别可实现高达7.59倍和2.45倍的加速比。

英文摘要：

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.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。