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.

翻译：大语言模型（LLMs）的文本生成因其自回归特性、庞大的参数量和有限的内存带宽而受限于内存瓶颈，常导致较低的令牌生成速率。推测解码已被提出作为加速LLM推理的解决方案。然而，现代开源LLM系列（如Llama 2 7B）通常缺乏现成的草案模型，因此需训练高质量草案模型以通过推测解码实现推理加速。本文提出一种简单的草案模型训练框架，用于直接对齐具备聊天能力的目标模型。利用该框架，我们训练了Llama 2 Chat Drafter 115M（仅为Llama 2 Chat 7B或更大模型原始规模的1.64%）。该训练框架仅包含预训练、蒸馏数据集生成以及基于知识蒸馏的微调，无需额外对齐步骤。在微调阶段，我们使用目标模型生成的指令-回复对在合理数据分布中进行蒸馏，并提出了新型总变差距离++（TVD++）损失函数，该函数融合了源自强化学习策略梯度方法的方差缩减技术。实验结果表明，搭载Llama 2 Chat Drafter 115M的推测解码在多种任务中无需特定任务微调即可实现最高2.3倍的块效率及2.4倍的自回归解码加速比。