EntMTP：利用熵引导的多令牌预测加速大语言模型推理

# 基于熵引导的多token预测加速LLM推理

来源：https://arxiv.org/html/2606.27550

###### 摘要

多token预测已被证明能在训练期间增加数据密度、提升下游文本生成质量，并成为自推测解码的事实标准。使用MTP头的现有基础模型和开源模型在整个生成序列中采用固定的基于树的注意力拓扑结构，这意味着推测深度（以及验证所需计算量）保持恒定，与上下文无关。这与自然语言的熵模式根本不符：低熵区域通常支持可靠的多步草稿，而高熵区域需要更保守的推测。为解决这一问题，我们提出熵引导多token预测（EntMTP），一种无需训练的计划器，它根据对局部生成熵的运行估计，在一组任务特定的帕累托最优树中选择基于树的注意力拓扑。通过将推测深度与上下文可预测性匹配，EntMTP能在不牺牲生成质量的情况下，最大化整个生成文本分布上的预期接受token吞吐量。在Humaneval、ShareGPT、GSM8k和Litbench基准测试中，EntMTP相比Hydra始终实现1.09-1.15倍加速，峰值加速达到约1.36倍（相比Medusa基线）。机器学习，ICML

Carrie Chen  
康奈尔大学  
[email protected]

## 1 引言

大型语言模型（LLM）在广泛任务上取得了惊人表现，然而自回归解码的顺序性质仍是其延迟瓶颈：每个生成token都需要完整前向传播所有模型参数，这使得推理在规模上受限于内存带宽且成本高昂。推测解码（Yaniv等人，2023；Chen等人，2023）通过使用廉价提案模型草拟k个候选token，并在目标模型的单次并行前向传播中验证它们来解决这一问题，同时精确保持目标分布。多token预测（MTP）头（Cai等人，2024；Ankner等人，2024）通过在目标模型最终隐藏状态上附加轻量模块，生成候选token树用于稀疏验证，从而无需单独草稿模型即可实现推测解码。

对于基于MTP的解码，性能主要取决于每一步实现的*接受深度*。先前工作（Cai等人，2024；Ankner等人，2024）使用这一指标在Alpaca子样本（Wang等人，2023）上离线选择草稿树拓扑；该选定拓扑在整个推理过程中重复使用。我们在附录A中表明这种全局选择过于粗糙。草稿token的接受情况随位置、局部上下文以及与训练分布的重叠程度而变化。更重要的是，接受行为和树拓扑上的吞吐量前沿都是**任务特定的**：在一个基准上帕累托最优的拓扑在另一个基准上可能高度次优。

我们提出EntMTP，一种轻量运行时树选择层，利用这种任务依赖性而非将其视为噪声。EntMTP首先离线构建Hydra树拓扑的吞吐量-帕累托前沿，使用贪心扩展和自推演接受估计。它仅保留改善成本-吞吐量权衡的前沿树，丢弃被支配拓扑。在推理时，预编译的TopologyBank为每个前沿树存储注意力掩码、位置偏移和收集索引。因此切换拓扑简化为O(1)指针交换，无需重构掩码或重新构建内核。在每个解码步骤，调度器从模型读取现成特征（例如EAGLE-2路径值，或基础top-1概率），并使用三种策略之一选择树：EntMTP-ll（阈值阶梯）、EntMTP∗（静态任务最优树规则）或EntMTPτ（带独立开启/关闭阈值的滞后二元τ规则，防止每步抖动）。

在batch_size=1设置下，EntMTPτ在所有三个评估任务（小学数学、代码补全、对话式LLM）上均优于Hydra和Medusa基线以及EntMTP∗。在GSM8K上，它达到112.0 tok/s，而Hydra为102.4 tok/s（+9.4%），加速比为自回归解码的3.13倍。在HumanEval上，它达到124.7 tok/s（比Hydra高14.0%，比AR高3.26倍）。最大绝对增益出现在ShareGPT：117.5 tok/s对比109.0 tok/s（+7.8%，AR的3.47倍）。所有设置下响应困惑度保持在Hydra的0.022以内，确认增益来自调度，而非验证规则的改变。

## 2 相关工作

### 2.1 推测解码

推测解码（Yaniv等人，2023；Chen等人，2023）通过交替进行低成本的*草稿*阶段和在目标模型上并行的*验证*阶段，无损加速自回归生成。给定前缀T₁:ⱼ，草稿分布q生成T̂ⱼ₊₁:ⱼₖ及其每token概率p̂ⱼ₊ᵢ，然后目标分布p通过单次前向传播评估整个候选序列。token从左到右以概率min(1, pⱼ₊ᵢ(t̂ⱼ₊ᵢ)/p̂ⱼ₊ᵢ(t̂ⱼ₊ᵢ))被接受；在第一次拒绝时，从残差norm(max(0, pⱼ₊ᵢ−p̂ⱼ₊ᵢ))中抽取替换token，并丢弃后缀。这确保接受token的边缘分布与目标分布匹配，因此实际加速完全取决于q对p的模仿程度以及验证的可并行化程度。

### 2.2 EAGLE和EAGLE-2

EAGLE（Li等人，2024b）用特征级别的自回归取代了token级别的草稿。给定目标模型的隐藏状态序列F₁:ᵢ和token序列T₁:ᵢ，一个轻量草稿器预测下一个隐藏状态F̂ᵢ₊₁ = D_θ(T₂:ᵢ₊₁, F₁:ᵢ)，然后由目标模型的冻结LM头解码：p_θ(tᵢ₊₁|t≤ᵢ) = softmax(W_LM F̂ᵢ₊₁)。草稿器训练以匹配目标模型的下一个隐藏状态，使用特征回归损失：L_reg = SmoothL1(Fᵢ₊₁, D_θ(T₂:ᵢ₊₁, F₁:ᵢ))。通过在隐藏状态空间中进行草稿而非直接预测独立未来token，EAGLE在类似参数成本下获得了更强的草稿器。

EAGLE-2（Li等人，2024a）保留了特征级草稿器，但使草稿树变为动态。它不是固定拓扑，而是使用草稿器自身的置信度对每个候选路径进行评分。对于草稿树中的节点tᵢ，其路径值定义为Vᵢ = ∏_{tⱼ∈Path(root, tᵢ)} pⱼ ≈ ∏_{tⱼ∈Path(root, tᵢ)} cⱼ，其中pⱼ是节点tⱼ处真实但未观测的接受概率，cⱼ是草稿器的局部置信度。推理时，EAGLE-2通过重复选择最高价值的前沿节点自上而下展开树：t* = argmax_{tᵢ∈F} Vᵢ，其中F表示当前前沿。由于评分cⱼ已由草稿器产生，这种动态树构建无需额外训练组件，并保持无损的目标模型验证。

### 2.3 独立多头预测

Medusa（Cai等人，2024）采用更直接的推测方法：在目标模型的最终隐藏状态上附加K个轻量预测头。每个头hᵢ独立于其他草稿token预测偏移i处的token，将草稿分布分解为∏i p_hᵢ(tⱼ₊ᵢ|hⱼ)。这些头发射top-kᵢ候选，其笛卡尔积被排列成草稿树，带有预计算的注意力掩码、位置偏移和收集索引。验证随后简化为目标模型的一次稀疏前向传播。

### 2.4 顺序多头预测

Hydra（Ankner等人，2024）保留了Medusa的基于隐藏状态的头设计，但恢复了草稿头之间的顺序依赖性。每个头hᵢ不是独立预测每个未来token，而是以较早头已草稿的token嵌入为条件，建模p(tⱼ₊ᵢ|t≤ⱼ, t̂ⱼ₊₁:ⱼ₊ᵢ₋₁)而非边缘分布p(tⱼ₊ᵢ|t≤ⱼ)。每个头实现为一个小的因果Transformer块，作用于基础隐藏状态和先前草稿的token嵌入。这保留了O(K)的头成本，同时允许草稿链以其自身前缀为条件。然而，Hydra继承了Medusa的静态树限制：一个单一拓扑离线选择，并在每个解码步骤重复使用。

## 3 优化任务特定的贪心草稿树

Hydra的接受行为是工作负载相关的，其变化方式验证器和基础模型无法吸收。在ShareGPT上，仅凭最近的接受历史就能预测下一个接受长度（EMA r=0.49）；在GSM8K和HumanEval上，最强的预测因子基于熵，且至少弱2倍（|r|≤0.22）（图3；完整表格见附录EntMTP：加速LLM推理的熵引导多token预测）。由于验证器和基础权重在各次运行中共享，这种离散性是工作负载诱导的token分布的特性，意味着当不同分布的接受信号相差超过2倍时，单一草稿拓扑不可能最优。因此，我们按照Ankner等人（2024）和Cai等人（2024）的两阶段离线过程，为每个任务优化一个单独的草稿树。

### 3.1 构建接受前沿

首先，通过贪心方式构建一系列规模递增的提案树T₁,...,T_N：

算法1 接受前沿的贪心构建

输入：最大深度D，预算B，提示P，验证器V  
1: T ← {(0)}, H ← []  
2: while |T| < B do  
3:   候选集C ← ∅  
4:   for 每条前缀链p ∈ 深度≤D-1的叶节点 do  
5:       扩展p -> 新节点n（以所有词汇扩展，成本为增广树的前向传播）  
6:       计算接受率r ← E_[P][accept_length(T ∪ {n})]  
7:       C ← C ∪ {(n, r)}  
8:   end for  
9:   选择n* = argmax_{n} r_n  
10:  T ← T ∪ {n*}  
11:  end while  
12: 返回 T

然后保留那些在（平均节点数，平均接受长度）空间中位于帕累托前沿的树。这与Ankner等人（2024）的过程相同，但我们在每个任务自己的验证集上运行它，而不是使用通用校准集。

### 3.2 贪心树与全局拓扑的比较

图1显示了在ShareGPT、GSM8K和HumanEval上，单个全局树（来自Ankner等人（2024）的默认Hydra树）与各自任务特定前沿的平均接受长度和节点数。在每个任务中，前沿包含一个比全局树更高效（更少节点，相等或更高接受长度）的树。例如，在GSM8K上，全局树有63个节点和τ=2.51，而前沿树有47个节点和τ=2.39（更少节点，相似接受长度）。在HumanEval上，全局树有63个节点和τ=3.06，前沿树有28个节点和τ=3.28（更少节点，更高接受长度）。

更重要的是，前沿的形状因任务而异：ShareGPT倾向于窄而深的树，HumanEval倾向于宽而浅的树，GSM8K介于两者之间。这证实了任务特定的优化是必要的：没有单个全局拓扑能匹配所有分布。

## 4 运行时树选择

给定任务特定的Pareto前沿，我们设计一个轻量调度器，在推理期间根据局部生成熵在各拓扑之间切换。

### 4.1 熵估计

我们使用归一化熵作为可预测性的代理。设p为验证器logits上的softmax分布。对于草稿树中的节点tᵢ，其局部熵定义为H(tᵢ) = -∑_v p(tᵢ=v) log p(tᵢ=v)。然后我们通过取最后D个已接受token的熵的指数移动平均值来估计运行熵s：

s ← β · s + (1-β) · H(t_当前)

β是一个平滑因子（我们使用0.9）。熵低表示确定性强，适合深层推测；熵高表示不确定性高，应使用浅树。

### 4.2 调度策略

我们提出三种调度策略，均基于运行熵s：

- **EntMTP∗**：选择整个任务中吞吐量最优的静态树（即前沿中平均接受长度/节点数比率最高的树）。这等价于对所有步骤使用单一树，但针对任务优化。
- **EntMTPτ**：使用带滞后的二元τ规则。当s ≤ τ_off时使用保守树T_-，当s ≥ τ_on时使用激进树T_+，中间状态使用前一个树（若s上升则保持不变直到s ≥ τ_on，若s下降则保持不变直到s ≤ τ_off）。设置τ_off ≤ τ_on以防止抖动。默认使用单一阈值(τ_on=τ_off=τ)。K路推广threshold_ladder将s划分为K个连续区间，选择T₁,...,T_K。两种策略都无需训练：τ通过在用于树搜索的相同100提示校准集上进行一维扫描选择。

##### 成本

将π插入解码循环基本免费。B中的每个树都附带其预计算的generate_hydra_buffers输出（注意力掩码、位置ID、检索索引），因此切换树就是一个字典指针交换，无需GPU工作。从已在寄存器中的验证器logits计算s只需对D≤4个标量进行一次累积乘积和一次.item()同步，每步增加<0.1ms。

## 5 评估方法

### 5.1 指标

EntMTP既不对原始LLM微调，也不放松Hydra典型的接受条件，因此它是一种无损加速方法（在所有运行中，延续困惑度与基础LM相差在0.02 nats内）。因此我们仅评估加速，使用两个指标：
- **加速比ρ**：相对于同一基础LM在相同提示集上标准自回归解码的每秒墙上时钟输出token数（包括提示预填充在计时区域内）。
- **平均接受长度τ**：每个草稿-验证周期生成token的期望数量。τ独立于硬件和运行时，将草稿拓扑质量与内核级效应分离。

### 5.2 设置

我们使用Vicuna-7B v1.3（Chiang等人，2023）作为基础LM，ankner/hydra-vicuna-7b-v1.3（Ankner等人，2024）作为Hydra验证器，FP16精度在单个NVIDIA A100上运行。所有方法共享温度T=0.7，后验阈值ε=0.09，混合系数α=0.3，最大输入1400个token，最大生成256个token。每个基准在从HumanEval-val、GSM8K-val和ShareGPT（Vicuna未过滤分割）抽取的100个提示（种子123）上计时。每次运行以一次热身生成开始以摊销JIT和KV缓存分配；计时不包括热身。调度器的保守树和激进树从每任务吞吐量前沿中选择；τ在{0.001, 0.005, 0.01, 0.02, 0.05}上在校准集选择。

## 6 结果

以下报告两种EntMTP策略在三个基准上的加速比和平均接受长度。

表1：相对于普通Vicuna-7B的加速比ρ和平均接受长度τ，在T=0.7下100提示集上。Hydra和Medusa使用其作者发布的默认树。EntMTP∗使用每任务吞吐量最优树；EntMTPτ每步在来自同一前沿的保守树和激进树之间切换。所有时间包括提示预填充。

固定的每任务吞吐量最优树（EntMTP∗）在token/s上比Hydra发布默认值高出7.1%-13.2%，同时使用≥2倍更少的草稿节点（HumanEval/GSM8K/ShareGPT上分别为28/46/30 vs. 63），并且比Medusa默认值高出7.7%-32.4%。添加运行时调度器（EntMTPτ）在EntMTP∗基础上进一步获得0.5%-2.1%的提升，并在每个基准上取得最佳墙上时钟吞吐量：在HumanEval上为普通Vicuna的3.26倍，在GSM8K上为3.13倍，在ShareGPT上为3.47倍。每一行的延续困惑度保持在基础LM的0.02 nats以内，因此增益是无损的。

### 6.1 增益来源

相对于Hydra默认值的静态增益在不同工作负载中一致分解：7%-13%的大部分来自更小的每任务树缩减了每步验证器成本，0-7%来自优化拓扑上更高的τ（HumanEval两者都涨，τ从3.06→3.28；ShareGPT牺牲τ降低3%以换取更小的树，仍净增+7% tok/s）。调度器相对于静态树的残余增益在GSM8K上最大（+2.1%），而在ShareGPT上最小（+0.5%）。这符合预期：GSM8K具有最高熵变化性，因此运行时切换最有益；ShareGPT更均匀，因此静态树已接近最优。

为了隔离调度器的影响，我们还测试了EntMTP∗（静态任务最优树）并报告其性能。在所有情况下，调度器都增加了价值，但最大增益出现在具有最不均匀熵分布的任务上。这证实了运行时调度器是有用的，并且无论基础树选择如何，其增益都是额外的。