UniVer: A Unified Perspective for Multi-step and Multi-draft Speculative Decoding¶
会议: NeurIPS 2026
arXiv: 2605.04543
代码: https://github.com/W-Y-P/UniVer
领域: LLM 效率
关键词: speculative decoding, optimal transport, multi-step verification, multi-draft tree, inference acceleration
一句话总结¶
UniVer 将投机推理中的树状验证形式化为条件最优传输（Conditional OT）问题，利用前缀接受概率作为动态缩放因子统一了横向多候选与纵向多步验证，在保证严格无损分布的前提下显著提升了接收长度与吞吐量。
研究背景与动机¶
自回归大语言模型的推理延迟主要受限于显存带宽与逐词解码的串行瓶颈。投机解码（Speculative Decoding）通过引入轻量级小模型前向预测多个未来 token、再由目标大模型并行核验的“推测-验证”范式，有效突破了显存带宽约束。然而，如何设计能在数学上严格保持原始目标分布无损、同时最大化期望接受长度的验证策略，一直是投机解码算法层面的核心挑战。近期的加速方法主要沿两个正交方向推进：其一是横向多候选维度，如 SpecTr 与 Greedy 方法将多 draft 验证建模为最优传输（Optimal Transport, OT）问题，在单步内全局协同分配概率质量；其二是纵向多步维度，如 Block Verification 与 Traversal Verification 探索跨深度序列或树状结构的前缀累积概率与回退，减少过早截断惩罚。
现有方案的根本缺陷在于两者的割裂。以 Greedy 为代表的最优传输方法在纵向视界上具有短视性（vertically myopic），将树中每一层的展开看作相互孤立的生成事件，破坏了跨步概率的连续耦合；而 Traversal 等多步树状验证方法在横向选择上又极其受限（horizontally limited），同层兄弟节点间仅依赖无放回递归拒绝采样（RRSw）做局部逐项对比与残差重归一化，未能将最优传输扩展到多层树结构中。直接在整棵候选树上求解全局 OT 又面临路径数量随深度指数爆炸的组合复杂度难题。
本文的切入角度是发现无需在整棵树上强行求解不可计算的全局最优传输，纵向层级依赖完全可以通过前缀接受概率进行抽象降维。核心 idea：将前缀接受概率形式化为动态缩放因子，将整树验证转化为可自顶向下逐层递推、横向并行求解的条件最优传输（Conditional OT）闭式解，在无损保证下实现横向多候选协同与纵向跨步深度的联合最优验证。
方法详解¶
整体框架¶
UniVer 将树状投机解码的验证流程解耦为两个核心阶段：自顶向下的“概率分配阶段（Allocation Phase）”与自底向上的“后序决策阶段（Decision Phase）”。在分配阶段，根节点以前缀接受概率 \(\tilde{p}_{root}=1\) 起步，逐层向下推进。对于每一层中的所有分叉父节点，算法以该节点的前缀接受概率作为缩放因子，并行计算该节点与其子节点集合之间的条件最优传输计划，得到各子节点的边缘概率并归一化为条件转移概率；在决策阶段，依据预计算得到的接受概率与残差回退概率，对整棵树执行单遍后序遍历（Post-order Traversal）并结合均匀随机数进行确定性比对，命中即可早停采样。
整个流程无缝支持 GPU 上的批量矩阵运算，消除了传统 RRSw 严格依赖串行拒绝采样的依赖开销。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["候选树生成<br/>Top-(m-1) 贪心选择 + 1 采样残差"] --> B["自顶向下条件概率分配<br/>动态缩放因子结合目标分布"]
B --> C["层内并行闭式 OT 求解<br/>计算边缘概率与归一化条件概率"]
C --> D["单遍后序遍历决策<br/>叶节点接受或非叶节点残差采样"]
D --> E["无损生成序列输出<br/>返回接受前缀与后续 Token"]关键设计¶
1. 前缀概率作为动态缩放因子的条件最优传输：化解全局组合爆炸
以往多候选最优传输仅适用于单步孤立分布，无法感知节点在整棵树中的前置到达概率，导致跨层累积误差。UniVer 证明了只需将父节点 \(v\) 的前缀有效接受概率 \(\tilde{p}_v \in (0, 1]\) 作为动态缩放因子，将局部验证目标转化为缩放后的目标分布 \(\tilde{p}_v \cdot \mathcal{M}_b(\cdot \mid v)\) 与草稿分布 \(\mathcal{M}_s(\cdot \mid v)\) 之间的条件传输问题。为了在多候选设置下保证局部无损与闭式计算，子节点集合 \(\mathcal{C}(v)=\{u_1, \dots, u_m\}\) 采用混合采样构造：前 \(m-1\) 个候选从草稿分布中确定性选取概率最高的 \(\text{Top}_{m-1}\)，第 \(m\) 个候选从残差草稿分布 \(\mathcal{M}_s^\neg\) 中随机采样。
2. 闭式边缘接受概率与残差重归一化：并行解耦消除串行循环
传统 RRSw 必须按顺序依次比较候选词并在拒绝后重新归一化词表，无法充分利用 GPU 并行能力。UniVer 针对混合采样方案推导出了条件最优传输的显式闭式解。定义归一化常数： $\(Z_v = 1 - \tilde{p}_v + \sum_{x \in \Sigma} \left[ \tilde{p}_v \mathcal{M}_b(x \mid v) - \mathcal{M}_s^\neg(x \mid v) \right]_+\)$ 残差采样候选 \(u_m\) 的边缘接受概率直接由截断比例确定： $\(p_v(u_m) = \min\left\{ 1, \frac{\tilde{p}_v \mathcal{M}_b(u_m \mid v)}{\mathcal{M}_s^\neg(u_m \mid v)} \right\}\)$ 而其余前 \(m-1\) 个确定性候选 \(u_k\) 及未入选的残差项 \(u_{\text{res}} \in \Sigma \setminus \mathcal{C}(v)\) 的接受质量分配为： $\(p_v(u) = \left[ \tilde{p}_v \mathcal{M}_b(u \mid v) - \mathcal{M}_s^\neg(u \mid v) \right]_+ \times \frac{1 - p_v(u_m)}{Z_v}\)$ 整层所有分支节点的这一计算相互独立，可一次性张量化并行完成，随后直接转换为后序遍历所需的条件接受概率 \(\tilde{p}_{u_j}\) 与非叶节点回退接受概率 \(\tilde{p}_v^{\text{res}}\)。
3. 后序单遍遍历验证机制：纵向回退与严格无损保障
在完成整棵树的条件接受概率分配后，UniVer 执行后序遍历做最终序列决策。每个节点预先抽取独立均匀分布随机数 \(\eta \sim U(0, 1)\)。遍历遇到叶节点时，若 \(\eta_v < \tilde{p}_v\)，则整条从根到叶的路径全部被接受，并从 \(\mathcal{M}_b(\cdot \mid v)\) 采样下一个 token 结束验证；若为非叶节点，则检查是否满足残差接受条件 \(\eta_v < \tilde{p}_v^{\text{res}}\)，若命中则从归一化残差分布中采样后续词并早停返回。论文严格证明了只要树满足局部无损条件（即边缘接受期望等于缩放后的目标分布 \(\mathbb{E}_{\mathcal{C}(v)}[p_v(t)] = \tilde{p}_v \mathcal{M}_b(t \mid v)\)），两阶段验证算法产生的最终输出序列就与自回归目标模型分布 \(\mathcal{M}_b\) 达到严格的统计一致，杜绝了分布漂移。
损失函数 / 训练策略¶
UniVer 属于无需训练或微调的免训练（training-free）推断时验证算法。它完全兼容现有的任意草稿模型（如基于隐藏状态演化的 EAGLE、EAGLE-2 或独立小模型 Llama-68M）。在执行阶段直接利用草稿与目标模型输出的 logits 进行张量化闭式计算，额外显存占用与计算开销极低。
实验关键数据¶
主实验¶
在基准测试集 Spec-Bench（包含 MT-Bench、翻译、摘要、问答、数学和 RAG 六大领域）上，采用平衡二叉树（深度 5，共 32 个叶节点）配置进行评估，采样温度设定为 \(T=1.0\)。主要指标为每次验证循环平均生成的接受 token 数（\(\tau\)）以及端到端吞吐量（TPS）。
| 模型与方法 | MT-Bench | 翻译 (WMT14) | 摘要 (CNN/DM) | 问答 (NQ) | 数学 (GSM8K) | RAG (DPR) | 平均 \(\tau\) (提升幅度) | 平均吞吐 TPS |
|---|---|---|---|---|---|---|---|---|
| Vicuna-7B-v1.3 (基准自回归: 35.9 TPS) | ||||||||
| Vanilla RRSw | 3.42 | 2.73 | 2.97 | 2.72 | 3.47 | 3.05 | 3.06 (基准) | 67.5 |
| Traversal Verification | 3.49 | 2.75 | 3.03 | 2.79 | 3.55 | 3.15 | 3.13 (+2.3%) | 68.9 |
| Greedy (单步 OT) | 3.56 | 2.80 | 3.12 | 2.87 | 3.66 | 3.24 | 3.21 (+4.9%) | 70.4 |
| UniVer (本文) | 3.68 | 2.83 | 3.21 | 2.94 | 3.76 | 3.32 | 3.29 (+7.5%) | 72.2 |
| Llama3.1-8B-Instruct (基准自回归: 34.5 TPS) | ||||||||
| Vanilla RRSw | 2.96 | 2.31 | 2.56 | 2.33 | 3.31 | 2.86 | 2.72 (基准) | 48.5 |
| Traversal Verification | 3.05 | 2.34 | 2.60 | 2.40 | 3.37 | 2.94 | 2.78 (+2.2%) | 49.5 |
| Greedy (单步 OT) | 3.09 | 2.54 | 2.70 | 2.46 | 3.45 | 3.04 | 2.88 (+5.9%) | 50.6 |
| UniVer (本文) | 3.19 | 2.57 | 2.77 | 2.57 | 3.51 | 3.11 | 2.95 (+8.5%) | 52.2 |
消融实验：温度与树拓扑结构影响¶
论文在 Vicuna-7B-v1.3 上详尽测试了不同采样温度（\(T=0.3, 0.6, 1.0\)）与不同树拓扑（平衡二叉树 2 drafts \(\times\) 5 steps、宽浅树 3 drafts \(\times\) 3 steps、不规则 EAGLE Tree）对接受长度 \(\tau\) 的影响。
| 温度配置 | 验证方法 | 2 drafts \(\times\) 5 steps (深树) | 3 drafts \(\times\) 3 steps (宽树) | EAGLE 树拓扑 (非平衡) |
|---|---|---|---|---|
| \(T=0.3\) | Vanilla RRSw | 3.34 (0.0%) | 3.04 (0.0%) | 3.33 (0.0%) |
| Traversal | 3.39 (+1.5%) | 3.08 (+1.3%) | 3.38 (+1.5%) | |
| Greedy | 3.39 (+1.5%) | 3.08 (+1.3%) | 3.38 (+1.5%) | |
| UniVer | 3.39 (+1.5%) | 3.09 (+1.6%) | 3.38 (+1.5%) | |
| \(T=0.6\) | Vanilla RRSw | 3.29 (0.0%) | 3.02 (0.0%) | 3.29 (0.0%) |
| Traversal | 3.32 (+0.9%) | 3.05 (+1.0%) | 3.33 (+1.2%) | |
| Greedy | 3.34 (+1.5%) | 3.07 (+1.7%) | 3.34 (+1.5%) | |
| UniVer | 3.37 (+2.4%) | 3.09 (+2.3%) | 3.35 (+1.8%) | |
| \(T=1.0\) | Vanilla RRSw | 3.06 (0.0%) | 2.89 (0.0%) | 3.11 (0.0%) |
| Traversal | 3.13 (+2.3%) | 2.91 (+0.7%) | 3.16 (+1.6%) | |
| Greedy | 3.21 (+4.9%) | 3.00 (+3.8%) | 3.21 (+3.2%) | |
| UniVer | 3.29 (+7.5%) | 3.03 (+4.8%) | 3.24 (+4.2%) |
关键发现¶
- 深层网络中的质量累积效益更显著：在节点数量扩展实验中，随着二叉树深度由 2 增加到 5（节点数 7 增加至 63），UniVer 相对于 Greedy 与 RRSw 的性能优势单调递增。这是因为纵向层级越深，单步局部截断造成的概率衰减越严重，条件 OT 的自顶向下质量重配优势得到充分释放。
- 高温度分布更具区分度：在 \(T=0\) 或极低温（\(T=0.3\)）下，模型输出概率集中于 argmax，所有验证算法收敛到完全等价的确定性贪心匹配（\(\tau \approx 3.40\)）；而在高熵、高发散度的 \(T=1.0\) 下，UniVer 的接受长度提升最为明显（最高达 8.5%），展现出最优传输对离散概率空间重叠度的精准利用。
- 无求解器开销，吞吐收益高度保真：由于推导出了完全闭式的解，无需迭代调用 Sinkhorn 等通用 OT 数值求解器，UniVer 的额外计算延迟微秒级，平均接受长度提升（7.5%）几乎 1:1 转化为端到端实际吞吐提升（7.0%）。
亮点与洞察¶
- 将前缀概率视作动态标度因子的数学降维：直接解全树 OT 会导致路径状态指数爆炸，而 UniVer 洞察到树的条件概率结构可让前缀概率 \(\tilde{p}_v\) 充当缩放因子，从而把全树全局优化优雅拆解为局部条件最优传输的递归复合。
- 理论完备且天然免调优：给出了严格的分布无损证明（Local & General Losslessness），并在定理层面给出了弱于贪心 OT 的条件最优性界，不需要修改已有架构或训练任何超参数。
- 可迁移性极强：该设计原则可以平滑扩展至多智能体推测通信、多模态自回归验证或扩散模型的多步潜变量剪枝决策中。
局限与展望¶
- 零温场景无增益：当采样温度趋近于 0 时，采样空间退化为单分支确定性贪心，UniVer 退化为标准贪心核验，不带来额外吞吐收益。
- 与草稿采样策略的耦合性：当前的条件 OT 闭式解强依赖于 \(\text{Top}_{m-1} + 1 \text{ 采样}\) 的特定混合策略，对于任意全随机采样或自适应动态宽度树，闭式解推导需进一步泛化。
- 未来方向：探索联合优化草稿模型的动态拓扑生成算法、采样方式与纵向概率传输机制的软硬件协同流水线。
相关工作与启发¶
- vs Greedy (Hu et al., ICLR 2025): Greedy 首次推导了单步多候选的最优传输闭式解，但忽略了纵向多步树结构的依赖（每步独立割裂截断）；UniVer 将其推广到任意深度的树上，证明了接受长度在数学上严格不劣于 Greedy。
- vs Traversal Verification (Weng et al., NeurIPS 2025): Traversal 虽引入了后序遍历与纵向跨步概率累积，但横向兄弟节点间依然采用低效的 RRSw 串行拒绝采样；UniVer 用条件最优传输彻底替换了横向 RRSw，消除了串行依赖并实现了横纵双维度的联合最优。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙利用前缀接受概率作为缩放因子，将整树验证转化为闭式条件最优传输，理论切入点极其优雅。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 6 大任务领域、多种开源主干架构（Vicuna、Llama3.1、Qwen2、Llama2）与不同树拓扑、温度下的全方位横向对比。
- 写作质量: ⭐⭐⭐⭐⭐ 符号系统严密，定理表述规范，从问题形式化到证明链路层层递进。
- 价值: ⭐⭐⭐⭐⭐ 作为即插即用、严格无损的推断加速验证算法，对大模型服务系统落地具备直接工程与理论价值。