DReSD: Dense Retrieval for Speculative Decoding

📄 arXiv: 2502.15572v2 📥 PDF

作者: Milan Gritta, Huiyin Xue, Gerasimos Lampouras

分类: cs.CL

发布日期: 2025-02-21 (更新: 2025-05-29)

备注: ACL (Findings) 2025

💡 一句话要点

DReSD：用于推测解码的稠密检索，显著提升LLM生成速度与质量

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

关键词: 推测解码 稠密检索 大型语言模型 近似最近邻搜索 语义检索

📋 核心要点

  1. 现有基于稀疏检索的推测解码方法依赖于精确字符串匹配和短上下文，限制了其性能。
  2. DReSD利用稠密检索，通过上下文token嵌入的近似最近邻搜索，寻找语义相关的token序列。
  3. 实验结果表明，DReSD在接受率、接受token长度和生成速度方面均优于稀疏检索方法。

📝 摘要（中文）

推测解码(SD)通过使用高效的草稿模型来预测接下来几个token，并由大型语言模型(LLM)在单个前向调用中验证这些token，从而加速LLM的生成过程，降低延迟并保持其输出质量。本文关注于基于检索的SD，其中草稿模型从非参数数据存储中检索下一个token。目前，稀疏检索(REST)是主要范式，因为它简单且可扩展，但其有效性受到短上下文和精确字符串匹配的限制。因此，我们提出了用于推测解码的稠密检索(DReSD)，它使用上下文token嵌入的近似最近邻搜索来检索语义上最相关的token序列用于SD。大量实验表明，与稀疏检索(REST)相比，DReSD的接受率平均提高87%，接受的token长度平均提高65%，生成速度平均提高19%。

🔬 方法详解

问题定义：现有的基于检索的推测解码方法，特别是使用稀疏检索（如REST）的方法，依赖于字符串的精确匹配，并且通常只使用较短的上下文。这导致无法有效利用语义信息，限制了草稿模型预测的准确性，从而降低了推测解码的效率。现有方法的痛点在于无法充分利用上下文的语义信息进行token检索。

核心思路：DReSD的核心思路是利用稠密向量表示来捕捉token的语义信息，并使用近似最近邻搜索（ANN）在向量空间中寻找与当前上下文最相关的token序列。通过将token嵌入到高维空间中，DReSD能够克服稀疏检索中精确匹配的限制，从而提高草稿模型预测的准确性。

技术框架：DReSD的整体框架包括以下几个主要模块：1) 上下文编码器：将输入上下文编码为稠密向量表示。2) 数据存储：存储预先计算好的token序列的稠密向量表示。3) 近似最近邻搜索：使用编码后的上下文向量在数据存储中进行ANN搜索，找到最相关的token序列。4) 验证模块：使用大型语言模型（LLM）验证检索到的token序列，并接受或拒绝它们。

关键创新：DReSD最重要的技术创新点在于使用稠密检索来替代传统的稀疏检索。与稀疏检索相比，稠密检索能够更好地捕捉token之间的语义关系，从而提高检索的准确性。此外，DReSD还使用了上下文编码器来更好地表示输入上下文，从而进一步提高了检索的性能。

关键设计：DReSD的关键设计包括：1) 上下文编码器的选择：可以使用预训练的语言模型（如BERT、RoBERTa）作为上下文编码器。2) 向量索引的选择：可以使用诸如FAISS、Annoy等高效的ANN索引库。3) 相似度度量：可以使用余弦相似度等度量来衡量向量之间的相似度。4) 数据存储的构建：需要预先计算并存储大量token序列的稠密向量表示。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DReSD在实验中表现出色，与稀疏检索(REST)相比，DReSD的接受率平均提高87%，接受的token长度平均提高65%，生成速度平均提高19%。这些结果表明，DReSD能够显著提高推测解码的效率，从而加速LLM的生成过程。

🎯 应用场景

DReSD可应用于各种需要快速生成文本的场景，例如聊天机器人、机器翻译、代码生成等。通过提高LLM的生成速度，DReSD可以显著提升用户体验，并降低计算成本。未来，DReSD可以进一步扩展到其他模态，例如图像和音频，从而实现更高效的多模态内容生成。

📄 摘要（原文）

Speculative decoding (SD) accelerates Large Language Model (LLM) generation by using an efficient draft model to propose the next few tokens, which are verified by the LLM in a single forward call, reducing latency while preserving its outputs. We focus on retrieval-based SD where the draft model retrieves the next tokens from a non-parametric datastore. Sparse retrieval (REST), which operates on the surface form of strings, is currently the dominant paradigm due to its simplicity and scalability. However, its effectiveness is limited due to the usage of short contexts and exact string matching. Instead, we introduce Dense Retrieval for Speculative Decoding (DReSD), a novel framework that uses approximate nearest neighbour search with contextualised token embeddings to retrieve the most semantically relevant token sequences for SD. Extensive experiments show that DReSD achieves (on average) 87% higher acceptance rates, 65% longer accepted tokens and 19% faster generation speeds compared to sparse retrieval (REST).