Direct Alignment of Draft Model for Speculative Decoding with Chat-Fine-Tuned LLMs

📄 arXiv: 2403.00858v4 📥 PDF

作者: Raghavv Goel, Mukul Gagrani, Wonseok Jeon, Junyoung Park, Mingu Lee, Christopher Lott

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

发布日期: 2024-02-29 (更新: 2024-05-13)

备注: 8 pages, 3 figures, Published at the ICLR 2024 Workshop on Understanding of Foundation Models (ME-FoMo)

💡 一句话要点

提出直接对齐草稿模型以加速LLM推理

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

关键词: 大规模语言模型 推测解码 知识蒸馏 模型训练 自然语言处理

📋 核心要点

  1. 现有的LLM推理方法因内存限制导致低效，尤其是在生成速率方面表现不佳。
  2. 论文提出了一种草稿模型训练框架，直接对齐聊天能力目标模型，简化了训练过程。
  3. 实验表明，Llama 2 Chat Drafter 115M在推测解码下实现了显著的性能提升，速度提高2.4倍。

📝 摘要（中文）

大规模语言模型（LLMs）的文本生成因其自回归特性、庞大的参数量和有限的内存带宽而受到内存限制，导致低的令牌生成速率。为了解决这一问题，论文提出了一种简单的草稿模型训练框架，能够直接对齐聊天能力目标模型。通过该框架，研究团队训练了Llama 2 Chat Drafter 115M，作为Llama 2 Chat 7B或更大模型的草稿模型，且仅占原始模型1.64%的大小。该框架包括预训练、蒸馏数据集生成和知识蒸馏微调，且不需要额外的对齐过程。实验结果表明，使用该草稿模型的推测解码在多项任务中实现了高达2.3的块效率和2.4倍的速度提升。

🔬 方法详解

问题定义：本论文旨在解决大规模语言模型推理过程中的内存限制问题，现有方法在生成速率上表现不佳，尤其是在缺乏高质量草稿模型的情况下。

核心思路：提出了一种简单的草稿模型训练框架，能够直接对齐目标聊天模型，避免了复杂的对齐过程。通过训练小型草稿模型来加速推理过程。

技术框架：该框架包括三个主要阶段：预训练、蒸馏数据集生成和知识蒸馏微调。预训练阶段为草稿模型提供基础，蒸馏数据集生成阶段则利用目标模型生成指令-响应对，最后通过知识蒸馏进行微调。

关键创新：最重要的创新在于提出了新的Total Variation Distance++ (TVD++)损失函数，该函数结合了强化学习中的方差减少技术，提升了模型的训练效果。

关键设计：在微调阶段，使用目标模型生成的指令-响应对作为蒸馏数据，确保训练数据的合理分布，且整个过程不需要额外的对齐步骤。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，Llama 2 Chat Drafter 115M在使用推测解码时实现了高达2.3的块效率，相较于自回归解码速度提升了2.4倍，且无需进一步的任务特定微调，展示了其在多项任务中的优越性能。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和智能助手等。通过加速LLM推理，能够提升用户体验，降低计算资源消耗，具有显著的实际价值和广泛的应用前景。

📄 摘要（原文）

Text generation with Large Language Models (LLMs) is known to be memory bound due to the combination of their auto-regressive nature, huge parameter counts, and limited memory bandwidths, often resulting in low token rates. Speculative decoding has been proposed as a solution for LLM inference acceleration. However, since draft models are often unavailable in the modern open-source LLM families, e.g., for Llama 2 7B, training a high-quality draft model is required to enable inference acceleration via speculative decoding. In this paper, we propose a simple draft model training framework for direct alignment to chat-capable target models. With the proposed framework, we train Llama 2 Chat Drafter 115M, a draft model for Llama 2 Chat 7B or larger, with only 1.64\% of the original size. Our training framework only consists of pretraining, distillation dataset generation, and finetuning with knowledge distillation, with no additional alignment procedure. For the finetuning step, we use instruction-response pairs generated by target model for distillation in plausible data distribution, and propose a new Total Variation Distance++ (TVD++) loss that incorporates variance reduction techniques inspired from the policy gradient method in reinforcement learning. Our empirical results show that Llama 2 Chat Drafter 115M with speculative decoding achieves up to 2.3 block efficiency and 2.4$\times$ speed-up relative to autoregressive decoding on various tasks with no further task-specific fine-tuning.