AdaEAGLE: Optimizing Speculative Decoding via Explicit Modeling of Adaptive Draft Structures

📄 arXiv: 2412.18910v1 📥 PDF

作者: Situo Zhang, Hankun Wang, Da Ma, Zichen Zhu, Lu Chen, Kunyao Lan, Kai Yu

分类: cs.AI, cs.CL

发布日期: 2024-12-25

💡 一句话要点

AdaEAGLE：通过显式建模自适应草稿结构优化推测解码

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

关键词: 推测解码 大型语言模型 自适应草稿结构 草稿长度预测 模型加速

📋 核心要点

  1. 现有推测解码方法依赖静态草稿结构，忽略了上下文信息，限制了解码速度的进一步提升。
  2. AdaEAGLE通过轻量级草稿长度预测器(LDLP)显式预测最佳草稿长度，从而实现自适应草稿结构。
  3. 实验表明，AdaEAGLE无需手动阈值即可达到可比的加速效果，结合阈值策略更可实现1.62倍加速。

📝 摘要（中文）

推测解码(SD)是一种流行的加速大型语言模型(LLM)推理的无损技术。我们发现，通过结合上下文感知的自适应草稿结构，可以显著提高具有静态草稿结构的SD框架的解码速度。然而，目前关于自适应草稿结构的研究受到其性能、建模方法和适用性的限制。在本文中，我们介绍了AdaEAGLE，这是第一个显式建模自适应草稿结构的SD框架。AdaEAGLE利用轻量级草稿长度预测器(LDLP)模块，在推理过程中显式预测最佳草稿token数量，以指导草稿模型。它在没有手动阈值的情况下实现了相当的加速效果，并允许更深入、更专业的优化。此外，结合基于阈值的策略，AdaEAGLE实现了比原始AR解码快1.62倍的速度，并且在保持输出质量的同时优于固定长度的SOTA基线。

🔬 方法详解

问题定义：现有推测解码方法通常采用静态的草稿结构，即草稿模型的输出长度是固定的。这种方式忽略了上下文信息，无法根据不同的输入动态调整草稿长度，从而限制了解码速度的进一步提升。此外，现有自适应草稿结构的研究在性能、建模方法和适用性方面存在局限性。

核心思路：AdaEAGLE的核心思路是显式地建模自适应草稿结构，通过预测最佳草稿长度来指导草稿模型的生成。通过上下文信息预测合适的草稿长度，使得草稿模型能够更有效地生成候选token，从而提高推测解码的效率。这种自适应的方式能够更好地利用计算资源，并减少无效的草稿token。

技术框架：AdaEAGLE框架主要包含两个核心模块：草稿模型和一个轻量级草稿长度预测器(LDLP)。首先，LDLP根据上下文信息预测最佳的草稿长度。然后，草稿模型根据LDLP预测的长度生成草稿序列。最后，目标模型验证草稿序列，并根据验证结果进行下一步的解码。整个过程无需手动设置阈值，可以进行更深入的优化。

关键创新：AdaEAGLE的关键创新在于显式地建模自适应草稿结构，并使用轻量级的LDLP模块来预测最佳草稿长度。与现有方法相比，AdaEAGLE能够根据上下文信息动态调整草稿长度，从而提高推测解码的效率。此外，AdaEAGLE无需手动设置阈值，简化了调参过程，并允许更深入的优化。

关键设计：LDLP模块是一个轻量级的神经网络，输入是上下文信息，输出是预测的草稿长度。LDLP可以使用多种网络结构，例如多层感知机或循环神经网络。损失函数可以选择均方误差或交叉熵损失，具体取决于草稿长度的表示方式（回归或分类）。草稿模型可以使用现有的预训练语言模型，例如GPT系列。AdaEAGLE还支持结合基于阈值的策略，进一步提高解码速度。

🖼️ 关键图片

img_0

📊 实验亮点

AdaEAGLE在多个数据集上进行了实验，结果表明，在没有手动阈值的情况下，AdaEAGLE实现了与现有方法相当的加速效果。结合基于阈值的策略，AdaEAGLE实现了比原始AR解码快1.62倍的速度，并且在保持输出质量的同时优于固定长度的SOTA基线。这些结果表明，AdaEAGLE能够有效地提高推测解码的效率。

🎯 应用场景

AdaEAGLE可应用于各种需要加速LLM推理的场景，例如在线对话系统、机器翻译、文本生成等。通过提高解码速度，AdaEAGLE可以降低延迟，提升用户体验，并降低计算成本。该研究对于推动LLM在实际应用中的部署具有重要意义，尤其是在资源受限的环境下。

📄 摘要（原文）

Speculative Decoding (SD) is a popular lossless technique for accelerating the inference of Large Language Models (LLMs). We show that the decoding speed of SD frameworks with static draft structures can be significantly improved by incorporating context-aware adaptive draft structures. However, current studies on adaptive draft structures are limited by their performance, modeling approaches, and applicability. In this paper, we introduce AdaEAGLE, the first SD framework that explicitly models adaptive draft structures. AdaEAGLE leverages the Lightweight Draft Length Predictor (LDLP) module to explicitly predict the optimal number of draft tokens during inference to guide the draft model. It achieves comparable speedup results without manual thresholds and allows for deeper, more specialized optimizations. Moreover, together with threshold-based strategies, AdaEAGLE achieves a $1.62\times$ speedup over the vanilla AR decoding and outperforms fixed-length SotA baseline while maintaining output quality.