[Arxiv 2024] EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees

Introduction

• 作者提出 Context-Aware Dynamic Draft Tree，对 EAGLE 中固定的 draft tree 做出了改进，通过动态扩展 draft tree 有效节省 draft model 的推理时间，还能通过 token re-rank 给 LLM 提供更少更高质量的 draft tokens，从而在保证接收率的前提下降低 LLM 的验证时间，相比 EAGLE 有了明显的加速
  

Method

• Expansion Phase. 在扩展 draft tree 时，只选择最深层上累积接收概率最高的 top-$k$ tokens 进行扩展，从而降低 draft model 的推理开销，其中累积接收概率为根节点到当前节点的接收概率之积，接收概率可以用 token 置信度近似 $V_i={\prod }_{t_j\in \text{Path}(\text{root},t_i)}{p}_j\approx {\prod }_{t_j\in \text{Path}(\text{root},t_i)}{c}_j$ (作者发现 draft model is well-calibrated)；例如下图中，只选择 current layer (orange blocks) 进行扩展
  
• Reranking Phase. draft tree 扩展完成后，作者从 token tree 中选出累积接收概率的 top-$m$ tokens 交给 LLM 验证，如果两个 token 累计接收概率相同，则优先选择位于更浅层的 token；这样选出的 top-$m$ tokens 仍然可以形成一棵树，可以用 tree attention 进行验证
  

Experiments

• Effectiveness.
  
• Ablation Study. 这里的 w/o value 指扩展 draft tree 时不使用累积接收概率而是直接使用单个 token 的 confidence
  

References

• Li, Yuhui, et al. “EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees.” arXiv preprint arXiv:2406.16858 (2024).
• code: https://github.com/SafeAILab/EAGLE