Recurrent Drafter for Fast Speculative Decoding in Large Language Models

📄 arXiv: 2403.09919v5 📥 PDF

作者: Yunfei Cheng, Aonan Zhang, Xuanyu Zhang, Chong Wang, Yi Wang

分类: cs.CL, cs.LG

发布日期: 2024-03-14 (更新: 2024-12-13)

备注: 13 pages, 7 figures

💡 一句话要点

提出ReDrafter以加速大型语言模型推理

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 推测解码 递归神经网络 动态树注意力 知识蒸馏 大型语言模型 推理加速 自然语言处理

📋 核心要点

  1. 现有大型语言模型推理速度较慢，限制了其在实时应用中的有效性。
  2. ReDrafter通过结合递归神经网络和动态树注意力算法，优化了推测解码过程。
  3. 在MT-Bench中，ReDrafter实现了最高2.8倍的推理速度提升，显示出其在实际应用中的潜力。

📝 摘要（中文）

我们提出了Recurrent Drafter（ReDrafter），一种先进的推测解码方法，能够在大型语言模型（LLMs）推理中实现最先进的加速。性能提升主要得益于三个关键方面：（1）利用递归神经网络（RNN）作为草稿模型，基于LLM的隐藏状态进行条件化；（2）在束搜索结果上应用动态树注意力算法，以消除候选序列中的重复前缀；（3）通过从LLM进行知识蒸馏进行训练。ReDrafter在MT-Bench中加速Vicuna推理，使用PyTorch实现，在Nvidia H100 GPU上实现了最高2.8倍的加速。为了验证其在实际环境中的可行性，我们还在Apple Silicon芯片的Metal GPU上实现了该方法，并进行了基准测试，取得了最高2.3倍的加速。

🔬 方法详解

问题定义：本论文旨在解决大型语言模型推理速度慢的问题。现有方法在推测解码过程中存在效率低下和重复计算的问题，导致推理时间过长。

核心思路：ReDrafter的核心思路是利用递归神经网络（RNN）作为草稿模型，基于LLM的隐藏状态进行条件化，从而提高推理速度。同时，通过动态树注意力算法消除候选序列中的重复前缀，进一步优化解码过程。

技术框架：ReDrafter的整体架构包括三个主要模块：首先，使用RNN生成草稿；其次，应用动态树注意力算法对束搜索结果进行处理；最后，通过知识蒸馏从LLM中获取知识以提升模型性能。

关键创新：ReDrafter的主要创新在于结合了RNN与动态树注意力算法，这种设计使得模型能够在保持生成质量的同时显著提高推理速度，与传统方法相比具有本质的区别。

关键设计：在模型设计中，RNN的结构和参数设置经过精心调整，以确保其能够有效地捕捉LLM的隐藏状态。此外，动态树注意力算法的实现细节也经过优化，以确保在束搜索中高效消除重复前缀。训练过程中采用的知识蒸馏策略也为模型的加速提供了支持。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

在实验中，ReDrafter在MT-Bench中对Vicuna推理实现了最高2.8倍的速度提升，使用PyTorch在Nvidia H100 GPU上进行测试。此外，在Apple Silicon芯片的Metal GPU上，ReDrafter也实现了最高2.3倍的加速，展示了其在实际应用中的有效性。

🎯 应用场景

ReDrafter的研究成果在多个领域具有广泛的应用潜力，尤其是在需要实时响应的自然语言处理任务中，如对话系统、智能助手和在线翻译等。其显著的加速效果将推动大型语言模型在移动设备和边缘计算环境中的应用，提升用户体验。

📄 摘要（原文）

We present Recurrent Drafter (ReDrafter), an advanced speculative decoding approach that achieves state-of-the-art speedup for large language models (LLMs) inference. The performance gains are driven by three key aspects: (1) leveraging a recurrent neural network (RNN) as the draft model conditioning on LLM's hidden states, (2) applying a dynamic tree attention algorithm over beam search results to eliminate duplicated prefixes in candidate sequences, and (3) training through knowledge distillation from the LLM. ReDrafter accelerates Vicuna inference in MT-Bench by up to 2.8x with a PyTorch implementation on Nvidia H100 GPUs. To demonstrate its practicality in real environments, we also validated its effectiveness for on-device applications by implementing the approach in MLX and benchmarking performance on Metal GPUs in Apple Silicon chips, achieving up to 2.3x speedup.