PARD-2：面向双模态投机解码的目标对齐并行草稿模型

# PARD-2：面向双模式投机解码的目标对齐并行草稿模型

来源：https://arxiv.org/html/2605.08632
Zihao An¹ Taichi Liu¹,² ¹footnotemark:1 Ziqiong Liu¹ Dong Li¹ Ruofeng Liu³ Emad Barsoum¹
¹Advanced Micro Devices, Inc. ²Rutgers University ³Michigan State University
\{Zihao.An, Taichi.Liu, Ziqiong.Liu, d.li, Emad.Barsoum\}@amd.com, [email protected]

###### 摘要

投机解码（Speculative Decoding）通过使用轻量级草稿模型提出候选 token，并由目标模型并行验证这些 token，从而加速大语言模型（LLM）的推理。然而，现有的草稿模型训练目标并未直接与推理时最大化连续 token 接受率的最终目标对齐。为了解决这一问题，我们重新制定了草稿模型的优化目标，将重点从单个 token 的预测准确率转移到整体的接受长度上。在本文中，我们在 PARD 的基础上提出了 PARD-2，这是一种带有置信度自适应 token（Confidence-Adaptive Token, CAT）优化的双模式投机解码框架。这种方法自适应地重新加权每个 token，以更好地与验证过程对齐。值得注意的是，PARD-2 使得单个草稿模型能够同时支持目标依赖（target-dependent）和目标独立（target-independent）模式。在多种模型和任务上的实验表明，PARD-2 实现了高达 6.94× 的无损加速，在 Llama3.1-8B 上比 EAGLE-3 快 1.9×，比 PARD 快 1.3×。我们的代码 available at https://github.com/AMD-AGI/PARD。

## 1 引言

随着大语言模型（LLM）的不断进步，其卓越的性能伴随着模型规模的快速增加。虽然这种缩放定律带来了显著的能力提升，但也使得自回归解码在推理时变得日益昂贵。

投机解码（Speculative Decoding, SD）\[17\] 最近作为一种降低 LLM 推理延迟的有效方法而出现。SD 使用轻量级草稿模型提出多个候选 token，然后由目标模型并行验证。一条有前景的研究路线是训练基于目标模型特征的轻量级自回归草稿器，包括 Medusa \[5\]、Hydra \[2\] 和 EAGLE-3 \[21\] 等方法，它们都取得了强劲的性能。然而，顺序草稿仍然需要多次前向传递，导致不可忽视的延迟 \[1, 30\]。为了进一步加速草稿过程，近期的工作探索了并行草稿以进一步降低草稿延迟：ParallelSpec \[30\] 训练并行草稿器在单次前向传递中生成多个 token，PARD \[1\] 将小型自回归模型适配为并行掩码 token 预测，而 DFlash \[7\] 则采用小型块扩散模型并行生成草稿 token。

> 参见标题 (a) Llama3.1-8B
> 参见标题 (b) Qwen3-8B

**图 1：vLLM 上的吞吐量和延迟权衡。** 在 (a) Llama-3.1-8B 和 (b) Qwen-3-8B 上，PARD-2 在各种批量大小（1 到 64）下始终实现更优的帕累托前沿。

然而，投机解码背后的一个常见假设是，所有草稿位置在训练时都应被同等对待，这对于训练收敛和接受长度来说是次优的 \[23, 7\]。与旨在均匀提高 token 预测准确率的标准化语言建模不同，投机解码最终关心的是目标模型能接受多少个草稿 token。我们的实验揭示了并行投机解码中的位置偏差：如图 2(a) 所示，后续草稿位置的 token 表现出 consistently 较低的接受率。随着草稿长度的增加，接受率往往难以维持，限制了并行草稿所能提供的实际加速效果。这一观察结果表明*均匀优化所有位置存在固有局限性*。虽然近期的方法如 DFlash \[7\] 和 DART \[23\] 通过位置感知的衰减权重来缓解这一问题，但它们的权重是固定的且主要依赖于位置。我们观察到，一个 token 的接受不仅由当前 token 的准确性决定，还严重受限于整个前缀的质量。这表明接受率是由当前 token 及其前缀上下文共同决定的。因此，一种同时考虑这两个因素的方法提供了提高接受长度和解码效率更有效的途径。

在本文中，我们介绍了 PARD-2，一种双模式投机解码框架，以缓解接受率的下降。我们提出了置信度自适应 token（CAT）优化，该优化分配 token 级别、依赖于上下文的置信度分数，以更好地将训练目标与投机解码中最大化连续 token 接受率的推理时目标对齐。具体而言，CAT 基于依赖于上下文的置信度分数动态重新加权 token 级目标，该分数计算为目标模型对前缀中所有先前 token 的置信度的累积乘积。这种设计鼓励草稿器最大化期望接受长度。

除了优化接受长度外，PARD-2 还解决了现有投机解码方法的目标依赖性问题。大多数投机解码方法是目标依赖的 \[21, 7\]，需要为每个目标模型从头训练一个新的草稿模型。建立在 PARD 基础之上，PARD-2 是首个允许单个草稿模型在推理期间动态切换目标依赖和目标独立模式的框架。与需要嫁接层的 EAGLE-3 和 DFlash 不同，PARD-2 保持独立架构，无需结构性开销即可实现这种灵活性。它在训练期间应用随机门控来控制目标隐藏状态的注入。因此，相同的草稿模型可以在目标依赖模式下运行以实现最大加速，同时也支持在一系列目标模型中泛化的目标独立模式。

总结一下，我们的主要贡献包括：

- 我们提出了 PARD-2，一个支持目标依赖和目标独立设置的双模式投机解码框架。据我们所知，这是首次在单个草稿模型中统一这些范式的工作。随机门控在训练期间注入目标隐藏状态，通过目标依赖优化实现峰值加速，同时保持对整个模型家族的通用兼容性。
- 我们重新审视了投机解码的基本目标，并证明其主要挑战是最大化连续 token 序列的接受率。为此，我们提出了一种新颖的优化策略 CAT。基于先前前缀，CAT 在目标模型依赖于上下文的置信度分数的指导下，自适应地重新加权其对单个 token 的关注度，从而显著提高预测和蒸馏效率。
- 我们在各种模型和基准测试上进行了广泛的实验，包括在 vLLM 框架中对 PARD-2 的实际验证。我们的结果显示，PARD-2 相比 PARD 平均加速 1.3×，相比自回归基线最高加速 6.94×。此外，它在高并发设置下提供最高的吞吐量，展现了极高的实际部署价值。

## 2 预备知识

### 2.1 投机解码

投机解码是一种用于加速 LLM 推理的无损解码策略。它不再仅仅使用目标模型 $\boldsymbol{\theta}_{\mathrm{target}}$ 生成每个 token，而是引入一个更小更快的草稿模型 $\boldsymbol{\theta}_{\mathrm{draft}}$ 提前提出多个候选 token，然后由目标模型并行验证。这种设计减少了昂贵的目标模型解码步骤的数量，同时保留了目标模型的精确输出分布。

> 参见标题 (a) 逐位置接受率和接受长度
> 参见标题 (b) 目标模型的置信度与接受率

**图 2：Llama3.1-8B 的接受行为。** (a) 在 HumanEval 基准测试上，PARD-2 在各个 token 位置上实现了比 PARD 更高的接受率和更长的接受长度，缓解了远端位置的性能下降。(b) 目标模型的置信度分数与实际接受率高度相关，支持将其作为 token 级接受的代理。

形式上，给定前缀 $X=(x_0,\ldots,x_{n-1})$，投机采样使用轻量级自回归草稿模型 $\boldsymbol{\theta}_{\mathrm{draft}}$ 提议长度为 $K$ 的 token 序列，记为 $\tilde{Y}=(\tilde{y}_n,\ldots,\tilde{y}_{n+K-1})$。提议概率分布分解为

$$
P(\tilde{Y}\mid X;\boldsymbol{\theta}_{\mathrm{draft}})=\prod_{k=0}^{K-1}P\!\left(\tilde{y}_{n+k}\mid x_0,\ldots,x_{n-1},\tilde{y}_n,\ldots,\tilde{y}_{n+k-1};\boldsymbol{\theta}_{\mathrm{draft}}\right). \tag{1}
$$

对于位置 $n+k$，令 $p_k(y)=P(y\mid x_0,\ldots,x_{n-1},\tilde{y}_n,\ldots,\tilde{y}_{n+k-1};\boldsymbol{\theta}_{\mathrm{target}})$ 和 $q_k(y)=P(y\mid x_0,\ldots,x_{n-1},\tilde{y}_n,\ldots,\tilde{y}_{n+k-1};\boldsymbol{\theta}_{\mathrm{draft}})$ 分别表示目标和草稿条件概率。在投机采样下，草稿 token $\tilde{y}_{n+k}$ 被接受的概率为

$$
a_k=\min\!\left(1,\,\frac{p_k(\tilde{y}_{n+k})}{q_k(\tilde{y}_{n+k})}\right),\qquad k=0,\ldots,K-1. \tag{2}
$$

忽略奖励 token，前 $k+1$ 个草稿 token 全部被接受的概率为 $\prod_{j=0}^{k}a_j$。因此，期望接受长度 $L$ 为

$$
\mathbb{E}[L\mid X,\tilde{Y}]=\sum_{k=0}^{K-1}\prod_{j=0}^{k}a_j. \tag{3}
$$

目标模型接受最长的有效前缀，并在第一次拒绝时从剩余分布中采样一个修正 token，从而保持与从目标模型采样的精确等价性。

### 2.2 并行草稿模型

虽然投机解码显著加速了 LLM 推理，但其草稿阶段仍然是顺序的，需要 $K$ 次顺序依赖的预测来生成 $K$ 个草稿 token。这种顺序延迟仍然可能限制端到端的加速效果。为了解决这个问题，最近的工作探索了并行预测多个 token 的并行草稿模型。DiffuSpec \[18\] 和 DFlash \[7\] 采用基于扩散的草稿器，通过迭代去噪生成 token。为了更好地匹配目标模型的自回归架构，PARD \[1\] 保留了自回归主干并引入了掩码占位符，能够在单次前向传递中进行并行掩码 token 预测。

特别是，PARD 引入了一种特殊的掩码 token $m$，并且每个未来 token 的预测仅依赖于前缀和前面的掩码占位符。其草稿概率分布为

$$
P(\tilde{Y}\mid X;\boldsymbol{\theta}_{\mathrm{PARD}})=\prod_{k=0}^{K-1}P\!\left(\tilde{y}_{n+k}\mid x_0,\ldots,x_{n-1},m_n,\ldots,m_{n+k-1};\boldsymbol{\theta}_{\mathrm{PARD}}\right). \tag{4}
$$

由于每个位置仅依赖于前缀和掩码 token，所有 $K$ 个预测都可以在单次前向传递中计算。这种方法不仅大幅降低了草稿延迟，还确保了目标独立性，使得草稿器可以在一系列目标模型中复用。

给定真实值 $Y=(y_n,\ldots,y_{n+K-1})$，PARD 使用交叉熵损失进行训练

$$
\mathcal{L}_{\mathrm{PARD}}=-\frac{1}{K}\sum_{k=0}^{K-1}\log P\!\left(y_{n+k}\mid x_0,\ldots,x_{n-1},m_n,\ldots,m_{n+k-1};\boldsymbol{\theta}_{\mathrm{PARD}}\right). \tag{5}
$$

> 参见标题 **图 3：PARD-2 概览。** PARD-2 的训练（中）和推理（右）设计。与 PARD（左）相比，PARD-2 集成了 CAT 优化、目标隐藏特征和知识蒸馏。PARD-2 支持在目标依赖和目标独立模式之间灵活切换。

## 3 方法

### 3.1 观察

草稿长度 $K$ 是并行草稿模型的一个关键设计选择。为了研究其影响，我们使用两种草稿长度 $K=8$ 和 $K=16$ 训练 PARD。如表 5 所示，增加 $K$ 几乎没有带来改进，甚至在几个基准测试上性能下降。这一观察结果与“更长的草稿长度自然意味着更大的接受长度和增强的解码效率”的常见直觉相矛盾。为了理解这一现象，我们分析了投机解码的验证机制。因为目标模型严格按顺序评估候选 token，任何后续 token 的接受都 heavily 依赖于所有先前 token 的成功验证。令 $a_j$ 表示目标模型接受第 $j$ 个草稿 token 的边际概率。式 (3) 可以按位置分解为

$$
\mathbb{E}[L\mid X,\tilde{Y}]=\sum_{k=0}^{K-1}\prod_{j=0}^{k}a_j=\sum_{k=0}^{K-1}\left(\prod_{j=0}^{k-1}a_j\right)a_k. \tag{6}
$$

这种分解揭示了决定位置 $k$ 的 token 是否被接受的两个关键因素。第一个因素，$\prod_{j=0}^{k-1}a_j$，是所有先前草稿 token 被接受的概率。第二个因素 $a_k$ 是到达位置 $k$ 后当前 token 被接受的概率。

然后我们将第一个因素定义为 $s_k$：

$$
s_k:=\prod_{j=0}^{k-1}a_j,\qquad s_0:=1. \tag{7}
$$

项 $s_k$ 是第 $k$ 个 token 对接受长度做出贡献的先决条件，因此可以解释为该 token 对加速的重要性。使用这种表示法，

$$
\mathbb{E}[L\mid X,\tilde{Y}]=\sum_{k=0}^{K-1}s_k a_k. \tag{8}
$$

第二个因素，$a_k$，反映了位置 $k$ 处草稿预测的局部质量。由于 token 级的训练目标...