1. 论文基本信息

1.1. 标题

利用推测加速检索增强型语言模型服务 (Accelerating Retrieval-Augmented Language Model Serving with Speculation)

1.2. 作者

Zhihao Zhang1†, Alan Zhu1, Lijie Yang1, Yihua Xu2, Lanting Li1, Phitchaya Mangpo Phothilimthana3, Zhihao Jia1‡

隶属机构:

1. 卡内基梅隆大学 (Carnegie Mellon University), 计算机科学学院 (School of Computer Science)
2. 加州大学伯克利分校 (University of California, Berkeley)
3. 谷歌 DeepMind (Google DeepMind)

1.3. 发表期刊/会议

预印本 (Preprint)，发布于 arXiv。

1.4. 发表年份

2024年1月25日

1.5. 摘要

检索增强型语言模型 (Retrieval-augmented Language Models, RaLM) 通过结合非参数知识库和参数化语言模型，展现了解决知识密集型自然语言处理 (Natural Language Processing, NLP) 任务的潜力。相较于微调 (fine-tuning) 完全参数化的模型，RaLM 在适应最新数据方面成本更低，并提供了更好的来源归因机制。在各种 RaLM 方法中，迭代式 RaLM 因其检索器 (retriever) 和语言模型之间更频繁的交互，提供了更好的生成质量。然而，迭代式 RaLM 通常由于频繁的检索步骤而面临高昂的开销 (overheads)。为此，本文提出了 RaLMSpec，一个受推测 (speculation) 启发而设计的框架，通过推测性检索 (speculative retrieval) 和批量验证 (batched verification)，为迭代式 RaLM 提供通用的加速，同时保持相同的模型输出。通过进一步结合预取 (prefetching)、最优推测步幅调度器 (optimal speculation stride scheduler) 和异步验证 (asynchronous verification)，RaLMSpec 可以自动最大限度地利用加速潜力。对于朴素 (naive) 的迭代式 RaLM 服务，在三个语言模型和四个下游问答 (Question Answering, QA) 数据集上的广泛评估表明，与基线 (baseline) 相比，当检索器分别为精确密集检索器 (exact dense retriever)、近似密集检索器 (approximate dense retriever) 和稀疏检索器 (sparse retriever) 时，RaLMSpec 可以分别实现 1.75-2.39x、1.04-1.39x 和 1.31-1.77x 的加速比。对于 KNN-LM 服务，当检索器分别为精确密集检索器和近似密集检索器时，RaLMSpec 可以分别实现高达 7.59x 和 2.45x 的加速比。

1.6. 原文链接

原文链接: https://arxiv.org/abs/2401.14021 PDF 链接: https://arxiv.org/pdf/2401.14021v1.pdf

2. 整体概括

2.1. 研究背景与动机

论文试图解决的核心问题是： 迭代式检索增强型语言模型 (iterative RaLM) 虽然能提供高质量的生成结果，但由于频繁地与外部知识库进行检索交互，导致其服务延迟 (serving latency) 高昂。这限制了其在实际应用中的效率。

为什么这个问题在当前领域是重要的：

1. 大规模语言模型 (Large Language Models, LLMs) 的知识限制： 尽管 LLMs (如 LLaMA-2, GPT-3) 在许多 NLP 任务中表现出色，但将海量知识编码到完全参数化的模型中需要巨大的训练和部署成本。当模型需要适应新数据或特定下游任务时，这种成本会进一步增加。
2. RaLM 的优势： 检索增强型方法通过结合非参数知识库，能够以更低的成本适应最新数据，并提供更好的信息来源归因 (source attribution) 机制，这在需要最新、可验证知识的应用中非常重要。
3. 迭代式 RaLM 的质量优势： 相比于只进行一次检索的单次 RaLM (one-shot RaLM)，迭代式 RaLM 在生成过程中能更频繁地与知识库交互，根据当前的生成上下文检索最相关的文档，从而产生更高质量的响应。
4. 效率与质量的权衡： 现有迭代式 RaLM 方法在生成质量上优于单次 RaLM，但频繁的检索操作带来了显著的计算开销和延迟，使得其实际部署成本高昂。

现有研究存在的具体挑战或空白： 现有的迭代式 RaLM 服务方法通常是顺序执行检索和生成步骤，这种固有串行性导致效率低下。一些尝试提高效率的工作（如 Alon et al., 2022）可能会牺牲模型输出质量。因此，如何在不影响生成质量的前提下，显著降低迭代式 RaLM 的服务延迟，是一个尚未充分解决的挑战。

这篇论文的切入点或创新思路是： 本文借鉴了计算机体系结构中的“推测执行 (speculative execution)”思想，并将其应用于 RaLM 的检索过程。核心思路是，通过维护一个局部缓存 (local cache) 进行“推测性检索”，预测接下来的检索结果，并与语言模型进行推测性生成。然后，通过“批量验证”步骤，批量地从真实的外部知识库中检索，验证推测结果的正确性。如果推测错误，则回滚并使用真实结果重新生成。这种方式利用了检索结果的“时间/空间局部性 (temporal/spatial locality)”，并通过批量检索提高效率，同时保证了最终输出的正确性。

2.2. 核心贡献/主要发现

论文最主要的贡献是：

1. 提出了 RaLMSpec 框架： 这是一个通用的、受推测启发的框架，旨在减少迭代式 RaLM 方法的服务延迟，同时保证模型输出与朴素实现完全相同。
2. 引入缓存驱动的推测性检索和批量验证机制： RaLMSpec 利用检索文档的时间和空间局部性，通过一个请求级别的局部缓存进行推测性检索，并结合批量验证来减少频繁检索的开销，并确保最终结果的正确性。
3. 开发了多种优化技术： 为了进一步提升性能，RaLMSpec 额外引入了三项关键技术：
   • 缓存预取 (Cache Prefetching)： 通过预先加载更相关的文档到局部缓存，提高推测成功率。
   • 最优推测步幅调度器 (Optimal Speculation Stride Scheduler, $\mathrm { O S ^ { 3 } }$ )： 动态调整推测步幅 (speculation stride)，以自适应地优化推测开销和检索节省之间的权衡。
   • 异步验证 (Asynchronous Verification)： 允许在进行验证的同时异步执行额外的推测步骤，以隐藏验证延迟。
4. 广泛的实证验证： 在不同的任务（朴素迭代式 RaLM 和 KNN-LM 服务）、语言模型（GPT-2, OPT, LLaMA-2/13B）、检索器类型（精确密集、近似密集、稀疏检索器）和下游数据集上，验证了 RaLMSpec 的显著加速效果和通用性。

论文得出了哪些关键的结论或发现：

1. 显著的加速比： RaLMSpec 在朴素迭代式 RaLM 服务中实现了 1.04x 至 2.39x 的加速比，在检索密集型的 KNN-LM 服务中，加速比更高达 7.59x。
2. 对检索器类型的敏感性： RaLMSpec 在检索延迟较高的场景（如使用精确密集检索器）中，加速效果最为显著，因为其主要优化目标是减少检索开销。
3. 优化组件的协同作用： 结合预取、最优推测步幅调度器和异步验证，RaLMSpec 能够持续且自动地在所有场景下实现最佳性能。其中，最优推测步幅调度器对性能提升至关重要，因为它能动态适应不同场景下推测开销和收益的权衡。
4. 模型输出质量的保证： RaLMSpec 的设计确保了加速过程不会改变模型的最终输出，从而避免了对生成质量的任何妥协。

3. 预备知识与相关工作

3.1. 基础概念

为了理解 RaLMSpec，我们需要了解以下核心概念：

• 大型语言模型 (Large Language Models, LLMs):

  • 概念定义: LLMs 是参数量巨大、通常基于 Transformer 架构的深度学习模型，通过在海量文本数据上进行自监督预训练，学习语言的统计规律和世界知识。它们能够执行文本生成、问答、翻译等多种自然语言处理任务。
  • 与本文关系: LLMs 是 RaLM 的基础组成部分，负责根据上下文和检索到的信息生成文本。本文中使用的 GPT-2、OPT 和 LLaMA-2 都是典型的 LLMs。

• 检索增强型语言模型 (Retrieval-Augmented Language Models, RaLM):

  • 概念定义: RaLM 是一种结合了参数化语言模型和非参数外部知识库 (如维基百科) 的模型。当需要回答问题或生成文本时，RaLM 会先从知识库中检索相关的文档或信息，然后将这些检索到的信息作为额外上下文输入给语言模型，以辅助生成更准确、更具知识性的响应。
  • 优势: 相比于纯参数化模型，RaLM 能够访问最新信息，减少“幻觉 (hallucination)”，并提供更好的事实归因。
  • 与本文关系: RaLMSpec 正是为了加速 RaLM 的服务过程而设计的。

• RaLM 的分类:

  • 单次 RaLM (One-shot RaLM): 在语言模型开始生成之前，只执行一次检索操作。检索到的文档与初始请求结合，一次性提供给语言模型进行后续生成。
  • 迭代式 RaLM (Iterative RaLM): 在生成过程中，语言模型会根据当前的上下文 (context) 频繁地、迭代地查询知识库，以检索更多相关的文档来辅助后续的生成步骤。这种方法能够更灵活地适应生成过程中不断变化的知识需求，通常能带来更好的生成质量，但代价是更高的检索开销。
  • 与本文关系: RaLMSpec 专注于解决迭代式 RaLM 的高开销问题。

• 检索器 (Retrievers):

  • 概念定义: 检索器是 RaLM 中负责从知识库中找出与查询最相关文档的组件。
  • 分类:
    • 稀疏检索器 (Sparse Retriever, SR): 基于词汇匹配的方法，如 TF-IDF (Term Frequency-Inverse Document Frequency) 或 BM25 (Best Match 25)。它们通过计算查询和文档之间词汇重叠的统计量来评估相关性。
    • 密集检索器 (Dense Retriever): 基于深度学习模型，将查询和文档都编码成低维向量 (embedding)，然后在向量空间中通过计算向量距离 (如余弦相似度) 来找到最相似的文档。
      • 精确密集检索器 (Exact Dense Retriever, EDR): 在向量空间中执行精确的最近邻搜索，确保找到真正的最近邻。例如使用 FAISS 库进行全量搜索。
      • 近似密集检索器 (Approximate Dense Retriever, ADR): 为了提高查询速度，执行近似的最近邻搜索，允许牺牲一定的检索精度。例如使用 HNSW (Hierarchical Navigable Small World) 等索引结构。
  • 与本文关系: RaLMSpec 在不同类型的检索器（EDR, ADR, SR）上进行了实验，以展示其通用性。

• 推测执行 (Speculative Execution) / 推测解码 (Speculative Decoding):

  • 概念定义: 推测执行是计算机体系结构中的一种优化技术，处理器在不确定某个条件分支是否会被执行时，会根据预测结果提前执行指令。如果预测正确，则节省了等待条件确认的时间；如果预测错误，则需要回滚并重新执行正确的指令，但总的来说，期望收益是正的。推测解码是将这一思想应用于 LLM 的生成过程，使用一个更小、更快的“草稿模型 (draft model)”来推测性地生成多个词元，然后由一个更大、更精确的“目标模型 (target model)”批量验证这些词元。
  • 与本文关系: RaLMSpec 借鉴了推测执行的理念，但将其应用于检索步骤，而非语言模型解码本身。

• K-近邻语言模型 (K-Nearest Neighbor Language Models, KNN-LM):

  • 概念定义: KNN-LM 是一种特殊的 RaLM，它在生成每个新词元时，都会从一个外部知识库中检索 K 个最近邻的上下文及其对应的目标词元。然后，它将这些检索到的 K 个词元的分布与基础语言模型预测的词元分布进行插值 (interpolate)，以产生最终的下一个词元分布。这种方法可以显著提高基础语言模型的困惑度 (perplexity) 性能，但对每个词元都需要进行检索，因此计算开销非常大。
  • 与本文关系: KNN-LM 是一个典型的、检索密集型的迭代式 RaLM 示例，RaLMSpec 在其上进行了专门的评估，并取得了显著的加速效果。

3.2. 前人工作

本文将相关工作分为几个类别，并在此基础上突出 RaLMSpec 的创新点。

• 检索增强型语言模型 (Retrieval-augmented language models):

  • 单次 RaLM (One-shot RaLM): Guu et al. (2020) 首次提出了将检索到的文档提供给语言模型以增强生成质量。随后的工作 (Shi et al., 2023; Park et al., 2023; Wang et al., 2023a; Zhu et al., 2023; Rubin & Berant, 2023; Wang et al., 2023b; Zhou et al., 2023) 也主要采用在生成前只检索一次的方式。
  • 朴素迭代式 RaLM (Naive Iterative RaLM): 另一类工作在处理单个请求时定期执行检索。Ram et al. (2023)、Lewis et al. (2020)、Jiang et al. (2023)、Borgeaud et al. (2022) 和 Khattab et al. (2022) 持续地从外部数据库中检索最新的上下文信息，并通过直接拼接或中间层交叉注意力 (cross-attention) 来利用这些信息改进生成质量。这些方法被本文称为“朴素迭代式 RaLM”，它们提供了更高质量的响应，但代价是过高的延迟开销。
  • KNN-LM (K-Nearest Neighbour Language Models): Khandelwal et al. (2019) 和 Drozdov et al. (2022) 提出通过插值 K 个检索文档的加权分布和语言模型本身的输出来生成下一个词元分布。KNN-LM 是一种极端检索密集型的迭代式 RaLM，例如 Khandelwal et al. (2019) 为每个生成的词元检索多达 1024 个文档，导致实践中无法承受的服务开销。
  • 差异化分析: RaLMSpec 的核心目标正是减少这些迭代式 RaLM 方法（包括朴素迭代式 RaLM 和 KNN-LM）的服务开销，同时保持其高生成质量。

• RaLM 的检索器 (Retrievers for RaLM):

  • 早期的检索器包括传统的稀疏检索器如 TF-IDF (Ramos et al., 2003) 或 BM25 (Robertson et al., 2009)。
  • 随着深度学习发展，Karpukhin et al. (2020) 训练了专门用于 RaLM 的密集检索器 (Dense Passage Retriever, DPR)。Izacard et al. (2021) 通过无监督对比学习 (unsupervised contrastive learning) 训练密集检索器。后续工作进一步探索了端到端预训练 (end-to-end pretraining)，将检索器和语言模型协同训练 (Izacard et al., 2022; Zhong et al., 2022; Khattab & Zaharia, 2020; Santhanam et al., 2021)。
  • 检索器在“检索开销”和“检索精度”之间存在权衡：精确密集检索器 (EDR) 准确但效率低，而近似检索器 (ADR) 查询速度快但精度较低 (He et al., 2021; Xiong et al., 2020)。
  • 差异化分析: RaLMSpec 为了展示其设计的通用性，在稀疏检索器、精确密集检索器和近似密集检索器等不同类型的检索器上进行了实验。

• 迭代式 RaLM 服务 (Iterative RaLM serving):

  • 与本文最相关的工作是 Alon et al. (2022)。该工作通过在 KNN-LM 不需要完全检索时使用预计算的自动机状态，减少了对外部知识库的调用次数，从而节省了延迟。
  • 差异化分析: Alon et al. (2022) 不保证保留相同的模型输出，因此可能会损害模型生成质量。相比之下，RaLMSpec 的核心目标是在保证模型输出质量不变的前提下，实现通用且高效的迭代式 RaLM 服务。这是 RaLMSpec 与现有工作的关键区别。
  • 推测在 LLM 服务中的应用: 推测 (Speculation) 在计算机体系结构领域有悠久历史 (Burton, 1985)。最近，推测解码 (speculative decoding) 的概念被引入大型语言模型服务，以减少延迟 (Leviathan et al., 2022; Stern et al., 2018; Chen et al., 2023; Miao et al., 2023; Xia et al.; Joao Gante, 2023; Yang et al., 2023)。
  • 差异化分析: RaLMSpec 是第一项将“推测性检索 (speculative retrieval)”概念引入 RaLM 服务的工作，并且与用于 LLM 的“推测性推理/解码 (speculative inference/decoding)”技术是正交 (orthogonal) 的，这意味着它们可以独立应用或结合使用。

3.3. 技术演进

LLM 领域的发展经历了从纯参数化模型到结合外部知识库的 RaLM 的演变。最初的 RaLM 侧重于单次检索，以减少训练成本和提高知识的时效性。然而，随着任务复杂性增加，人们发现迭代式检索能更好地适应动态生成过程中的信息需求，从而带来质量提升。但这种质量提升是以高昂的检索延迟为代价的。为了解决这一效率瓶颈，研究者们开始探索各种服务优化技术，其中推测执行 (speculative execution) 作为一种通用的系统优化思想被引入。早期工作将其应用于 LLM 解码本身（推测解码），而本文则创新性地将其应用于 RaLM 的检索环节，以期在保持原始模型输出质量的同时，显著提高迭代式 RaLM 的服务效率。

3.4. 差异化分析

RaLMSpec 的主要创新和差异点在于：

1. 目标明确： 专注于加速迭代式 RaLM 的检索开销，而现有推测解码工作主要关注 LLM 解码开销。
2. 保证质量： 通过批量验证和回滚机制，RaLMSpec 严格保证加速后的模型输出与原始朴素实现完全一致，这区别于一些可能牺牲质量来换取速度的现有 RaLM 服务优化方案。
3. 推测对象不同： RaLMSpec 将推测思想应用于检索过程（即推测性检索），而非 LLM 的词元生成过程（推测性解码）。两者是互补且正交的。
4. 利用局部性： 核心是利用检索文档的时间/空间局部性，通过局部缓存进行快速推测，这是其性能提升的物理基础。
5. 自适应优化： 引入了最优推测步幅调度器，能够根据运行时特征（检索器类型、语言模型、推测准确率）动态调整策略，从而在不同场景下实现最佳性能。

4. 方法论

4.1. 方法原理

RaLMSpec 的核心思想是利用检索文档的时间局部性 (temporal locality) 和空间局部性 (spatial locality)。这意味着在迭代式检索过程中，同一个文档或相邻的文档很可能被重复检索。基于这一观察，RaLMSpec 提出了一种缓存驱动的推测性检索机制 (caching-based speculative retrieval mechanism)，并结合批量验证 (batched verification) 来降低迭代式 RaLM 的服务延迟，同时保证生成质量不变。

其直觉在于：

1. 检索成本高昂： 每次从外部知识库检索都涉及到查询编码、向量相似度搜索（对密集检索器）或索引查找（对稀疏检索器），以及数据传输，这些都是耗时的操作。
2. 推测性检索快： 如果能在局部缓存中快速预测出检索结果，就可以避免与外部知识库的昂贵交互。局部缓存的查找成本远低于外部知识库。
3. 批量操作高效： 验证阶段通过批量查询外部知识库，利用并行性，使得 N 次检索的成本远低于顺序执行 N 次单次检索的成本。
4. 正确性保证： 如果推测错误，系统会回滚到错误点，并使用真实检索结果重新执行，确保最终输出与非推测版本完全一致。

4.2. 核心方法详解 (逐层深入)

RaLMSpec 的工作流程可以分为推测性检索、批量验证以及进一步的优化组件。

4.2.1. 推测性检索 (Speculative Retrieval)

对于每个新的请求，RaLMSpec 会维护一个局部缓存 (local cache) 来存储已检索到的文档。这个局部缓存充当一个“检索”缓存，与传统系统中的精确匹配缓存有所不同，它会使用与原始检索器相同的评分机制对缓存内的条目进行排序。

1. 初始化与缓存预取 (Initialization and Cache Prefetching):

   • 当一个请求开始时，局部缓存为空。
   • 首先使用初始查询 $q$ 从外部知识库 ($\mathcal{C}$) 中进行检索，并用检索到的键值对填充局部缓存 $Q$。
   • 缓存预取 (Cache Prefetching) 是一种优化策略，通过在每个验证步骤中，不仅添加 top-1 文档，还添加来自知识库的 top-$k$ 文档（如 top-20 或 top-256）到局部缓存中。这旨在增加缓存命中率和推测成功率。

2. 推测过程:

   • 在后续的检索步骤中，RaLMSpec 不再直接查询外部知识库，而是推测性地从局部缓存 $Q$ 中进行检索。
   • 它使用与知识库检索相同的评分指标，在局部缓存内对条目进行排序，并取出最匹配的文档 $\hat{d}_i$。
   • 然后，语言模型 $f(\cdot)$ 会使用当前上下文 $X$ 和推测的文档 $\hat{d}_i$ 生成 $k$ 个新的词元 $\hat{X}_i$。
   • 生成的词元会添加到当前上下文 $X$ 中，形成新的上下文用于下一个推测步骤。
   • 这个推测性检索-生成循环会连续执行 $s$ 次，其中 $s$ 是推测步幅 (speculation stride)，表示在进行批量验证之前连续执行的推测步骤的数量。

   算法 1 (部分) - 推测性检索步骤: Algorithm 1 RaLMSpec Pipeline. Input: Input tokens $X = \{ x _ { 0 } , x _ { 1 } , \cdot \cdot \cdot , x _ { t - 1 } \}$ , external corpus $\mathcal { C }$ , language model $f ( \cdot )$ 2: Output: RaLM generated outputs 3: Initialize local cache $Q = \{ \}$ , speculation stride $s$ , model generation stride $k$ 4: $q = \operatorname { encode } ( X ) .$ . $Q$ insert(C.retrieve `( q )` ) $\triangleright$ cache prefetching 5: while EOS not in $X$ do 6: for $i = 1$ to $s$ do 7: $q _ { i } = \operatorname { encode } ( X ) , \hat { d } _ { i } = Q . \operatorname { retrieve } ( q _ { i } ) $ $\triangleright$ speculative retrieval 8: $\hat { X } _ { i } = f ( X , \hat { d } _ { i } , k ) $ $\triangleright$ model generation step that generates $k$ new tokens 9: $X = [ X , \hat { X } _ { i } ] $ 10: end for

   • 符号解释:

     • $X$: 当前的输入词元序列。

     • $x_0, \dots, x_{t-1}$: 序列中的具体词元。

     • $\mathcal{C}$: 外部知识库。

     • $f(\cdot)$: 语言模型，接受当前上下文 $X$、检索到的文档 $d_i$ 和生成词元数 $k$ 作为输入。

     • $Q$: 每个请求维护的局部缓存。

     • $q$: 编码后的查询向量。

     • $\operatorname{encode}(X)$: 将当前输入词元序列 $X$ 编码为查询 $q$。

     • $C.\operatorname{retrieve}(q)$: 从外部知识库 $\mathcal{C}$ 中检索与查询 $q$ 最相关的文档。

     • $s$: 推测步幅，即在验证前执行的推测步骤数量。

     • $k$: 语言模型一次生成的新词元数量。

     • $q_i$: 第 $i$ 次推测步骤中，根据当前上下文 $X$ 编码的查询。

     • $\hat{d}_i$: 第 $i$ 次推测步骤中，从局部缓存 $Q$ 中推测性检索到的文档。

     • $\hat{X}_i$: 第 $i$ 次推测步骤中，由语言模型 $f$ 使用上下文 $X$ 和推测文档 $\hat{d}_i$ 生成的 $k$ 个新词元。

     • $[X, \hat{X}_i]$: 将新生成的词元 $\hat{X}_i$ 拼接到当前上下文 $X$ 的末尾。

     • EOS: 结束生成 (End-Of-Sequence) 标记。

       关键特性: 对于大多数密集和稀疏检索器，如果知识库中的 top-ranked 文档存在于局部缓存中，那么从局部缓存检索时，该文档仍将位居榜首。这确保了局部缓存推测的有效性。

4.2.2. 批量验证与回滚 (Batched Verification and Rollback)

在连续进行 $s$ 次推测性检索和生成后，RaLMSpec 会进入一个验证步骤 (verification step)。

1. 批量检索 (Batched Retrieval):

   • 系统会收集在 $s$ 个推测步骤中产生的所有查询 $\{q_1, \dots, q_s\}$。
   • 然后，它会使用这些查询批量地从外部知识库 $\mathcal{C}$ 中执行检索，得到真实的文档结果 $\{d_1, \dots, d_s\}$。批量检索通过并行化提高了效率。

2. 验证与纠正 (Validation and Correction):

   • RaLMSpec 逐一比较推测的文档 $\hat{d}_i$ 和真实的文档 $d_i$。
   • 如果存在不匹配，则找到第一个不匹配的位置 $m = \arg \min_i \hat{d}_i \neq d_i$。
   • 如果发现不匹配 ($m \le s$)，生成过程将回滚 (roll back) 到第 $m$ 个推测步骤之前。
   • 然后，语言模型将使用正确的真实文档 $d_m$ 重新开始生成，并继续后续的生成过程。
   • 这种回滚机制确保了最终的模型输出与原始的非推测实现完全一致。

3. 局部缓存更新 (Local Cache Update):

   • 在验证步骤中，从知识库检索到的真实文档 $d_1, \dots, d_s$ 会被用来更新局部缓存 $Q$。
   • 除了更新 top-1 文档外，还可以选择进行预取 (prefetching)，即更新 top-$k$ 个文档到局部缓存中，以期提高未来的推测成功率。

   算法 1 (部分) - 批量验证步骤: Algorithm 1 RaLMSpec Pipeline. ... 10: end for 11: $d _ { 1 } , \cdot \cdot \cdot , d _ { s } = \mathcal { C } . \mathrm { retrieve } ( q _ { 1 } , \cdot \cdot \cdot , q _ { s } )$ $\triangleright$ batched verification 12: $m = \arg \operatorname* { m i n } _ { i } \hat { d } _ { i } \neq d _ { i } $ 13: if $m \le s$ then do correction if needed 14: Roll $X$ back to the $m$ -th speculation step 15: $\hat { X } = f ( X , d _ { i } , k ) $ 16: $X = [ X , \hat { X } ] $ 17: end if 18: end while

   • 符号解释:

     • $d_1, \dots, d_s$: 从外部知识库 $\mathcal{C}$ 中批量检索到的真实文档序列。

     • $\mathcal{C}.\mathrm{retrieve}(q_1, \dots, q_s)$: 使用批量查询 $\{q_1, \dots, q_s\}$ 从外部知识库 $\mathcal{C}$ 中检索对应的真实文档。

     • $m = \arg \min_i \hat{d}_i \neq d_i$: 找到第一个推测文档 $\hat{d}_i$ 与真实文档 $d_i$ 不匹配的位置索引。如果所有文档都匹配，则 $m > s$。

     • Roll $X$ back to the $m$-th speculation step: 将当前生成序列 $X$ 恢复到第 $m$ 个推测步骤之前的状态。

     • $\hat{X} = f(X, d_i, k)$: 语言模型 $f$ 使用回滚后的上下文 $X$ 和正确的真实文档 $d_i$ (这里应为 $d_m$) 重新生成词元 $\hat{X}$。

     • $X = [X, \hat{X}]$: 将重新生成的词元拼接到上下文 $X$ 中。

       该图像是一个示意图，展示了迭代 RaLM 和 RaLMSpec 框架的工作流程对比。图中可见，RaLMSpec 通过推测性检索和批量验证相结合，显著减少了知识检索的开销，提升了响应速度。

Figure 1 演示了迭代式 RaLM (a) 和 RaLMSpec (b) 的工作流程对比，以及 RaLMSpec 带来的延迟节省 (c)。

• 图 1(a) 显示了现有迭代式 RaLM 的流程：检索、生成、再检索、再生成，每个检索步骤都是串行的，导致高开销。

• 图 1(b) 概述了 RaLMSpec：它在多个快速的推测性检索步骤 (1, 3, 5) 之后，进行一个批量验证步骤 (6) 来保证正确性。推测性检索从局部缓存中进行，速度快。

• 图 1(c) 比较了 RaLMSpec 和现有方法的延迟，显示 RaLMSpec 显著降低了延迟。

  该图像是示意图，展示了RaLMSpec框架中的推测检索和缓存更新机制。上半部分展示了在推测检索过程中，语言模型（LM）如何使用局部缓存进行候选项检索，针对查询 $q_0$ 和 $q_1$ 的过程。下半部分说明了缓存更新策略，包括 Top-1 和 Top-k 更新方式，分别对应于模型验证和预取机制的实施。

Figure 2 详细说明了推测性检索和批量验证中局部缓存的使用。

• 局部缓存 $Q$ 存储了检索到的文档，并使用与原始检索器相同的评分机制。
• 推测性检索从 $Q$ 中进行，速度快。
• 在验证步骤中，$Q$ 会被从知识库中检索到的 top-1 或 top-$k$ 文档更新（后者即为预取）。

4.2.3. 异步验证 (Asynchronous Verification)

除了同步的推测和验证，RaLMSpec 还支持异步验证 (asynchronous verification)。

• 原理: 不在验证步骤期间暂停推测，而是允许在验证前一步进行时，异步地启动一个额外的推测步骤。

• 收益: 当验证延迟小于语言模型的解码延迟时，异步验证尤其有效。如果推测成功，模型的生成可以继续进行，验证的延迟实际上被隐藏在推测步骤的并行执行之后。

• 最佳策略: 在这种情况下，当推测步幅 $s=1$ 时，异步验证是推测性检索的最佳策略。

• 回滚处理: 如果异步验证失败，模型会根据纠正后的信息重新生成输出，这相当于回退到朴素实现，没有额外的推测开销。

  该图像是示意图，展示了RaLMSpec框架中验证线程与推测线程的工作流程。推测线程通过生成和验证文档来降低延迟，其中若生成的文档与预期不符，则需重新生成。该机制与单线程处理相比实现了更高的效率。

Figure 3 展示了异步验证如何通过在一个有效的推测步骤背后隐藏验证延迟来节省延迟。当推测文档与真实文档不匹配时，语言模型会使用真实文档重新生成输出。

4.2.4. 最优推测步幅调度器 (Optimal Speculation Stride Scheduler, $\mathrm { O S ^ { 3 } }$ )

推测步幅 $s$ 是一个关键超参数，它平衡了推测开销和延迟节省之间的权衡。过大的 $s$ 会导致推测失败时回滚开销过高，过小的 $s$ 则无法充分利用推测的优势。$\mathrm { O S ^ { 3 } }$ 旨在自适应地为不同配置找到最优的 $s$ 值。

1. 目标函数 (Objective Function): 目标是验证固定数量的文档，同时最小化延迟。因此，将目标设定为每单位时间成功验证的文档数量最大化。

2. 期望匹配长度 (Expected Matching Length):

   • 设 $\gamma(X) = P(d_i = \hat{d}_i \mid X)$ 为在给定当前上下文 $X$ 下，单步推测的准确率（即推测文档与真实文档匹配的概率）。
   • 那么，对于推测步幅 $s$，期望的匹配文档数量 (Expected # of verified document) 为： $$E[\text{verified documents} \mid X, s] = \frac { 1 - \gamma ( X ) ^ { s } } { 1 - \gamma ( X ) }$$
   • 符号解释:
     • $E[\text{verified documents} \mid X, s]$: 在给定上下文 $X$ 和推测步幅 $s$ 下，期望成功验证的文档数量。
     • $\gamma(X)$: 在给定上下文 $X$ 下，单步推测准确率，即推测文档 $\hat{d}_i$ 与真实文档 $d_i$ 匹配的概率。
     • $s$: 推测步幅。

3. 延迟计算 (Latency Calculation):

   • 设 $a$ 为一个推测步骤的延迟（推测性检索 + 语言模型解码），$b$ 为一个验证步骤的延迟。
   • 同步验证 (Synchronous Verification): 总延迟为 $s a + b$。
   • 异步验证 (Asynchronous Verification):
     • 如果所有推测文档都匹配（概率为 $\gamma(X)^s$），则受益于异步验证，延迟为 $(s - 1)a + \max(a, b)$。
     • 否则（概率为 $1 - \gamma(X)^s$），发生不匹配，异步验证没有收益，延迟为 $s a + b$。
     • 因此，期望延迟为： $$\text{Expected Latency} = [ \gamma ( X ) ^ { s } ( ( s - 1 ) a + \operatorname* { m a x } ( a , b ) ) + ( 1 - \gamma ( X ) ^ { s } ) ( s a + b ) ]$$
     • 符号解释:
       • $a$: 单次推测步骤的延迟 (推测性检索 + LM 解码)。
       • $b$: 单次验证步骤的延迟 (批量检索)。
       • $\max(a, b)$: $a$ 和 $b$ 中的较大值，表示异步隐藏延迟的上限。

4. 目标函数 (Objective Functions for $s$):

   • 同步验证下优化目标: $$\operatorname{Maximize} \quad \frac { 1 - \gamma ( X ) ^ { s } } { ( 1 - \gamma ( X ) ) ( s a + b ) }$$
   • 异步验证下优化目标: $$\operatorname{Maximize} \quad \frac { 1 - \gamma ( X ) ^ { s } } { ( 1 - \gamma ( X ) ) \left[ \gamma ( X ) ^ { s } \left( ( s - 1 ) a + \operatorname* { m a x } ( a , b ) \right) + \left( 1 - \gamma ( X ) ^ { s } \right) \left( s a + b \right) \right] }$$
   • 通过估计 $a, b, \gamma(X)$，即可求解最优的 $s$。

5. 参数估计 (Parameter Estimation for $\mathrm { O S ^ { 3 } }$):

   • a, b 的估计: 直接从最近步骤的性能分析结果中进行估计。对于精确密集检索器和稀疏检索器，小批量下的批量检索延迟几乎恒定；对于近似密集检索器，延迟与批量大小近似线性关系，但存在较大的截距，因此批量检索仍更高效。
   • $\gamma(X)$ 的估计: 使用特定窗口大小 $w$ 内的最大似然估计 (maximum log-likelihood estimation)，以兼顾局部性和减少方差。
     • 设 s(t) 是第 $t$ 个最近验证步骤的推测步幅 (也是批量大小)。
     • 设 M(s(t), X) 是相应的匹配文档数量。
     • 推测准确率 $\hat{\gamma}(X)$ 估计为： $$\hat { \gamma } ( X ) = \frac { \sum _ { t } M ( s ( t ) , X ) } { \sum _ { t } M ( s ( t ) , X ) + \sum _ { t } \mathbb { 1 } \left( M ( s ( t ) , X ) < s ( t ) \right) }$$
     • 符号解释:
       • $\hat{\gamma}(X)$: 对单步推测准确率的估计值。
       • $w$: 用于估计的窗口大小。
       • s(t): 第 $t$ 个最近验证步骤的推测步幅（批量大小）。
       • M(s(t), X): 第 $t$ 个最近验证步骤中匹配的文档数量。
       • $\mathbb{1}(\cdot)$: 指示函数 (indicator function)，如果括号内的条件为真则为 1，否则为 0。$\mathbb{1}(M(s(t), X) < s(t))$ 表示第 $t$ 个验证步骤中存在不匹配。
     • 为了防止过分乐观的估计和除零错误，对 $\hat{\gamma}$ 设置一个常数上限 $\gamma_{\max}$。

4.2.5. KNN-LM 的特殊处理

对于 KNN-LM 这种检索密集型工作负载，RaLMSpec 对缓存更新和验证协议进行了调整：

• 缓存填充: 无法简单地通过添加相同文档来填充缓存，因为未来解码步骤很可能不会再次检索到完全相同的文档。取而代之的是，利用观测到的空间局部性，用当前检索项之后的 $n$ 个相邻项来填充缓存（例如，实验中使用 $n=10$）。
• 推测成功定义: 将“成功的推测步骤”定义为推测的下一个词元与真实下一个词元匹配，而不是要求检索到的 K 个文档集合与真实 K 个文档集合完全匹配。当 K 值很大（例如 $K=1024$）时，匹配所有 K 个文档是指数级困难的，但匹配最终解码的词元更容易实现。这种放松的定义对于 KNN-LM 的推测成功至关重要。

5. 实验设置

5.1. 数据集

本文在两种主要的 RaLM 服务任务上进行了评估，使用了不同的数据集：

1. 朴素迭代式 RaLM 服务 (Naive iterative RaLM serving):

   • 任务类型: 知识密集型开放域问答 (knowledge-intensive open-domain question-answering tasks)。
   • 数据集:
     • Wiki-QA (Yang et al., 2015): 一个基于维基百科的问答数据集。
     • Web Questions (Berant et al., 2013): 包含从 Google 搜索日志中提取的真实世界用户问题。
     • Natural Questions (Kwiatkowski et al., 2019): 包含来自 Google 搜索的用户问题，以及维基百科文章中的长答案和短答案。
     • Trivia QA (Joshi et al., 2017): 一个大型的、远程监督的阅读理解挑战数据集，问题通常需要更深度的知识。
   • 知识库 (External Knowledge Base): 对于所有上述任务，均使用 维基百科语料库 (Wikipedia corpus) (Chen et al., 2017) 作为外部知识库。

2. KNN-LM 服务 (KNN-LM serving):

   • 任务类型: 语言建模，通常通过困惑度 (perplexity) 评估。
   • 数据集:
     • WikiText-103 (Merity et al., 2016): 一个大型的、高质量的维基百科文章数据集，常用于语言模型评估，也是原始 KNN-LM 工作 (Khandelwal et al., 2019) 中使用的数据集。

5.2. 评估指标

本文主要关注服务延迟 (serving latency)，并将其转化为加速比 (speed-up ratio) 进行衡量。

1. 延迟 (Latency):

   • 概念定义: 延迟是指从请求发出到接收到完整响应所需的时间。在 RaLM 服务中，它通常包括检索文档的时间、语言模型生成词元的时间以及两者之间协调的开销。
   • 度量单位: 墙钟时间 (wall-clock time)，以秒 (seconds) 为单位。
   • 特别说明: 由于 Python 全局解释器锁 (Global Interpreter Lock, GIL) 的存在，异步验证 (asynchronous verification) 的潜力未能完全实现，因此其延迟是基于理想运行时间进行模拟的，不包含额外的开销。
   • 与本文关系: 论文的核心目标是降低 RaLM 的服务延迟，因此延迟是主要的评估指标。

2. 加速比 (Speed-up Ratio):

   • 概念定义: 加速比衡量了优化方法相对于基线的性能提升程度。它定义为基线方法的延迟除以优化方法的延迟。
   • 数学公式: $$\text{Speed-up Ratio} = \frac{\text{Baseline Latency}}{\text{Optimized Method Latency}}$$
   • 符号解释:
     • $\text{Speed-up Ratio}$: 加速比，一个大于 1 的值表示优化方法更快。
     • $\text{Baseline Latency}$: 基线方法完成任务所需的总时间。
     • $\text{Optimized Method Latency}$: 优化方法（RaLMSpec 及其变体）完成任务所需的总时间。
   • 与本文关系: 论文使用加速比来量化 RaLMSpec 相较于基线方法所带来的性能提升。

3. 困惑度 (Perplexity, PPL):

   • 概念定义: 困惑度是评估语言模型性能的常用指标，衡量一个概率模型预测样本的能力。直观上，困惑度可以理解为模型对测试集平均每个词元的“不确定性”程度。困惑度越低，表示语言模型对文本序列的预测能力越强，模型性能越好。
   • 数学公式: 对于一个包含 $N$ 个词元（$w_1, w_2, \dots, w_N$）的测试序列 $W$，其困惑度 PPL(W) 的计算公式为： $$PPL(W) = P(w_1, w_2, \dots, w_N)^{-\frac{1}{N}} = \sqrt[N]{\frac{1}{P(w_1, w_2, \dots, w_N)}}$$ 其中，$P(w_1, w_2, \dots, w_N)$ 是模型对整个序列的联合概率，通常通过链式法则计算： $$P(w_1, w_2, \dots, w_N) = \prod_{i=1}^N P(w_i | w_1, \dots, w_{i-1})$$
   • 符号解释:
     • PPL(W): 文本序列 $W$ 的困惑度。
     • $W = (w_1, w_2, \dots, w_N)$: 包含 $N$ 个词元的文本序列。
     • $N$: 文本序列中的词元总数。
     • $P(w_1, w_2, \dots, w_N)$: 语言模型对整个序列 $W$ 的联合概率。
     • $P(w_i | w_1, \dots, w_{i-1})$: 语言模型在给定前 i-1 个词元时预测第 $i$ 个词元 $w_i$ 的概率。
   • 与本文关系: 虽然困惑度是 KNN-LM 工作中评估模型质量的关键指标，但 RaLMSpec 的设计保证了最终输出与基线模型完全相同，因此它关注的是保持困惑度不变的前提下降低延迟，而非直接优化困惑度。

5.3. 对比基线

本文根据不同的服务任务选择了不同的基线模型。

1. 朴素迭代式 RaLM 服务 (Naive Iterative RaLM Serving):

   • 基线名称: RaLMSeq
   • 实现细节: 直接遵循 Ram et al. (2023) 中的实现。在该实现中，语言模型每生成四个词元 (tokens) 就会触发一次检索。最新检索到的文档块会被直接预置到提示 (prompt) 中，并替换掉之前的文档。
   • 代表性: 代表了当前迭代式 RaLM 的典型朴素实现，频繁的串行检索是其主要性能瓶颈。

2. KNN-LM 服务 (KNN-LM Serving):

   • 基线名称: 未明确命名，但指代原始 KNN-LM 实现。
   • 实现细节: 遵循 Khandelwal et al. (2019) 中的实现。在该实现中，每生成一个词元 (token) 就会执行一次检索。
   • 代表性: KNN-LM 是一个极端检索密集型的工作负载，其设计使得检索开销非常高，是评估 RaLMSpec 在这种场景下加速能力的重要基线。

5.4. 语言模型

为了证明框架的通用性，本文选择了来自三种标准自然语言生成 (NLG) 模型类别的语言模型，并涵盖了不同的模型规模：

1. GPT-2 (Radford et al., 2019):
   • 具体模型: GPT2-medium
2. OPT (Zhang et al., 2022):
   • 具体模型: OPT-1.3B
3. LLaMA-2 (Touvron et al., 2023):
   • 具体模型: LLaMA-2-7B (用于主要实验) 和 LLaMA-2-13B (用于额外评估，以展示对更大模型的有效性)。
4. KNN-LM 的基础语言模型:
   • 具体模型: 一个 16 层的解码器专用 (decoder-only) Transformer 模型，具有 247M 可训练参数。这是 Khandelwal et al. (2019) 中使用的相同模型。

5.5. 检索器

为了展示方法的一致性，本文针对不同类型的检索器进行了测试，包括基于向量的密集检索器和基于词袋的稀疏检索器。

1. 精确密集检索器 (Exact Dense Retriever, EDR):
   • 具体模型: DPR (Dense Passage Retriever) (Karpukhin et al., 2020)
   • 实现: 基于标准的 FAISS 库 (Johnson et al., 2019)。
2. 近似密集检索器 (Approximate Dense Retriever, ADR):
   • 具体模型: DPR-HNSW (DPR with Hierarchical Navigable Small World) (Malkov & Yashunin, 2018)
   • 实现: 基于标准的 FAISS 库。
3. 稀疏检索器 (Sparse Retriever, SR):
   • 具体模型: BM25 (Best Match 25) (Robertson et al., 2009)
   • 实现: 使用 Pyserini 库 (Lin et al., 2021)。

5.6. 实现细节

• 最大输入提示长度 (Maximum input prompt length): 512 词元。
• 最大生成长度 (Maximum generation length): 128 词元。
• 检索文档块的最大长度 (Maximum length of retrieved document chunk): 对于朴素迭代式 RaLM 服务，设置为 256 (与 Ram et al., 2023 一致)。
• 推测步幅 (Speculation Stride, $s$):
  • 当最优推测步幅调度器 ($\mathrm { O S ^ { 3 } }$) 未启用时，RaLMSpec 使用固定步幅 $s=3$。
  • 当 $\mathrm { O S ^ { 3 } }$ 启用时，RaLMSpec 初始化步幅为 $s=1$，然后由调度器自适应调整。
• $\mathrm { O S ^ { 3 } }$ 参数:
  • 窗口大小 $w=5$。
  • 推测准确率上限 $\gamma_{\max} = 0.6$。
• 预取大小 (Prefetch Size):
  • 默认使用 20。
  • 在消融研究中也测试了 256。
• 异步验证 (Asynchronous Verification):
  • 由于 Python 全局解释器锁 (GIL) 的存在，异步验证的潜力未能完全实现。
  • 其延迟是模拟的，通过计算没有额外开销的理想运行时间。
  • 异步验证并非主要的加速来源，具体细节在消融研究中提供。
  • KNN-LM 服务中未启用异步验证，留作未来工作。
• 延迟测量: 除了异步验证是模拟的，所有其他延迟测量均采用墙钟时间 (wall-clock time)。
• 实现语言: 所有实现均使用 Python 编写。
• 计算资源: 在 Oracle 云的 VM.GPU.A10 实例上进行测试，该实例包含一个 A10 GPU 和 15 个 CPU。

6. 实验结果与分析

6.1. 核心结果分析

本节详细分析 RaLMSpec 在不同任务、模型和检索器类型下的实验结果。

6.1.1. 朴素迭代式 RaLM 服务 (Naive Iterative RaLM Serving)

实验在 GPT2-medium、OPT-1.3B 和 LLaMA-2-7B 三个语言模型上，针对 Wiki-QA、Web Questions、Natural Questions 和 Trivia-QA 四个 QA 数据集，以及精确密集检索器 (EDR)、近似密集检索器 (ADR) 和稀疏检索器 (SR) 三种不同类型的检索器进行了广泛评估。

该图像是一个示意图，展示了不同方法（RaLMSpec、RaLMSeq等）在多个数据集（GPT2、OPT、LLAMA-2等）上响应延迟（Latency）的比较。图中数据以条形图形式呈现，提供了各模型的延迟数据，以及不同检索模式下的性能表现。

Figure 4 展示了 RaLMSeq (基线)、RaLMSpec (基础版) 和 RaLMSpec+PSA (RaLMSpec 结合预取、最优推测步幅调度器和异步验证) 在不同配置下的延迟对比。总延迟被分解为语言模型生成延迟 (G) 和检索延迟 (R)，以直观展示 RaLMSpec 主要优化了检索部分。

• 精确密集检索器 (EDR) 场景：

  • RaLMSpec+PSA 实现了最显著的加速比。与 RaLMSeq 相比，对 GPT2 达到 2.39x，对 OPT 达到 2.33x，对 LLaMA-2 达到 1.75x。
  • 分析： EDR 的检索延迟 (R) 在总延迟中占据了非常大的比例。RaLMSpec 通过推测性检索和批量验证，大幅减少了对昂贵 EDR 的调用次数，因此加速效果最为明显。图中可以看出，RaLMSpec 显著压缩了蓝色部分的检索延迟。

• 近似密集检索器 (ADR) 场景：

  • RaLMSpec+PSA 实现了 1.05x (GPT2)、1.39x (OPT) 和 1.04x (LLaMA-2) 的加速比。
  • 分析： ADR 本身检索速度就快于 EDR，导致检索延迟 (R) 在总延迟中所占比例较小。因此，RaLMSpec 的加速空间相对有限。值得注意的是，基础版 RaLMSpec (没有 PSA 组件) 在某些情况下甚至比基线表现更差（加速比小于 1）。这主要是因为当检索延迟不高时，固定的推测步幅可能不是最优的，推测失败带来的回滚开销甚至会超过检索节省的收益。

• 稀疏检索器 (SR) 场景：

  • RaLMSpec+PSA 实现了 1.53x (GPT2)、1.77x (OPT) 和 1.31x (LLaMA-2) 的加速比。
  • 分析： 稀疏检索器通常介于 EDR 和 ADR 之间，其检索延迟也带来了一定的加速空间。与 ADR 类似，初始 RaLMSpec 版本在某些情况下表现不佳，凸显了自适应调度器的重要性。

总体分析： RaLMSpec+PSA 在所有场景下都持续实现了最佳性能。这表明其提出的三个优化组件（预取、最优推测步幅调度器、异步验证）协同工作，能够充分挖掘加速潜力。其中，加速比最显著的场景是检索延迟占主导地位时（如使用 EDR），因为 RaLMSpec 主要优化了检索环节。当语言模型生成延迟 (G) 占据主导时，检索延迟 (R) 的节省在整体延迟中的影响会被稀释。

6.1.2. 消融研究 (Ablation Study)

以下是原文 Table 1 的转录。

以下是原文 Table 1 的结果：

Table 1 展示了各个组件 (P: 预取, S: 最优推测步幅调度器, A: 异步验证) 对加速比的贡献。加速比是相对于基线 (RaLMSeq) 计算的，并对四个数据集取平均。

• 最优推测步幅调度器 (S) 的重要性： 在大多数情况下，启用 $\mathrm { O S ^ { 3 } }$ (RaLMSpec+S) 带来的收益最大。尤其是在 ADR 和 SR 场景中，基础版 RaLMSpec (固定步幅 $s=3$) 性能不佳（加速比小于 1），而 RaLMSpec+S 显著提升了性能，甚至超过了 1x。这表明对于检索延迟相对较低的场景，动态调整推测步幅以平衡开销和收益是至关重要的。
• 预取 (P) 的贡献： 预取 (RaLMSpec+P) 也能提高性能，主要是通过缓存更多条目来提高推测准确率。但在某些情况下，如果预取大小选择不当，其开销可能超过收益（参见 Table 2）。
• 异步验证 (A) 的贡献： 异步验证 (RaLMSpec+A) 也有助于性能提升，但通常不如 $\mathrm { O S ^ { 3 } }$ 或预取显著。它通过在验证期间引入额外的推测步骤来隐藏延迟。但如果验证过早失败，其优势就无法实现。
• 协同作用： 结合所有三个组件 (RaLMSpec+PSA) 几乎总能实现最佳性能。这表明这些组件相互补偿，共同发挥了 RaLMSpec 的最大潜力。

预取大小的消融研究： 以下是原文 Table 2 的转录。

以下是原文 Table 2 的结果：

Table 2 比较了不同预取大小 (20 vs. 256) 对加速比的影响。

• 对于 EDR，将预取大小从 20 增加到 256 略微提高了 GPT2 的性能，但降低了 OPT 和 LLaMA-2 的性能。
• 对于 ADR 和 SR，增加预取大小到 256 普遍导致性能下降。
• 分析： 这是因为增加预取大小会引入更高的检索开销。如果预取带来的推测成功率提升不足以抵消其自身的开销，整体性能就会下降。这再次强调了在不同场景下，超参数选择的重要性，以及自适应调度器的价值。

6.1.3. KNN-LM 服务 (KNN-LM Serving)

KNN-LM 是一种检索密集型工作负载，每次生成一个词元就需要进行检索。RaLMSpec 对其缓存更新规则和验证协议进行了修改（如将推测成功定义为下一个词元匹配而非 K 个文档匹配）。

该图像是一个柱状图，展示了不同最近邻数目（k）对应的加速比（Speed Up）。随着最近邻数目的增加，算法在不同配置下的性能表现也随之变化，最大加速比在 $k=1$ 时达 7.59，且不同参数设置下的结果明显不同。

Figure 5(a) 显示了使用精确密集检索器 (EDR) 的 KNN-LM 服务的加速比。

• RaLMSpec 实现了高达 7.59x 的加速比，尤其是在 $k=1$ 时。

• 即使在 $k=1024$ 这种极端密集检索的情况下，也能达到 3.88x 的加速比。

• 图中的不同线条代表不同的推测步幅设置（s=3, s=5, OS^3）。最优推测步幅调度器 ($\mathrm { O S ^ { 3 } }$) 几乎在所有 $k$ 值下都表现最佳。对于 EDR，通常较大的步幅能带来更好的性能。

  该图像是一个条形图，展示了在不同邻居数量（k）下，使用近似密集检索器时的加速比（Speed Up）。条形图中红色、蓝色和灰色分别表示不同的设置（s=3、s=5和OS^3），对应的加速比数值如2.39、2.45等展示了不同配置下的性能。在邻居数量k=1时，加速比最高，接近2.45。

Figure 5(b) 显示了使用近似密集检索器 (ADR) 的 KNN-LM 服务的加速比。

• RaLMSpec 实现了高达 2.45x 的加速比，同样在 $k=1$ 时表现最佳。

• 在 $k=1024$ 时，也能达到 2.37x 的加速比。

• 对于 ADR，最优推测步幅调度器 ($\mathrm { O S ^ { 3 } }$) 同样能持续实现最佳性能，尤其是在 $k$ 值较大时，其优势更加明显。与 EDR 不同，ADR 在不同 $k$ 值下对步幅的选择可能不同。

  分析： KNN-LM 的特性是极端频繁的检索（每个词元都检索），这使得检索延迟在总延迟中占据了极高的比例。RaLMSpec 通过其推测性批量检索机制，极大地减少了对外部知识库的调用次数，因此能够实现比朴素迭代式 RaLM 更高的加速比。

6.1.4. LLaMA-2-13B 服务评估

以下是原文 Table 3 的转录。

以下是原文 Table 3 的结果：

Table 3 展示了 RaLMSpec+PSA 在 LLaMA-2-13B 模型上，对四个下游数据集的加速比。

• 使用精确密集检索器 (EDR) 时，RaLMSpec+PSA 仍能实现高达 1.85x 的加速比。
• 对于近似密集检索器 (ADR)，改进幅度较小，约为 1.02x-1.04x。
• 分析： 即使对于更大的 LLaMA-2-13B 模型，RaLMSpec 仍然有效。对于 ADR，由于语言模型的生成延迟已经远远超过了检索延迟，因此检索延迟的节省在总端到端延迟节省中的比例变小，导致加速效果不那么显著。这再次印证了 RaLMSpec 的加速主要取决于检索开销在总开销中的占比。

6.2. 消融实验/参数分析 (Appendix 中的详细消融)

6.2.1. 不同系统组件组合的消融研究 (A.3 Ablation Study on Different System Components)

以下是原文 Table 4 的转录。

以下是原文 Table 4 的结果：

Table 4 展示了在 LLaMA-2-7B 和 Wiki QA 数据集上，不同组件组合（P: 预取, S: 最优推测步幅调度器, A: 异步验证）对平均服务延迟的影响（B: 基线 RaLMSeq）。

• EDR (精确密集检索器):

  • 在 EDR 场景下，检索延迟是主要瓶颈，因此预取 (P) 带来了显著改进 (从 144.39s 降至 82.23s)。
  • 令人意外的是，单独的预取 (P) 甚至略优于单独的 $\mathrm { O S ^ { 3 } }$ (S)。预取+异步验证 (PA) 甚至略优于 RaLMSpec+PSA。
  • 分析： 这可能是因为对于 EDR 这种检索延迟极高的场景，默认的固定步幅 $s=3$ 已经接近最优，并且 $\mathrm { O S ^ { 3 } }$ 需要一个“热身阶段”来适应最优步幅（从 $s=1$ 开始），从而引入了额外的开销。而预取在这种情况下能直接提高推测成功率，减少回滚，效果显著。

• ADR (近似密集检索器) 和 SR (稀疏检索器):

  • 对于 ADR 和 SR，$\mathrm { O S ^ { 3 } }$ (S) 成为了最重要的组件，它能将延迟从基线的 8.06s (ADR) / 10.75s (SR) 降低到 8.14s (ADR) / 10.38s (SR)，甚至在结合异步验证 (SA) 后，ADR 达到 7.83s，SR 达到 8.26s。
  • 分析： 对于 ADR 和 SR，检索本身效率较高，语言模型解码延迟占比更高。此时，不当的固定推测步幅可能导致过高的回滚开销，因此 $\mathrm { O S ^ { 3 } }$ 动态调整步幅的能力变得至关重要。预取 (P) 在这些场景下甚至可能带来负面影响（如 ADR 的 14.25s），这表明预取带来的收益可能无法覆盖其自身的检索开销。异步验证 (A) 也对性能有所贡献。

• 总结： 结合所有组件 (PSA) 通常能全面发挥 RaLMSpec 的潜力。$\mathrm { O S ^ { 3 } }$ 对于自适应地找到最优步幅至关重要，尤其是在检索开销不占主导地位的工作负载中。预取在检索开销大的场景下效果显著，而异步验证则提供额外的收益。

6.2.2. 推测步幅消融研究 (A.4 Ablation Study on Speculation Stride)

以下是原文 Table 5 的转录。

以下是原文 Table 5 的结果：

Table 5 展示了在 LLaMA-2-7B 和 Wiki QA 数据集上，不同固定推测步幅 ($s=2, 4, 8$) 与最优推测步幅调度器 ($\mathrm { O S ^ { 3 } }$) 的性能对比。

• EDR (精确密集检索器):
  • 对于 EDR，较大的推测步幅 ($s=4, 8$) 表现更好（81.06s, 81.90s），优于 $s=2$ (92.17s)。
  • 分析： 这是因为 EDR 的检索延迟远大于语言模型生成延迟。激进的推测步幅带来的回滚开销相对较小，而匹配更长的序列可以带来更大的检索节省。$\mathrm { O S ^ { 3 } }$ (85.19s) 略逊于固定 $s=4, 8$ 的情况，原因在于 $\mathrm { O S ^ { 3 } }$ 需要从 $s=1$ 开始进行“热身”和适应，这在初期会引入非最优步幅的开销。
• ADR (近似密集检索器) 和 SR (稀疏检索器):
  • 对于 ADR 和 SR，较小的推测步幅 ($s=2$) 表现最佳（ADR: 9.86s, SR: 10.65s）。
  • 随着步幅增大 ($s=4, 8$)，延迟显著增加（例如 ADR 在 $s=8$ 时达到 25.88s）。
  • 分析： ADR 和 SR 的检索效率较高，如果推测步幅过大，推测失败时的回滚开销会显著影响性能。在这种情况下，$\mathrm { O S ^ { 3 } }$ (ADR: 8.14s, SR: 10.38s) 能够动态选择更小的、更合适的步幅，因此表现优于所有固定步幅。
• 总结： 实验结果有力地证明了最优推测步幅调度器的重要性。它能够自适应地应对不同检索器类型和延迟特性，找到接近最优的推测策略，避免了手动调优的复杂性和次优性能。

6.3. 数据呈现 (表格)

由于附录中的表格（Table 6, 7, 8）结构复杂且数据量庞大，为遵循“完整转录”的指令并避免因结构复杂性导致渲染问题，我将为每张附录表格创建单独的 HTML 表格。

6.3.1. 表 6: 使用精确密集检索器 (EDR) 时 GPT-2、OPT 和 LLaMA-2 的平均延迟（秒）

以下是原文 Table 6 的结果：

6.3.2. 表 7: 使用近似密集检索器 (ADR) 时 GPT-2、OPT 和 LLaMA-2 的平均延迟（秒）

以下是原文 Table 7 的结果：

6.3.3. 表 8: 使用稀疏检索器 (SR) 时 GPT-2、OPT 和 LLaMA-2 的平均延迟（秒）

以下是原文 Table 8 的结果：

7. 总结与思考

7.1. 结论总结

本文提出了 RaLMSpec，一个创新性的框架，旨在加速迭代式检索增强型语言模型 (RaLM) 的服务过程。其核心思想是将计算机体系结构中的推测执行理念引入 RaLM 的检索环节。RaLMSpec 利用检索文档的时间和空间局部性，通过在请求级别维护一个局部缓存进行推测性检索。为了保证结果的正确性，RaLMSpec 引入了批量验证机制，并在推测失败时进行回滚，确保最终模型输出与原始朴素实现完全一致。

为了进一步优化性能，RaLMSpec 集成了三项关键技术：

1. 缓存预取 (Cache Prefetching)：通过预先加载更多相关文档来提高推测成功率。

2. 最优推测步幅调度器 (Optimal Speculation Stride Scheduler, $\mathrm { O S ^ { 3 } }$ )：动态调整推测步幅，以自适应地平衡推测开销和检索节省。

3. 异步验证 (Asynchronous Verification)：通过并行执行推测和验证来隐藏部分延迟。

   实验结果表明，RaLMSpec 在多种任务（朴素迭代式 RaLM 和 KNN-LM 服务）、语言模型（GPT-2, OPT, LLaMA-2/13B）、检索器类型（精确密集、近似密集、稀疏检索器）和下游数据集（Wiki-QA, Web Questions, Natural Questions, Trivia-QA, WikiText-103）上均取得了显著的加速效果。对于朴素迭代式 RaLM，加速比最高可达 2.39x；而对于检索更为密集型的 KNN-LM，加速比更是高达 7.59x。特别地，$\mathrm { O S ^ { 3 } }$ 被证明是至关重要的组件，它能够自适应地优化不同场景下的性能。

总而言之，RaLMSpec 为解决迭代式 RaLM 的高延迟问题提供了一个通用且高效的解决方案，在不牺牲模型生成质量的前提下，显著提升了服务效率。

7.2. 局限性与未来工作

论文作者指出了 RaLMSpec 的一些局限性：

1. 异步验证的实际实现限制： 由于 Python 全局解释器锁 (GIL) 的存在，异步验证的全部潜力未能完全实现。在实验中，异步验证的延迟是模拟的，而非真实的墙钟时间。这意味着其在实际 Python 环境中的真实收益可能低于模拟结果。

2. KNN-LM 异步验证的缺失： 在 KNN-LM 服务中，异步验证并未启用，这被留作未来的工作。

   基于这些局限性，论文提出了未来的研究方向：

3. 探索在没有 GIL 限制或通过其他语言/运行时（如 C++ 或 Rust，或使用 PyTorch/TensorFlow 的并行计算特性）实现异步验证，以充分发挥其潜力。

4. 将异步验证技术集成到 KNN-LM 服务中，以进一步提升其性能。

7.3. 个人启发与批判

个人启发：

1. 跨领域思想的迁移价值： RaLMSpec 再次证明了将计算机体系结构（如推测执行）的优化思想迁移到机器学习系统优化中的巨大潜力。这种思维方式能够从根本上识别和解决系统瓶颈。
2. 局部性原理的普适性： 检索文档的时间和空间局部性是 RaLMSpec 成功的基石。这提示我们在设计任何涉及数据访问的系统时，都应优先考虑和利用数据的局部性特征。
3. 批量操作的效率优势： 批量验证是 RaLMSpec 实现加速的关键。它强调了在现代硬件（尤其是 GPU）上利用并行性进行批量处理的效率优势，单次昂贵操作的批量化往往能带来超线性的性能提升。
4. 自适应调优的重要性： 最优推测步幅调度器 ($\mathrm { O S ^ { 3 } }$) 的引入，解决了固定超参数在不同场景下性能不稳定的问题。这启发我们，在复杂的 ML 系统中，自适应、动态的运行时调度器对于实现鲁棒且最优的性能至关重要。
5. 不牺牲质量的优化： RaLMSpec 严格保证了输出质量与原始模型一致，这在对准确性和可信度有高要求的知识密集型应用中非常关键。速度提升而不影响质量，是系统优化的黄金准则。

个人批判与可以改进的地方：

1. 异步验证的真实性能验证： 模拟异步验证性能是一个值得注意的局限。虽然作者解释了 GIL 的原因，但在实际部署中，异步验证的真实性能（包括其带来的额外复杂性和潜在的同步开销）需要更严谨的实验验证，例如在非 Python 环境或使用专门的并行计算库进行实现和测试。
2. $\mathrm { O S ^ { 3 } }$ 的开销和收敛性： 最优推测步幅调度器 ($\mathrm { O S ^ { 3 } }$) 需要在线估计 $a, b, \gamma(X)$ 等参数。这些估计本身的开销（例如，维护历史记录和计算统计量）以及在动态环境中的收敛速度和稳定性如何，是值得深入探讨的问题。如果估计不准确或收敛缓慢，可能导致次优的步幅选择，甚至引入额外的开销。
3. 局部缓存的管理策略： 论文中对局部缓存的维护机制描述相对简单。对于长时间运行的请求或大量并发请求，缓存的大小、替换策略（例如 LRU, LFU）、一致性模型等都会影响性能。随着上下文长度的增加，局部缓存可能变得庞大，其自身的查询和更新开销也可能变得不可忽略。
4. 推测失败的惩罚机制： 虽然论文讨论了回滚，但对于不同检索器类型和模型大小，推测失败带来的实际惩罚（例如，回滚和重新生成的时间成本）可能差异很大。更精细地建模和优化这种惩罚，可能会进一步提升 $\mathrm { O S ^ { 3 } }$ 的性能。
5. 更深层次的推测： RaLMSpec 专注于检索环节的推测。未来工作可以探索将推测性检索与现有的推测性解码技术（针对 LLM 解码）相结合，从而实现端到端的更大加速。这需要更复杂的调度和验证机制。
6. 能耗和资源利用率： 论文主要关注延迟。在实际部署中，能耗和资源利用率（如 GPU 内存、CPU 使用率）也是关键指标。例如，批量验证虽然减少了检索次数，但可能要求更大的批量大小，这会增加 GPU 内存压力。对这些方面的权衡进行分析将使研究更完善。