arXiv:2503.03434cs.CLcs.AI2025-03ACL被引 5

用检索增强生成草稿,提升大模型推理速度与泛化能力。

RASD: Retrieval-Augmented Speculative Decoding

  • 通过检索构建最优树结构,替代小模型生成草稿。
  • 在多任务上实现当前最优推理加速,最长接受长度显著提升。
  • 兼容多种生成与检索方法,适合需要高效推理的部署场景。

推测解码通过生成草稿令牌供目标模型验证来加速大语言模型的推理。现有方法依赖轻量级草稿模型或额外结构生成草稿令牌,并从数据库中检索上下文。由于草稿模型规模小且训练数据有限,其在域外场景下性能下降明显。此外,草稿阶段的时间开销限制了验证阶段的最大接受长度,影响整体效率。本文提出RASD(检索增强推测解码),采用检索方法增强基于模型的推测解码。我们引入树剪枝与树融合机制:首先基于草稿模型的概率分布设计剪枝方法,构建最优检索树;其次使用最长前缀匹配算法将草稿模型生成的树与检索树合并,形成统一验证树。实验表明,RASD在DocQA、Summary、Code和域内QA等任务上达到当前最优推理加速效果。同时,RASD具备强可扩展性,可无缝集成于各类推测解码方法,包括生成式与检索式方法。

原文摘要 · Abstract (English)

Speculative decoding accelerates inference in large language models (LLMs) by generating draft tokens for target model verification. Current approaches for obtaining draft tokens rely on lightweight draft models or additional model structures to generate draft tokens and retrieve context from databases. Due to the draft model's small size and limited training data, model-based speculative decoding frequently becomes less effective in out-of-domain scenarios. Additionally, the time cost of the drafting phase results in a low upper limit on acceptance length during the verification step, limiting overall efficiency. This paper proposes RASD (Retrieval-Augmented Speculative Decoding), which adopts retrieval methods to enhance model-based speculative decoding. We introduce tree pruning and tree fusion to achieve this. Specifically, we develop a pruning method based on the draft model's probability distribution to construct the optimal retrieval tree. Second, we employ the longest prefix matching algorithm to merge the tree generated by the draft model with the retrieval tree, resulting in a unified tree for verification. Experimental results demonstrate that RASD achieves state-of-the-art inference acceleration across tasks such as DocQA, Summary, Code, and In-Domain QA. Moreover, RASD exhibits strong scalability, seamlessly integrating with various speculative decoding approaches, including both generation-based and retrieval-based methods.

推理加速检索增强推测解码

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。