加速知识检索：伯克利&DeepMind联合研究，RaLMSpec让语言模型服务飞速提升2-7倍

近年来，随着大型语言模型（LLM）的出现，在多样化的 NLP 任务上取得了令人瞩目的成果。然而，知识密集型任务仍是 NLP 领域中的一项挑战，因为这些任务不仅要求模型要理解和生成自然语言，还要能够访问和利用大量的外部知识。然而，要想将大量知识编码进一个完全参数化的模型中，不仅在训练上需要更多努力，在部署时也同样如此。特别是当基础模型需要适应新数据或不同的下游任务时，这一挑战会更加严峻。

为了应对这一挑战，最近的研究提出了检索增强型语言模型（RaLM），通过增强检索将参数化的语言模型与非参数化的知识库结合起来，这种方法在低成本适应最新数据和更好的源归因机制方面表现出色。在各种 RaLM 方法中，由于在检索器和语言模型之间更频繁的交互，迭代 RaLM 提供了更好的生成质量。然而，迭代 RaLM 也会因为频繁的检索步骤而遇到高昂开销。

▲图1 (a)迭代式 RaLM 的工作流程 (b)RaLMSpec 框架 (c)RaLMSpec 在保持模型质量的同时实现了较低的延迟

由此，本文提出了一种名为 RaLMSpec 的框架，它采用了推测性检索和批量验证的方法，以减少迭代 RaLM 的服务开销，同时保证模型输出的一致性。通过进一步结合预提取、最优推测步长调度器和异步验证，RaLMSpec 能够自动充分利用加速潜力。在对三种语言模型在四个下游 QA 数据集上的评估中，RaLMSpec 展现出了作为通用加速框架在服务迭代 RaLM 方面的有效性。

论文题目:

Accelerating Retrieval-Augmented Language Model Serving with Speculation

论文链接:

https://arxiv.org/abs/2401.14021

声明：本期论文解读非人类撰写，全文由赛博马良「AI论文解读达人」智能体自主完成，经人工审核、配图后发布。

公众号「夕小瑶科技说」后台回复“智能体内测”获取智能体内测邀请链接。

---

所以，这便提出了一个问题：我们能否在不影响生成质量的情况下降低迭代 RaLM 的开销？

RaLMSpec 框架

由于生成性语言模型的自回归特性，现有的迭代 RaLM 方法检索步骤通常是用单个查询来概括当前上下文执行的。RaLMSpec 通过推测性检索（speculative retrieval）和批量验证（batched verification）来减少迭代 RaLM 的服务开销，同时保证模型输出的一致性。此外，RaLMSpec 还通过本地缓存、预提取、最优推测步长调度器和异步验证等技术进一步减少 RaLM 服务延迟。

▲算法1 R阿LMSpec 工作流

RaLMSpec 的核心设计包括推测性检索和批量验证。

• 推测性检索：利用本地缓存来存储检索文档，而不是每次都从知识库中检索。

• 批量验证：则在一系列推测性检索步骤之后进行，以确保推测结果的正确性。

• 如果在验证步骤中发现推测结果与真实检索结果不符，RaLMSpec 会回滚到第一个推测错误的位置，并使用正确的文档重新进行语言模型解码。

推测性检索

RaLMSpec 的核心之一是本地缓存与推测性检索技术。这一技术基于计算机架构中的推测执行概念，即通过在本地缓存中存储过去检索的文档来推测性地检索，而非每次都直接从知识库中检索。这种方法利用了检索文档的时空局部性，即在生成过程中可能会多次检索相同或连续的文档。

RaLMSpec 用更高效但准确度较低的推测性检索步骤，替代了现有 RaLM 方法中昂贵的迭代检索步骤。随后，RaLMSpec 使用批量验证步骤来纠正任何不正确的推测结果，从而保持模型的生成质量。

▲图2 推测式检索

1. 维护本地缓存：针对每个请求会创建并维护一个本地缓存，用于存储与该请求相关的文档条目。在这个本地缓存中，算法使用与原始检索器相同的评分指标对文档进行排序，以便在后续的处理中更有效地处理查询。

2. 验证步骤中的预提取：在验证步骤中，会从知识库中检索排名前 1 或前 k 的文档，并将它们填充到本地缓存中。这个步骤被称为预提取（prefetching），是为了提前准备一些可能相关的文档，以加速后续的查询处理。

批量验证

RaLMSpec 采用批量验证来纠正推测性检索的不正确结果，并保证模型输出的正确性。经过一系列推测性检索步骤后，RaLMSpec 通过执行批量检索来启动验证步骤，批量查询对应于推测性检索步骤中的提问。如果推测文档与验证步骤中检索到的真实文档不匹配，RaLMSpec 会自动纠正这种不匹配，通过回滚到第一个错误推测的位置并使用真实文档重新运行语言模型解码。这种批量检索通过并行处理提高了效率。

1. 隐藏验证延迟：异步验证通过将验证步骤延迟隐藏在推测步骤之后，实现了延迟的节省。这意味着在进行验证的同时，可以不需要等待验证步骤完成，就能继续进行下一个推测步骤。

2. 重新生成输出：如果检测到测错误，即推推测文档与真实文档不匹配，语言模型将使用真实文档重新生成输出。

如图 3 所示 ，RaLMSpec 支持异步验证，允许在进行验证的同时异步执行额外的推测步骤。这种技术在验证延迟小于语言模型解码延迟的情况下特别有益。此外，RaLMSpec引入了最优推测步长调度器（OS3），它通过将推测步长（即两个验证步骤之间连续推测步骤的数量）作为超参数，动态调整推测步长以最小化推测开销。

▲图3 异步验证

此外，RaLMSpec 还维护了一个本地检索缓存，用于存储每个提问的历史文档，并通过从本地缓存来执行推测性检索，而非在知识库中检索。通过这种方式，RaLMSpec 可以显著提高推测成功率。

实验结果

• 自然语言生成模型：实验中选用了 GPT-2、OPT 和 LLaMA-2 这三种。

• 开放领域问答数据集：Wiki-QA、Web Questions、Natural Questions 和 Trivia-QA，以及用于 KNN-LM 评估的 WikiText-103 数据集。

RaLMSpec 在不同检索器上的性能

通过在不同的语言模型、检索器和数据集上进行评估，RaLMSpec 展示了其在减少服务延迟方面的显著性能。在使用精确密集检索器时，RaLMSpec 能够实现最高 2.39 倍的加速比。即使在使用近似密集检索器和稀疏检索器时，RaLMSpec 也能实现显著加速。结果说明，RaLMSpec 在处理检索延迟方面特别有效，但其加速效果受限于检索延迟与端到端延迟的比例。

▲图4 RaLM 系列框架在 LLM 上的延迟比较

不同组件的贡献

表 1 的消融研究显示，启用最优推测步长调度器（）在三个组件中带来的性能提升最为显著，而将所有三个元素结合使用则一致实现了最佳性能。

这表明控制推测开销与延迟减少之间的权衡对于实现最佳加速效果至关重要。此外，预提取可以通过在本地缓存中缓存更多条目来提高推测准确率，但是增加预提取大小可能会引入更高的检索开销。异步验证通过在进行验证时引入额外的推测步骤来提高性能，但如果在早期阶段验证失败，则无法实现异步验证的好处。

▲表1 各组分加速比与基线的消融结果

不同推测步长的影响

实验结果表明，对于精确的密集检索器，较大的推测步长更为有利，因为检索延迟远远超过了语言模型生成延迟，因此推测步长会带来较小的开销，同时有机会获得更长的匹配。

相反，对于近似密集和稀疏检索器，较小的推测步长更为合适，因为检索更加高效，进行更多推测的成本效益比较低。因此，最优推测步长调度器（）旨在解决这种动态性，并在所有情况下实现接近最优的步长调度。

▲表2 不同预提取 size 加速比与基线的消融结果

KNN-LM 服务的 RaLMSpec 评估

KNN-LM 服务面临的主要挑战是高昂的推理开销，因为每一步生成 token 时都需要进行检索。RaLMSpec 通过修改缓存更新规则和验证协议，实现了显著的加速比。

结果表明，特别是在检索步骤频繁的情况下，RaLMSpec 能够有效地减少 KNN-LM 服务的延迟。

▲图5 使用 Wikipedia-QA 的 RaLMSpec 在 KNN-LM 上的加速比结果

结论

RaLMSpec 框架的提出，针对迭代型检索增强语言模型（iterative RaLM）在生成过程中频繁与知识库交互导致的高开销问题，提供了一种创新的解决方案。通过引入推测性检索和批量验证，RaLMSpec 能够在保证生成质量的前提下，显著减少服务延迟。

尽管 RaLMSpec 取得了显著成果，但仍有些潜在的研究方向值得进一步探索：

• 缓存策略的优化：当前的局部缓存策略虽然有效，但可能存在进一步优化的空间。例如，研究更智能的缓存更新机制，以适应不同的检索模式和数据分布。

• 异步验证的扩展：由于存在 Python 全局解释器锁（GIL），异步验证的潜力尚未被完全实现。未来可以探索在不同编程语言或运行时环境中实现异步验证，以充分利用并行处理能力。

• 跨模型和任务的泛化：虽然 RaLMSpec 在 QA 任务上表现出色，但其在其他知识密集型 NLP 任务中的效果尚未验证。

• 推测步长的动态调整：虽然已经提出了最优推测步长调度器，但如何根据实时反馈动态调整推测步长，以适应不断变化的数据和模型状态，仍是一个值得研究的问题。

• 未来可以探索将 RaLMSpec 与其他模型加速技术（如模型剪枝、量化等）结合使用，以进一步提高效率和降低资源消耗。

但总之，RaLMSpec 通过其创新的设计，为迭代型检索增强 LLM 的加速提供了一个有力的框架，能够在保持生成质量的同时显著减少服务延迟。未来的研究也可以在此基础上，继续探索更多的优化方向和应用场景。

声明：本期论文解读非人类撰写，全文由赛博马良「AI论文解读达人」智能体自主完成，经人工审核、配图后发布。

公众号「夕小瑶科技说」后台回复“智能体内测”获取智能体内测邀请链接。