Nearest Neighbor Speculative Decoding for LLM Generation and Attribution

📄 arXiv: 2405.19325v3 📥 PDF

作者: Minghan Li, Xilun Chen, Ari Holtzman, Beidi Chen, Jimmy Lin, Wen-tau Yih, Xi Victoria Lin

分类: cs.CL

发布日期: 2024-05-29 (更新: 2025-04-25)

期刊: Advances in Neural Information Processing Systems (2024), vol. 37, page 80987-81015

🔗 代码/项目: GITHUB

💡 一句话要点

提出近邻推测解码（NEST），提升LLM生成质量、溯源能力和推理速度。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 近邻检索 推测解码 半参数语言模型 LLM 溯源 知识密集型任务 生成质量 推理加速

📋 核心要点

  1. 现有LLM存在幻觉问题，且缺乏溯源能力，而kNN-LM等半参数模型推理速度慢，文本流畅度不足。
  2. NEST通过token级别检索，计算半参数混合分布，识别有希望的文本片段，并使用推测解码加速生成。
  3. 实验表明，NEST在知识密集型任务中显著提升了生成质量和溯源率，并加速了推理过程，达到1.8倍。

📝 摘要（中文）

大型语言模型（LLM）常常产生幻觉，并且缺乏对其生成内容进行溯源的能力。半参数语言模型，如kNN-LM，通过使用非参数数据存储中最近邻匹配来改进给定提示的LM输出，从而缓解这些限制。然而，这些模型通常推理速度慢，并且生成不流畅的文本。本文介绍了一种新颖的半参数语言建模方法——近邻推测解码（NEST），它能够将任意长度的真实文本片段整合到LM生成中，并提供对其来源的溯源。NEST在每个推理步骤执行token级别的检索，以计算半参数混合分布，并识别语料库中有希望的span延续。然后，它使用近似推测解码程序，接受检索到的span的前缀或生成新的token。NEST显著提高了基础LM在各种知识密集型任务中的生成质量和溯源率，超越了传统的kNN-LM方法，并且与上下文检索增强方法相比具有竞争力。此外，NEST显著提高了生成速度，应用于Llama-2-Chat 70B时，推理时间加速了1.8倍。代码将在https://github.com/facebookresearch/NEST/tree/main发布。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）生成内容时存在的幻觉问题，以及缺乏对其生成内容进行溯源的问题。现有的半参数模型，如kNN-LM，虽然可以缓解这些问题，但存在推理速度慢和生成文本不流畅的痛点。

核心思路：论文的核心思路是利用近邻检索和推测解码相结合的方式，在生成过程中引入外部知识，并加速推理过程。通过在每一步生成时检索语料库中最相关的文本片段，并利用这些片段进行推测解码，从而提高生成质量和溯源能力，同时提升推理速度。

技术框架：NEST的主要流程如下：1) Token级别检索：在每个token生成步骤，使用当前上下文查询非参数数据存储，检索最相关的文本片段。2) 半参数混合分布计算：将检索到的文本片段与LM的输出进行混合，得到一个半参数混合分布。3) 推测解码：使用推测解码算法，尝试接受检索到的文本片段的前缀作为生成结果，或者生成新的token。如果接受了检索到的片段，则可以实现溯源。

关键创新：NEST的关键创新在于将近邻检索与推测解码相结合，从而在提高生成质量和溯源能力的同时，显著提升了推理速度。与传统的kNN-LM方法相比，NEST能够处理任意长度的文本片段，并且通过推测解码加速了生成过程。与上下文检索增强方法相比，NEST避免了对整个上下文进行编码，从而降低了计算复杂度。

关键设计：NEST的关键设计包括：1) 检索策略：如何有效地检索到相关的文本片段？可能涉及到使用高效的索引结构和相似度度量方法。2) 混合权重：如何合理地混合检索到的文本片段和LM的输出？可能需要根据检索结果的质量和LM的置信度来动态调整混合权重。3) 推测解码策略：如何有效地利用检索到的文本片段进行推测解码？可能需要设计合适的接受准则和回退机制。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

NEST在知识密集型任务中显著提高了生成质量和溯源率，超越了传统的kNN-LM方法，并且与上下文检索增强方法相比具有竞争力。应用于Llama-2-Chat 70B时，推理时间加速了1.8倍，表明NEST在实际应用中具有显著的性能优势。

🎯 应用场景

NEST可应用于各种需要高质量、可溯源的文本生成场景，例如：知识密集型问答、内容创作、代码生成等。通过提供生成内容的来源，NEST可以提高用户对LLM输出的信任度，并减少幻觉的产生。此外，NEST的加速效果使其更适用于对延迟敏感的应用场景。

📄 摘要（原文）

Large language models (LLMs) often hallucinate and lack the ability to provide attribution for their generations. Semi-parametric LMs, such as kNN-LM, approach these limitations by refining the output of an LM for a given prompt using its nearest neighbor matches in a non-parametric data store. However, these models often exhibit slow inference speeds and produce non-fluent texts. In this paper, we introduce Nearest Neighbor Speculative Decoding (NEST), a novel semi-parametric language modeling approach that is capable of incorporating real-world text spans of arbitrary length into the LM generations and providing attribution to their sources. NEST performs token-level retrieval at each inference step to compute a semi-parametric mixture distribution and identify promising span continuations in a corpus. It then uses an approximate speculative decoding procedure that accepts a prefix of the retrieved span or generates a new token. NEST significantly enhances the generation quality and attribution rate of the base LM across a variety of knowledge-intensive tasks, surpassing the conventional kNN-LM method and performing competitively with in-context retrieval augmentation. In addition, NEST substantially improves the generation speed, achieving a 1.8x speedup in inference time when applied to Llama-2-Chat 70B. Code will be released at https://github.com/facebookresearch/NEST/tree/main.