arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2503.03434cs.CLcs.AI
收藏

RASD：检索增强的投机解码

RASD: Retrieval-Augmented Speculative Decoding

Alibaba Cloud, Alibaba Group（阿里云，阿里巴巴集团）

Guofeng Quan, Wenfeng Feng, Chuzhan Hao, Guochao Jiang, Yuewei Zhang, Hao Wang

阅读 PDF
arXiv 页面 ↗HTML 原文 ↗

AI总结：

本文提出RASD框架，通过基于概率分布的树剪枝和最长前缀匹配的树融合技术，将检索树与草稿模型生成的树合并，有效解决了域外场景下投机解码效果下降和接受长度受限的问题，在多项任务上实现了最先进的推理加速。

摘要：

投机解码通过生成草稿token供目标模型验证，加速大语言模型（LLMs）的推理。当前获取草稿token的方法依赖轻量级草稿模型或额外模型结构，以生成草稿token并从数据库检索上下文。由于草稿模型规模小且训练数据有限，基于模型的投机解码在域外场景中经常效果下降。此外，草稿阶段的时间成本导致验证步骤中接受长度的上限较低，限制了整体效率。本文提出RASD（检索增强的投机解码），采用检索方法增强基于模型的投机解码。我们引入树剪枝和树融合来实现这一点。具体而言，我们开发了一种基于草稿模型概率分布的剪枝方法，以构建最优检索树。其次，我们采用最长前缀匹配算法，将草稿模型生成的树与检索树合并，形成用于验证的统一树。实验结果表明，RASD在DocQA、Summary、Code和In-Domain QA等任务上实现了最先进的推理加速。此外，RASD展现出强大的可扩展性，能无缝集成各种投机解码方法，包括基于生成和基于检索的方法。

英文摘要：

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.

机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。

↑