发表时间: 2025-11 · MLSys'26
作者/机构: Jameson Sandler, Jacob K. Christopher, Thomas Hartvigsen, Ferdinando Fioretto
本文旨在解决大型语言模型(LLM)推理加速中的两个核心瓶颈。当前,推理性解码(Speculative Decoding)已成为加速LLM推理的标准方法,但其效率受限于两大问题:
1. 起草阶段的自回归依赖:传统的起草模型(drafter)通常是自回归的,一次只能生成一个词元(token),这限制了并行计算的潜力,构成了延迟瓶颈。
2. 起草与验证模型的不对齐:起草模型与验证模型(verifier)之间的分布差异导致起草的词元被频繁拒绝,降低了加速效果,构成了对齐瓶颈。
为同时解决这两个瓶颈,本文提出了SpecDiff-2,一个新颖的推理性解码框架。其核心创新点和主要贡献如下:
基于离散扩散的并行起草机制:引入离散扩散语言模型(DLM)作为非自回归的起草模型。DLM通过少数几步去噪过程并行生成整个草稿序列,其延迟主要取决于去噪步数而非序列长度,从而有效解决了起草阶段的延迟瓶颈。
训练时对齐方法:条带蒸馏(Streak-distillation):为解决扩散起草模型与自回归验证模型之间的根本性差异(前者学习序列的联合分布,后者学习前缀条件下的下一个词元分布),本文提出了一种新颖的微调方法。该方法直接以提升理论加速率为目标,通过一个代理目标函数鼓励起草模型生成更长的、能被验证模型接受的词元序列(“条带”)。
测试时对齐方法:自选择接受(Self-selection acceptance):该机制在测试时进一步提升吞吐量。它利用扩散模型能以极低成本并行生成多个候选草稿的特性,让验证模型在正式逐词元验证前,先评估所有候选草稿的预期接受长度,并选择最优的一个进行验证。
实现业界领先的加速效果:在多个涵盖推理、编码和数学任务的基准测试中,SpecDiff-2展现了最先进的性能。与之前的基准方法相比,其每秒生成词元数(tokens-per-second)平均提升高达55%,相较于标准解码实现了最高5.5倍的平均加速,且完全不损失模型输出的准确性。如下图1所示,SpecDiff-2中经过对齐的起草模型比基础的SpecDiff起草模型性能提升超过40%。
推理性解码(SD)算法:该算法基于两个语言模型:一个小型起草模型 $Q$ 和一个更大型的目标验证模型 $P$。它们共同生成一个词元序列 $x_{1:n} = (x_1, ..., x_n)$。对于给定的前缀 $s$,模型分别提供下一个词元的后验概率 $q(· | s)$ 和 $p(· | s)$。生成过程以大小为 $\gamma$ 的块进行:起草模型提出草稿 $x_{1:\gamma} \sim Q(· | s)$,而验证模型并行地对所有位置进行评分。词元从左到右使用标准的无损接受规则进行提交。
标准接受规则:对于在位置 $i$ 的一个起草词元 $x_i$,其接受规则为:
其中,$s$ 代表 $x_i$ 之前的前缀,即 $s = x_{1:i-1}$(当 $i > 1$ 时)。此规则意味着当 $q(x_i | s) \le p(x_i | s)$ 时,词元会被立即接受。
拒绝与替换:当在位置 $k$ 首次出现拒绝时,草稿中剩余的部分将被丢弃,并从归一化的残差分布中抽取一个替换词元:
这些规则保证了最终生成的文本与仅使用验证模型 $P$ 进行的常规解码在分布上完全一致,同时实现了草稿的并行提出和验证【10, Fast inference from transformers via speculative decoding by Leviathan et al., 2023, International Conference on Machine Learning】。关于其分布保真性的证明见附录C.1。
接受率决定加速效果:在“起草-验证”解码模式中,实际获得的加速效果由接受率决定。
SD接受率定义:对于前缀 $s$、起草模型 $Q$ 和验证模型 $P$,推理性解码的接受率 $\alpha(s)$ 定义为:
预期接受长度:进一步地,位置 $j$ 的接受率 $\gamma_j(s)$ 定义为在前 $j-1$ 个草稿词元均被接受的条件下,第 $j$ 个草稿词元被接受的概率。根据标准的“接受连乘”恒等式,对于一个长度为 $\gamma$ 的草稿,预期提交的词元数量为:
其中 $\oplus$ 表示序列拼接。因此,提升 $\gamma_j(s)$ 能单调地增长接受的词元序列长度,并在固定的起草成本下提高吞吐量。
对齐的挑战:接受率是衡量 $P$ 和 $Q$ 之间“对齐”程度的指标,与两者之间的总变分距离(TV)成反比。因此,基于蒸馏的对齐方法,如直接最小化TV距离,已成为加速自回归起草模型推理性解码的常用手段【31, DistillSpec: Improving speculative decoding via knowledge distillation by Zhou et al., 2023】。然而,这些方法通常将每个位置的接受率视为可交换的,并且只优化紧随前缀的第一个位置的对齐(即最大化 $\gamma_1(s \oplus x_1)$)。当损失函数定义在窗口大小 $\gamma=1$ 上时,这种简化的最小化方法与真正的接受标准等价(推导见附录C.2)。
扩散模型的特殊性:对于扩散起草模型而言,不同位置的对齐水平差异显著,因此必须对整个草稿窗口进行仔细优化才能获得预期的加速。图2的经验性数据显示,对于基础的扩散起草模型(“Base”),以及通过自回归目标函数对齐的扩散模型(“AR-distillation”),其位置接受率 $\gamma_j$ 随着草稿索引 $j$ 的增加而迅速下降。特别是,“AR-distillation”将增益集中在早期位置,导致后续词元的校准效果很差。相比之下,本文提出的条带感知目标(“Streak-distillation”)能够在草稿窗口的后续索引上正确对齐,其平均接受率比AR方法高出3.2倍。这凸显了为扩散模型设计能够对齐整个窗口的对齐程序的必要性。
Speculative Diffusion Decoding (SpecDiff)回顾:在介绍我们的新颖对齐技术之前,本节回顾了SpecDiff框架【4, Speculative diffusion decoding: Accelerating language generation through diffusion by Christopher et al., 2025, NAACL】。SpecDiff用一个低延迟的扩散起草模型 $Q_{\text{diff}}$ 替代了自回归起草模型 $Q$。这利用了并行起草的一个关键特性:窗口大小为 $\gamma$ 的草稿词元是同时生成的,因此起草成本主要取决于扩散模型的去噪步数,而不是 $\gamma$,从而解决了前述的延迟瓶颈。
掩码离散扩散语言模型 (MDM):离散扩散语言模型通过“腐蚀-去噪”步骤实现非自回归的文本生成。一个前向过程从一个干净的序列 $x$ 生成一个部分腐蚀的序列 $x(t)$(其中 $t \in [0, 1]$),然后一个去噪器 $Q_{\text{diff}}$ 预测以 $x(t)$ 为条件的词元分布。在掩码扩散模型(MDM)【15, Simple and effective masked diffusion language models by Sahoo et al., 2024】中,引入了一个特殊的 [MASK] 词元,每个位置 $i$ 以 $1-\epsilon(t)$ 的概率被独立替换为 [MASK]。去噪器并行地为所有位置输出 $Q_{\text{diff}}(\cdot | x(t); \theta)_i$,并通过在当前被掩码的集合 $M(t)$ 上使用掩码交叉熵进行训练:
起草过程:给定前缀 $s$ 和窗口大小 $\gamma$,起草过程首先用 [MASK] 词元填充草稿位置 $x_{1:\gamma}$;然后应用少量的去噪步骤(实践中通常为一步)来产生一个联合提议 $x_{1:\gamma} \sim Q_{\text{diff}}(\cdot | s \oplus [\text{MASK}]_{1:\gamma})$,这个过程在所有位置上是并行的。本文用 $Q_{\text{diff}}(\cdot|s)$ 来表示这个联合过程。
对齐挑战重述:尽管推理性扩散提供了低且可预测的起草延迟,但它本身并未解决对齐问题。扩散起草模型被训练来建模数据块上的联合去噪分布,而验证模型评估的是基于前缀条件的下一个词元后验概率。因此,即使联合样本在全局看起来很流畅,词元级别的校准不足仍会降低接受率。
本节介绍本文的关键贡献:两种专为扩散起草模型设计的对齐机制。一种是训练时(微调)的目标函数,其效果随蒸馏计算量扩展(第5.1节);另一种是测试时的选择规则,其效果随验证模型计算量扩展(第5.2节)。重要的是,这两种机制都只作用于扩散起草模型 $Q_{\text{diff}}$,而验证模型 $P$ 始终保持冻结。这些机制共同实现了稳健的起草-验证模型对齐,并且这种对齐能力能泛化到微调数据之外。
目标:本节的目标是以一种直接提升公式(3)中定义的预期接受条带长度 $\text{TokensDraft}(\gamma, s)$ 的方式,来对齐扩散起草模型与验证模型,同时保持训练过程的可行性。
贪心接受代理:我们首先构建一个接受率的代理——“贪心接受”(greedy acceptance),它能提供平滑的训练信号。在给定前缀 $s$ 的情况下,接受一个起草词元的概率被设定为验证模型的概率:
这个方案不使用起草模型的后验概率进行验证,仅作为推导蒸馏目标的分析工具。
位置接受率的简化:在贪心接受方案下,沿着教师路径 $x_{1:j-1} \sim P(\cdot|s)$ 的位置接受率 $\tilde{\gamma}j$ 简化为如下乘积形式:
其中,期望是关于 $x$ 为条件,因为所有词元都是并行生成的。}$ 的。需要注意的是,扩散起草模型并不以 $x_{1:j-1
关键恒等式:由于公式(4)中的内部和是两个词汇表 $V$ 上分类分布的点积,它可以被重写为关于任一分布的期望。因此,对验证模型 $P$ 的依赖可以通过将其吸收到 $x_j$ 的采样分布中来变得隐式,从而得到:
其中外部期望仍然是关于从验证模型链中抽取的教师前缀 $x_{1:j-1}$,而内部期望可以通过从验证模型 $P$(蒸馏期间)或起草模型 $Q_{\text{diff}}$(验证期间)采样 $x_j$ 来评估。
条带蒸馏目标:这个恒等式非常重要,因为它允许对条带长度代理进行路径化重表达,并直接导出一个可行的训练目标。将公式(5a)代入公式(3)的接受连乘结构中,我们得到了一个 $\text{TokensDraft}(\gamma, s)$ 的路径化估计器,其中内积沿着教师轨迹进行评估。在所有前缀和验证模型续写上最大化这个量,就得到了条带蒸馏(streak-distillation)目标。
定义 5.1 (条带蒸馏):设 $P$ 为一个冻结的验证模型,$Q_{\text{diff}}$ 为一个扩散起草模型,其位置边际分布为 $q_j(\cdot|s;\theta)$。条带蒸馏目标是:
这个条带蒸馏目标反映了公式(3)中的理想目标:保留了接受连乘的结构,并将每个因子替换为在验证模型词元处评估的贪心接受项。因此,目标是优化整个草稿窗口,这是通过定理5.1中可行的代理目标而不是棘手的、路径依赖的接受/拒绝过程来完成的。同样的分析也揭示了在推理时的一个补充手段:在固定的验证模型下,使用面向条带的代理在多个候选草稿中进行选择,可以进一步扩大接受的前缀长度。
动机:虽然条带蒸馏在训练时改善了起草-验证模型的对齐,但通过在测试时利用验证模型的计算能力,可以获得额外的增益。先前的方法(如EAGLE-2)通过自回归的多路径扩展来实现这一点,即用 $Q_{ar}$ 生成多个提议路径以增加验证模型接受某个续写的可能性【12, EAGLE-2: Faster inference of language models with dynamic draft trees by Li et al., 2024b】。然而,当使用自回归起草模型时,这种测试时方法有其缺点:首先,起草的复杂度增加;其次,多路径的生成必须以并行的树状方式进行,增加了复杂性。
扩散模型的多草稿优势:相比之下,本文认为扩散起草模型支持一种更简单、更高效的多草稿机制。单次去噪过程就能暴露所有位置的边际分布,从中可以以可忽略的额外成本采样出多个联合草稿。这一观察启发了一种测试时机制,即在进行无损验证之前,使用验证模型在这些草稿中进行选择。
自选择接受机制:所提出的测试时程序,称为自选择接受(self-selection acceptance),具有以下特性:(1) 它只需对标准的推理性扩散进行最小的改动;(2) 扩散起草模型提供了时间和计算上都高效的起草方式,对于 $K$ 条路径,起草计算的复杂度为 $O(1)$(对于固定的 $\gamma$);(3) 所提出的方法随 $K$ 扩展,允许以可忽略的串行开销获得额外的吞吐量,如图5所示。
实现流程:首先,扩散模型提供了高效生成多条路径的独特能力。在推理性扩散设置中(见第4节),对 [MASK] 进行单步去噪,$Q_{\text{diff}}$ 会产生位置边际分布 ${q_j(\cdot | s)}{j=1}^{\gamma}$。通过独立采样 $x} \sim \prod_{j=1}^{\gamma} q_j(\cdot|s)$ 可以得到一个联合草稿,其成本相对于神经网络前向传播可忽略不计。重复此采样 $K$ 次,可以得到 ${x^{(k){1:\gamma}}$,同样具有可忽略的串行开销。所提出的自选择方法使用一个面向条带的验证模型分数对这些候选草稿进行排序,并选择最佳草稿进行无损验证。}^{K
草稿评估与选择:为了评估 $K$ 个生成的草稿中哪一个与验证模型 $P$ 诱导的分布 $p_j(\cdot|s)$ 对齐得最好,本文使用了定义5.1中条带蒸馏目标的一个理论上等价的变体。从该条带目标出发,并利用公式(5b)的等价性,我们得到:
其中 $x \sim Q_{\text{diff}}(\cdot|s)$ 现在是由扩散模型 $Q_{\text{diff}}$ 生成的草稿。这种替代形式允许我们直接使用验证模型 $P$ 来确定一个草稿 $x$ 预期的接受词元数。
吞吐量最大化选择:重要的是,公式(7)量化了草稿 $x$ 预期对推理吞吐量的贡献程度。因此,对于从 $Q_{\text{diff}}(\cdot|s)$ 中采样的 $K$ 个草稿,可以通过以下方式选择吞吐量最大化的草稿,记为 $x^{\text{max}}$:
该方法通过基于每个生成草稿的预期吞吐量进行选择,从而在构造上最大化了预期加速。在操作上,验证模型使用词元级后验概率 $p_j(x_j^{(k)} | s \oplus x_{<j}^{(k)})_{j=1}^{\gamma}$ 对每个候选草稿进行评分,然后在应用接受规则之前选择 $x^{\text{max}}$。如附录B.8所述,所有 $K$ 个草稿的分数可以通过树状注意力【25, DySpec: Faster speculative decoding with dynamic token tree structure by Xiong et al., 2024】高效计算。
算法 1: 自选择接受
输入: 前缀 s; 验证模型 P(·|s, θ); 起草模型 Q_diff(·|s); 整数 K, γ
输出: 从 P(·|s) 生成的最终序列 y ∈ V*
(q_1, ..., q_γ) ← Q_diff(·|s) (草稿边际概率)
并行地对 k 从 1 到 K 执行:
x_k(1:γ) ~ (q_1, ..., q_γ) (采样K个草稿)
τ_k ← TokensDraft(x_k(1:γ), s) (计算预期吞吐量)
(x_max, y) ← (arg max_k τ_k, []) (选择最佳草稿)
y ← []
对于 i 从 1 到 γ:
p_i ← P(x_max_i | s ⊕ y)
b ~ Bernoulli(p_i)
如果 b = 1:
y ← y ⊕ x_max_i (接受起草的词元)
否则:
x'_i ~ P(· | s ⊕ y) / (1 - P(x_max | s ⊕ y)) (替换词元)
y ← y ⊕ x'_i
返回 y
返回 y
使用贪心接受进行无损验证:一旦 $P$ 选定了草稿 $x^{\text{max}}$,就必须对其进行验证以确定接受和拒绝哪些词元。部署贪心接受规则变得非常简单,因为 $x^{\text{max}}$ 上的验证模型概率已在之前的排序阶段被缓存。此外,请注意,从未考虑 $x^{\text{max}}$ 的起草概率,这使得该方法支持跨分词器(cross-tokenizer)应用,并可扩展到那些不输出校准概率分布的扩散起草模型【15, Simple and effective masked diffusion language models by Sahoo et al., 2024; 17, Simplified and generalized masked diffusion for discrete data by Shi et al., 2024】。
评估目标:本节对SPECDIFF-2与当前最先进的推理性解码方法进行全面评估,涵盖自回归(AR)和基于扩散的起草模型。研究目标有两个:(i) 量化我们的测试时和训练时对齐方法在推理性解码中带来的增益;(ii) 探究基于扩散的起草在低预算下扩展高效推理的适用场景。所有方法的输出都与验证模型完全一致(无损)。
模型与数据集:
- 基准模型:
1. Speculative Sampling (SpS):经典的AR起草-验证方法。
2. EAGLE / EAGLE-2:采用验证器对齐的早期接受机制的AR起草模型,代表了AR推理性解码的SOTA。
3. SpecDiff:原始的、使用未对齐扩散起草模型的方法,用于隔离评估本文提出的对齐方法的影响。
硬件配置:
软件配置:
- 实现:依赖于EAGLE和EAGLE-2的公开实现以保证可比性。
模型配置:
- 验证模型:
- QWEN2.5-72B-INSTRUCT
- LLAMA-2-70B-CHAT
- 附录中还包括QWEN2.5-14B-INSTRUCT和LLAMA-2-13B-CHAT。
起草模型:
扩散起草模型:
AR起草模型 (用于基准):
模型配对总结:如下表1所示。
| 验证模型 | 起草模型 | 共享分词器 |
|---|---|---|
| Qwen2.5-72B-Instruct | DiffuCoder-7B | Qwen2.5 |
| Qwen2.5-72B-Instruct | Qwen2.5-7B-Instruct | Qwen2.5 |
| Llama-2-70B-Chat | DiffuLlama-7B | Llama-2 |
| Llama-2-70B-Chat | Llama-2-7B-Chat | Llama-2 |
表1. 验证器-起草器配对和共享的分词器。
关键超参数:
- 扩散步数:设置为1步。附录A.2的消融实验表明,增加步数会线性增加起草延迟,而接受的词元数量增长停滞。
- 起草温度 (Temperature):设置为1.5。如图5所示,该温度在草稿质量和多草稿(K个)之间的方差之间取得了良好平衡,从而充分利用自选择机制。
由于所有方法均采用无损解码方案,输出与验证模型完全一致,因此省略了PPL、ROUGE等文本质量指标。
核心发现:如表2所示,SpecDiff-2在所有设置中始终获得最高的加速比和最长的平均接受条带长度(TokensDraft),显著优于现有SOTA模型。
- 平均性能:SpecDiff-2在所有设置中平均实现了4.22倍的加速,比EAGLE-2高出超过30%。
- 最佳性能:在起草模型专长的领域(如使用DIFFUCODER进行代码生成),加速比飙升至超过标准自回归生成的5倍。
关键趋势分析:
1. 更大容量的起草模型带来更高吞吐量:7B参数的起草模型相比EAGLE-2的1B模型,产生了更长的接受条带,并直接转化为更高的加速比。
2. 扩散起草的并行性优势:尽管扩散起草模型更大,但其并行生成草稿的特性使得起草延迟与需要多次前向传播的、更小的EAGLE/EAGLE-2模型相当。
任务敏感性:
- 结构化任务(数学、代码):SpecDiff-2在这些任务上表现最佳,平均加速比为4.71倍,而EAGLE-2为3.43倍。这表明扩散起草模型在目标分布结构更受限(如代码、步骤化推理)时表现出色。
- 开放式问答:在语义多样性更高的开放式问答中,SpecDiff-2的优势有所减小,加速比为3.24倍,EAGLE-2为2.80倍。尽管如此,SpecDiff-2仍然比SOTA提升了16%。
这些结果表明,SpecDiff-2特别适用于加速结构化的LLM推理任务。
动机:推理加速算法的一个重要动机是解决因测试时计算扩展(如思维链CoT)而引入的运行时延迟问题。在有时间限制的场景下,更快的解码速度能将相同的物理时间转化为更多的有效推理词元,从而提升任务准确率。
实验设置:
- 模型:Qwen2.5-72B-Instruct。
- 任务:Math500基准测试,模型被指示以思维链(CoT)方式进行思考。
- 流程:给定一个时间预算 $b$ 秒,时间到后,模型被提示“思考时间结束,总结你的答案”,然后评估最终答案的正确性。
实验结果 (图6,图中未提供,根据文本描述):
- 推理时间与准确率:增加推理预算(x轴)能直接提升Math500任务的准确率,因为更多的推理时间允许模型进行更深入的思考。
- SpecDiff-2的增益:由于SpecDiff-2提高了吞吐量,在相同的物理时间内,加速后的系统能够生成更多的推理词元,从而达到更高的准确率。
- 在 $b=15$ 秒的限制下,加速后的Qwen2.5-72B模型相比基础模型准确率提升了 +63%。
- 相比未对齐的SpecDiff,准确率进一步提升了 +11%。
基线对比:表3将SpecDiff-2与无对齐的SpecDiff进行对比,显示了本文提出的训练时和测试时对齐技术的显著效果。
训练时扩展(条带蒸馏):
- 结果 (图7):随着条带蒸馏训练步数的增加,吞吐量持续提升。在训练了60,000步后,Qwen模型的加速比提升了30%,与DistillSpec在自回归起草模型上取得的改进相当。最终,经过蒸馏的起草模型吞吐量比EAGLE-2基线高出32%。
测试时扩展(自选择接受):
- 结果 (图5):在Math500上,随着并行草稿数量 $K$ 的增加,加速比平滑提升。这是因为更多的草稿增加了验证模型选中一个对齐样本的概率。
- 在起草温度 $T=2.0$ 时增益最大,当 $K=8$ 时,吞吐量额外提升了 +20%。
- 在低温度 $T=0.1$ 时,由于草稿间方差小,扩展效果不明显。
综合效果:
早期LLM推理加速:早期的工作主要集中在改变模型架构,如模型剪枝【18, A simple and effective pruning approach for large language models by Sun et al., 2024】和量化【6, GPTQ: Accurate post-training quantization for generative pre-trained transformers by Frantar et al., 2023】,这些方法以牺牲生成质量为代价来提高效率。
无损加速与推理性解码:推理性解码【10, Fast inference from transformers via speculative decoding by Leviathan et al., 2023; 2, Accelerating large language model decoding with speculative sampling by Chen et al., 2023】作为一种无损加速方法应运而生。现有研究主要分为两类,分别解决推理性解码的两个瓶颈之一:
起草模型对齐:通过专门的训练程序或测试时起草增强来改善。
起草模型延迟:通过并行化来解决。
本文的定位:本文的工作同时解决了上述两个瓶颈,并报告了对这些先前工作的显著改进。
接受规则的扩展:本文提出的贪心接受规则表明,替代性的接受标准既实用又有利。对这些规则的正式分析(如偏见/方差权衡、最坏情况下的接受率)是一个开放方向。由于贪心接受天然支持异构起草模型和不匹配的分词器,系统地研究跨家族和跨分词器的起草是一个关键的下一步。
多样化的起草策略:自选择起草机制为并行生成多个草稿的不同程序奠定了基础。一个自然的扩展是部署 $K$ 个专门针对子分布(如代数运算、数值计算、代码)的独特起草模型,并在测试时进行选择。这引发了关于多样性-对齐权衡、成本感知的选择策略以及如何在 $K$、温度和提议深度之间分配固定计算预算的开放问题。
扩散起草模型的最佳尺寸:本文的实证结果表明,扩散的并行起草特性使其能够扩展到比自回归架构大得多的起草模型,同时保持相似的起草延迟。然而,在较小的验证模型上速度提升较小,表明通过按比例缩放DLM起草模型与验证模型的尺寸,可能会实现更大的加速。推导扩散起草模型的缩放定律和计算-延迟边界是一个开放方向。
硬件优化:目前缺乏专门为半自回归推理优化的有效内核,并且在长扩散草稿的背景下支持KV缓存对于高效推理至关重要。
本研究提出了SpecDiff-2,这是首个针对使用扩散起草模型的推理性解码进行模型对齐的工作。为了同时解决起草延迟和起草-验证模型对齐两大瓶颈,本文引入了专为扩散起草模型设计的训练时蒸馏技术(条带蒸馏)和测试时并行草稿生成与选择机制(自选择接受)。
通过利用这些新颖的扩散对齐方法,所提出的框架在词元吞吐量上取得了显著提升,与之前的基准方法相比,生成速度平均提升高达+55%,相较于标准解码实现了最高5.5倍的平均加速,且完全不损失准确性。SpecDiff-2在推理、编码和数学等多个基准测试中均达到了最先进的性能,为无损加速领域树立了新的标杆。
我们相信,使用扩散起草模型代表了推理时加速的重要一步,为加速大型语言模型开辟了多个有前景的研究方向。
A.1 较小验证模型上的结果:在QWEN-2.5-14B和LLAMA-2-13B等较小验证模型上,7B参数的扩散起草模型显得过大。如表4所示,DIFFUCODER仍能在多数设置中优于EAGLE-2,但性能较弱的DIFFULLAMA则无法持续超越。这表明起草模型的尺寸应与验证模型成比例,以达到最佳效果。
A.2 LLaMA模型温度消融:如图8(图中未提供)所示,对于DIFFULLAMA起草模型,在扩展草稿数量时,保持temperature=1.0可获得最佳性能。
A.3 扩散步数消融:如图9(图中未提供)所示,当扩散步数 $T > 1$ 时,加速效果迅速下降。这是因为起草延迟随 $T$ 线性增加,而接受的词元数量增长却停滞不前。因此,实验中采用 $T=1$。
A.4 Gamma (草稿窗口) 消融:如图10(图中未提供)所示,不同的窗口大小 $\gamma$ 会影响整体加速。实验中为DIFFUCODER设置 $\gamma=32$,为DIFFULLAMA设置 $\gamma=16$。最佳窗口大小主要取决于扩散起草模型的质量,这可能与预训练数据的规模和质量有关。
A.5 联合对齐扩展:如图11(图中未提供)所示,结合条带蒸馏和测试时自选择对齐技术,相比基础的扩散起草模型和EAGLE-2基线,性能提升了+39%,在Math500数据上对QWEN2.5-72B模型实现了4.3倍的吞吐量提升。
A.6 条带蒸馏 vs. DistillSpec:如图12(图中未提供)所示,将DistillSpec简单调整用于扩散模型(即对所有位置的接受率取平均)的效果远不如条带蒸馏。这是因为这种朴素的实现方式忽略了位置依赖性,过度强调了后部词元,与理论上的吞吐量定义不符,从而导致了实际吞吐量的差距。
B.1 硬件设置:实验在单台或两台NVIDIA A100 80GB GPU上执行。模型加载和LoRA合并在CPU上完成,以避免GPU内存碎片化。
B.2 软件环境:使用Python 3.10, PyTorch 2.3+, Transformers 4.46.x, PEFT等库。启用了tf32以利用A100的Tensor Core。
B.3 模型加载:验证模型使用AutoModelForCausalLM.from_pretrained(...)在GPU上以bfloat16格式加载。起草模型首先在CPU上加载,合并LoRA适配器(如果提供),然后将合并后的单一模型对象移动到GPU并转换为目标数据类型。
B.4 接受与分布收集:通过一个DraftProbCollector钩子在起草模型的生成循环中收集草稿词元的概率分布。验证端则应用贪心接受规则,并记录每次通过的接受条带长度。
B.5 评估协议:对每个基准样本,分别执行推理性解码和基线(仅验证模型)生成,记录各自的端到端时间。吞吐量以(生成的新词元数 - 1)/ 总时间计算,加速比为两者吞吐量的比率。
B.6 数据集流式传输:基准测试数据直接从Hugging Face流式传输,包括MATH-500, LiveCodeBench, MT-Bench等。
B.7 蒸馏数据集:
temp=1.0下生成多个续写构成语料库。B.8 测试时验证:
- 树状结构:$K$个不同的草稿可能在不同位置有冗余,因此构建一个词元树。
- 并行验证:将树状结构展平为一个序列块,并附带一个自定义的注意力掩码,传递给验证模型。该掩码确保每个词元的注意力只作用于其正确的路径前缀。
- 实现:使用PyTorch张量操作实现树的展平、掩码构建和结果的批处理收集。KV缓存在不同推理性解码轮次之间被复用。
- 效率:通过轻量级的性能分析工具确保树构建和掩码生成的开销不会主导运行时间。
定理 C.1 (SD-Invariance):通过推理性采样从 $p(x)$ 和 $q(x)$ 中抽取的词元与直接从 $p(x)$ 中抽取的词元分布相同。
证明:(如【10, Fast inference from transformers via speculative decoding by Leviathan et al., 2023, International Conference on Machine Learning】中所述)令 $\alpha$ 为接受概率。调整后的分布 $p'(x) = \text{norm}(\max(0, p(x)-q(x)))$ 的归一化常数为 $1-\alpha$。对于任意词元 $x$,其最终被采样的概率为(接受 $x$ 的概率 + 拒绝 $x$ 并从修正分布中采样得到 $x$ 的概率):
定理 C.2 (Position-Wise Alpha Identity):若位置接受率 $\alpha_1, ..., \alpha_\gamma$ 是可交换的(如AR起草模型所假设),则最大化TokensDraft简化为最大化 $\alpha_1(s \oplus x_1)$(紧随前缀s的接受率)。
证明:如果所有 $\alpha_j$ 都相同(设为 $\alpha$),则 $\text{TokensDraft} = \sum_{j=1}^{\gamma} \alpha^j = \alpha \frac{1-\alpha^\gamma}{1-\alpha}$。由于该表达式是关于 $\alpha$ 的单调递增函数,最大化TokensDraft等价于最大化 $\alpha$。在AR模型中,所有位置的生成过程类似,因此优化第一个位置的接受率 $\alpha_1$ 被视为优化所有位置的代理。
以下是Qwen2.5-72B-Instruct模型在Math500语料库上的生成示例,温度为0.0,最大生成3096个词元。
问题 1: 代数
问题:希拉里有十一枚硬币,都是一角和五分的。这些硬币总共价值75美分。她有多少枚五分硬币?
解:设 $d$ 为一角硬币的数量, $n$ 为五分硬币的数量。
将第二个方程除以5,得到 $2d + n = 15$。用这个方程减去第一个方程,得到 $d = 4$。代回:$4 + n = 11 \Rightarrow n = 7$。
答案:7
问题 2: 中级代数
问题:求 $f(x, y) = x\sqrt{1-y^2} + y\sqrt{1-x^2}$ 的最大值。
解:令 $x = \cos\theta, y = \cos\phi$。则 $\sqrt{1-y^2} = \sin\phi$ 且 $\sqrt{1-x^2} = \sin\theta$:
$\sin(\theta+\phi)$ 的最大值为1,这是可以达到的(例如,$x=1, y=0$)。
答案:1
问题 3: 代数
问题:对于多少个 $x$ 值,表达式 $\frac{x^2+1}{x^2-9}$ 未定义?
解:当分母为零时,表达式未定义,即 $x^2 - 9 = 0 \Rightarrow (x-3)(x+3) = 0$。因此 $x=3$ 或 $x=-3$。
答案:2