Graph-Structured Speculative Decoding

📄 arXiv: 2407.16207v1 📥 PDF

作者: Zhuocheng Gong, Jiahao Liu, Ziyue Wang, Pengfei Wu, Jingang Wang, Xunliang Cai, Dongyan Zhao, Rui Yan

分类: cs.CL

发布日期: 2024-07-23


💡 一句话要点

提出图结构推测解码(GSD)加速LLM推理,显著提升token接受率和推理速度。

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

关键词: 推测解码 大型语言模型 推理加速 有向无环图 图结构 LLaMA-2 多假设生成

📋 核心要点

  1. 现有推测解码方法依赖于单个草稿假设,限制了token接受率和加速效果。
  2. GSD通过生成多个假设并利用有向无环图(DAG)结构来高效管理和合并这些假设,从而提高token接受率。
  3. 实验表明,GSD在包括LLaMA-2在内的多个LLM上实现了显著的加速,最高可达1.96倍。

📝 摘要(中文)

推测解码是一种通过利用小型语言模型生成假设序列,然后由大型语言模型(LLM)验证来加速LLM推理的有前景的技术。该方法的有效性很大程度上取决于草稿模型的性能和效率之间的平衡。本研究侧重于通过生成多个假设而不是一个假设来提高被最终输出接受的草稿token的比例。这使得LLM有更多选择,并选择满足其标准的最长序列。分析表明,草稿模型产生的假设共享许多共同的token序列,这表明存在优化计算的潜力。利用这一观察结果,我们引入了一种创新方法,利用有向无环图(DAG)来管理起草的假设。这种结构使我们能够有效地预测和合并重复出现的token序列,从而大大降低了草稿模型的计算需求。我们将这种方法称为图结构推测解码(GSD)。我们将GSD应用于包括700亿参数的LLaMA-2模型在内的一系列LLM,并观察到1.73倍至1.96倍的显著加速,显著超过了标准推测解码。

🔬 方法详解

问题定义:现有推测解码方法通常只生成一个草稿假设序列,如果这个序列与大型语言模型(LLM)的预测不一致,就会导致大量的token被拒绝,从而限制了加速效果。痛点在于如何提高草稿token的接受率,同时避免引入过多的计算开销。

核心思路:核心思路是生成多个草稿假设,而不是仅仅一个。这样,LLM在验证阶段就有更多的选择,可以选择最符合其标准的序列。为了避免多个假设带来的计算冗余,论文利用了这些假设之间存在大量重复token序列的特性,通过图结构来共享这些重复计算。

技术框架:GSD的核心是一个有向无环图(DAG),其中节点代表token,边代表token之间的转移概率。草稿模型生成多个假设序列,这些序列被映射到DAG上。DAG结构能够有效地合并重复的token序列,从而避免重复计算。在验证阶段,LLM从DAG中选择最长的、符合其标准的路径作为最终输出。整个流程包括草稿生成、DAG构建、LLM验证和路径选择四个主要阶段。

关键创新:最重要的技术创新点在于利用图结构来管理和合并多个草稿假设。与传统的推测解码方法相比,GSD能够更有效地利用草稿模型的计算结果,并显著提高token的接受率。本质区别在于,GSD不再是简单地验证单个假设,而是从多个假设中选择最优的组合。

关键设计:DAG的构建方式是关键。论文可能采用了某种启发式算法来确定哪些token序列应该被合并,以及如何有效地搜索DAG以找到最优路径。具体参数设置和损失函数(如果存在)未知。网络结构方面,草稿模型可以是任何小型语言模型,但其性能直接影响GSD的加速效果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,GSD在多个LLM上实现了显著的加速。例如,在700亿参数的LLaMA-2模型上,GSD实现了1.73倍至1.96倍的加速,显著优于标准推测解码方法。这些结果验证了GSD的有效性,并表明其具有广泛的应用潜力。

🎯 应用场景

GSD可广泛应用于各种需要加速LLM推理的场景,例如在线对话系统、文本生成、机器翻译等。通过提高推理速度,GSD可以降低计算成本,提升用户体验,并促进LLM在资源受限设备上的部署。未来,GSD可以与其他加速技术相结合,进一步提升LLM的推理效率。

📄 摘要(原文)

Speculative decoding has emerged as a promising technique to accelerate the inference of Large Language Models (LLMs) by employing a small language model to draft a hypothesis sequence, which is then validated by the LLM. The effectiveness of this approach heavily relies on the balance between performance and efficiency of the draft model. In our research, we focus on enhancing the proportion of draft tokens that are accepted to the final output by generating multiple hypotheses instead of just one. This allows the LLM more options to choose from and select the longest sequence that meets its standards. Our analysis reveals that hypotheses produced by the draft model share many common token sequences, suggesting a potential for optimizing computation. Leveraging this observation, we introduce an innovative approach utilizing a directed acyclic graph (DAG) to manage the drafted hypotheses. This structure enables us to efficiently predict and merge recurring token sequences, vastly reducing the computational demands of the draft model. We term this approach Graph-structured Speculative Decoding (GSD). We apply GSD across a range of LLMs, including a 70-billion parameter LLaMA-2 model, and observe a remarkable speedup of 1.73$\times$ to 1.96$\times$, significantly surpassing standard speculative decoding.