Speculative Decoding with CTC-based Draft Model for LLM Inference Acceleration

📄 arXiv: 2412.00061v1 📥 PDF

作者: Zhuofan Wen, Shangtong Gui, Yang Feng

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-11-25

💡 一句话要点

提出基于CTC的推测解码方法，提升LLM推理加速时的草稿接受率

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

关键词: 大型语言模型 LLM推理加速 推测解码 CTC 草稿模型

📋 核心要点

  1. 现有推测解码方法中，草稿模型忽略了草稿token间的关联，导致接受率不高，影响整体推理速度。
  2. 论文提出基于CTC的草稿模型，通过加强草稿token间的关联性，生成更高质量的草稿序列。
  3. 实验结果表明，该方法相比现有方法，显著提高了草稿的接受率，从而加速了LLM的推理过程。

📝 摘要（中文）

大型语言模型（LLM）的推理加速在许多应用场景中至关重要，而推测解码已显示出其在加速推理方面的优势。推测解码通常引入一个草稿模型来辅助基础LLM，草稿模型生成草稿，基础LLM验证草稿是否被接受。在这种框架下，最终的推理速度取决于草稿模型的解码速度和草稿模型提供的草稿的接受率。目前，广泛使用的草稿模型通常以非自回归方式生成接下来几个位置的草稿token，而不考虑草稿token之间的相关性。因此，它具有较高的解码速度，但接受率不尽如人意。在本文中，我们专注于如何提高草稿模型的性能，旨在通过高接受率来加速推理。为此，我们提出了一种基于CTC的草稿模型，该模型加强了草稿阶段草稿token之间的相关性，从而生成更高质量的草稿候选序列。实验结果表明，与强大的基线相比，该方法可以实现更高的接受率，从而实现更快的推理速度。

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLM）推理加速问题，特别是在推测解码框架下，草稿模型生成的草稿接受率不高的问题。现有草稿模型通常采用非自回归方式生成草稿token，忽略了token之间的相关性，导致生成的草稿质量不高，最终影响了整体的推理速度。

核心思路：论文的核心思路是利用CTC（Connectionist Temporal Classification）来建模草稿token之间的相关性，从而提高草稿模型的生成质量和草稿的接受率。通过加强token间的依赖关系，使得草稿模型能够生成更符合语言模型分布的草稿序列。

技术框架：整体框架仍然是推测解码，包含一个基础LLM和一个草稿模型。草稿模型基于CTC进行训练和推理，生成草稿序列。基础LLM负责验证草稿序列的正确性，并决定接受或拒绝。如果接受，则继续生成后续token；如果拒绝，则回退到基础LLM进行自回归生成。

关键创新：最重要的技术创新点在于将CTC引入到草稿模型的训练中。传统的草稿模型通常采用独立的token预测方式，而CTC能够建模序列之间的对齐关系，从而学习到token之间的依赖关系。这使得草稿模型能够生成更连贯、更符合语言规律的草稿序列，显著提高了草稿的接受率。与现有方法的本质区别在于，该方法考虑了草稿token之间的序列关系，而现有方法则忽略了这一点。

关键设计：论文中，CTC被用于训练草稿模型，目标是最大化观察到正确token序列的概率。具体来说，草稿模型输出的是每个token位置上的概率分布，CTC通过前向-后向算法计算所有可能的对齐路径的概率，并选择概率最高的路径作为最终的预测序列。关键参数包括CTC的权重、草稿模型的网络结构（例如Transformer）以及训练数据等。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，与现有推测解码方法相比，该方法显著提高了草稿的接受率，从而实现了更快的推理速度。具体的性能数据需要在论文中查找，但摘要中明确指出该方法优于强大的基线模型，证明了其有效性。

🎯 应用场景

该研究成果可广泛应用于需要加速LLM推理的场景，例如在线对话系统、机器翻译、文本摘要等。通过提高LLM的推理速度，可以降低计算成本，提升用户体验，并促进LLM在资源受限设备上的部署。未来，该方法有望进一步扩展到其他序列生成任务中。

📄 摘要（原文）

Inference acceleration of large language models (LLMs) has been put forward in many application scenarios and speculative decoding has shown its advantage in addressing inference acceleration. Speculative decoding usually introduces a draft model to assist the base LLM where the draft model produces drafts and the base LLM verifies the draft for acceptance or rejection. In this framework, the final inference speed is decided by the decoding speed of the draft model and the acceptance rate of the draft provided by the draft model. Currently the widely used draft models usually generate draft tokens for the next several positions in a non-autoregressive way without considering the correlations between draft tokens. Therefore, it has a high decoding speed but an unsatisfactory acceptance rate. In this paper, we focus on how to improve the performance of the draft model and aim to accelerate inference via a high acceptance rate. To this end, we propose a CTC-based draft model which strengthens the correlations between draft tokens during the draft phase, thereby generating higher-quality draft candidate sequences. Experiment results show that compared to strong baselines, the proposed method can achieve a higher acceptance rate and hence a faster inference speed.