AdaEDL: Early Draft Stopping for Speculative Decoding of Large Language Models via an Entropy-based Lower Bound on Token Acceptance Probability

📄 arXiv: 2410.18351v1 📥 PDF

作者: Sudhanshu Agrawal, Wonseok Jeon, Mingu Lee

分类: cs.CL, cs.LG

发布日期: 2024-10-24

备注: Workshop on Efficient Natural Language and Signal Processing at NeurIPS 2024

💡 一句话要点

AdaEDL：基于熵的下界概率自适应停止推测解码，提升大语言模型推理效率。

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

关键词: 推测解码 大语言模型 推理加速 自适应草稿长度 熵估计

📋 核心要点

  1. 现有推测解码方法采用静态草稿长度，忽略了token接受率的动态变化，导致效率降低。
  2. AdaEDL通过熵估计token接受概率下界，自适应调整草稿长度，提前停止低接受概率的草稿过程。
  3. 实验表明，AdaEDL在多种设置下优于静态草稿长度方法，提升高达57%，且在高采样温度下更稳健。

📝 摘要（中文）

推测解码是一种强大的技术，旨在规避现代大型语言模型（LLM）的自回归约束。推测解码技术的目标是提高大型目标模型的平均推理时间，而不牺牲其准确性，它通过使用更高效的草稿模型来提出草稿token，然后并行验证这些token。每一轮草稿中产生的草稿token数量称为草稿长度，它通常是一个静态超参数，基于草稿token的接受率统计数据选择。然而，设置静态草稿长度可能会对性能产生负面影响，尤其是在草稿成本高且接受的token数量差异很大的情况下。自适应熵基草稿长度（AdaEDL）是一种简单、无需训练和参数的标准，它通过基于当前观察到的草稿logits的熵来近似草稿token的预期接受概率的下界，从而允许提前停止token草稿过程。我们表明，在各种设置和数据集中，AdaEDL始终优于静态草稿长度推测解码10%-57%，并且优于其他无训练草稿停止技术高达10%。同时，我们表明AdaEDL比这些技术更稳健，并在高采样温度场景中保持性能。与依赖于训练数据集特定草稿停止预测器的技术相比，由于它无需训练，AdaEDL可以无缝集成到各种预先存在的大语言模型系统中。

🔬 方法详解

问题定义：推测解码旨在加速大型语言模型的推理过程，但现有方法通常采用固定的草稿长度。这种静态策略忽略了不同token被接受的概率差异，导致在接受率较低时仍然生成大量草稿token，浪费计算资源。尤其是在草稿模型计算成本较高或token接受率波动较大的情况下，固定草稿长度会显著降低效率。

核心思路：AdaEDL的核心思想是根据已生成草稿token的熵来估计其被目标模型接受的概率下界。熵越高，表示模型对该token的不确定性越大，被接受的概率越低。通过设定一个阈值，当估计的接受概率下界低于该阈值时，提前停止草稿过程，避免生成无用的草稿token。这种自适应调整草稿长度的方法能够更有效地利用计算资源。

技术框架：AdaEDL可以无缝集成到现有的推测解码框架中。其主要流程如下：1) 使用草稿模型生成草稿token；2) 计算已生成草稿token的熵；3) 基于熵估计token被接受的概率下界；4) 如果估计的概率下界低于预设阈值，则停止生成草稿token；5) 将已生成的草稿token提交给目标模型进行验证。

关键创新：AdaEDL的关键创新在于提出了一种基于熵来估计token接受概率下界的方法。与需要训练的草稿停止预测器不同，AdaEDL无需训练，可以直接应用于各种预训练的大型语言模型。此外，AdaEDL能够自适应地调整草稿长度，从而更好地适应不同的场景和数据集。

关键设计：AdaEDL的关键设计在于如何根据熵来估计token接受概率的下界。论文中具体使用了何种函数关系来建立熵与接受概率下界之间的联系，以及如何选择合适的阈值，这些细节在原文中应该有更详细的描述（未知）。由于该方法是 training-free 的，因此没有损失函数和网络结构相关的设计。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，AdaEDL在多种设置和数据集中均优于静态草稿长度推测解码，性能提升范围为10%-57%。同时，AdaEDL也优于其他无训练草稿停止技术，提升高达10%。在高采样温度场景下，AdaEDL表现出更强的鲁棒性，能够保持性能优势。

🎯 应用场景

AdaEDL可广泛应用于各种需要加速大型语言模型推理的场景，例如在线对话系统、机器翻译、文本摘要等。通过提高推理效率，AdaEDL可以降低计算成本，提升用户体验，并促进大型语言模型在资源受限环境中的部署。该方法无需训练，易于集成，具有很高的实际应用价值。

📄 摘要（原文）

Speculative decoding is a powerful technique that attempts to circumvent the autoregressive constraint of modern Large Language Models (LLMs). The aim of speculative decoding techniques is to improve the average inference time of a large, target model without sacrificing its accuracy, by using a more efficient draft model to propose draft tokens which are then verified in parallel. The number of draft tokens produced in each drafting round is referred to as the draft length and is often a static hyperparameter chosen based on the acceptance rate statistics of the draft tokens. However, setting a static draft length can negatively impact performance, especially in scenarios where drafting is expensive and there is a high variance in the number of tokens accepted. Adaptive Entropy-based Draft Length (AdaEDL) is a simple, training and parameter-free criteria which allows for early stopping of the token drafting process by approximating a lower bound on the expected acceptance probability of the drafted token based on the currently observed entropy of the drafted logits. We show that AdaEDL consistently outperforms static draft-length speculative decoding by 10%-57% as well as other training-free draft-stopping techniques by upto 10% in a variety of settings and datasets. At the same time, we show that AdaEDL is more robust than these techniques and preserves performance in high-sampling-temperature scenarios. Since it is training-free, in contrast to techniques that rely on the training of dataset-specific draft-stopping predictors, AdaEDL can seamlessly be integrated into a variety of pre-existing LLM systems.