RheoSampling: Resolving the One-Hot Dilemma in Stochastic Dynamic-Tree Speculative Decoding¶
会议: NeurIPS 2026
arXiv: 2609.21827
OpenReview: https://openreview.net/forum?id=PgNNLmDUZe
领域: LLM 效率
关键词: 投机微调与解码, 动态树推测解码, 随机采样, 最优传输, 无损加速
一句话总结¶
RheoSampling 针对随机解码下动态树推测解码将投机分布退化为 one-hot 概率导致接受率骤降的难题，提出将树构建代理概率与验证真实采样概率解耦的双重身份机制，结合最优传输验证算法与稀疏截断，在严格保持数学无损性的前提下实现动态树结构与随机探索的双赢。
研究背景与动机¶
投机解码（Speculative Decoding）通过轻量草稿模型并行推测多个未来 token、再由目标大模型单次前向验证，已成为打破自回归解码显存带宽瓶颈的主流范式。为了在每轮验证中捕获更多潜在的高概率分支，前沿方法逐步从单链推测转向树状结构。以 EAGLE-2/3 为代表的动态树方案，凭借“扩展-重排序-全局剪枝”（expand-then-rerank）机制，能够根据上下文自适应生成紧凑的高质量候选拓扑，在确定性贪心解码（\(T=0\)）下取得了显著的加速比。
然而，一旦进入具有随机性的温度采样场景（\(T>0\)），现有的动态树范式便遭遇了根本性障碍。因为动态树在逐层扩展时纯粹依赖 top-\(K\) 确定性筛选，直接将草稿模型预测的连续概率分布硬截断为离散点质量（point masses）；在后续验证中，草稿概率被直接当作 1.0 的 one-hot 概率处理，导致目标分布中长尾 token 的有效探索空间被完全抹杀，接受率产生断崖式下跌。这暴露出了一个尖锐的“one-hot 困境”：动态树为了维持高质量的上下文感知拓扑而牺牲了随机探索，而静态树虽然能进行随机采样却受困于僵化的预设结构。
这一困境的深层症结在于，现存框架强行让同一个概率分布同时承担“指导树扩展与剪枝”和“计算验证接受率”这两项相互冲突的任务。如果简单地用随机采样替换 top-\(K\) 候选，被采样 token 的真实概率会直接参与路径评分，导致其能否在全局剪枝中存活强依赖于采样出的 token 身份；这种条件存活偏置直接扭曲了留存节点的后验分布，从根本上摧毁了推测解码的无损性（Losslessness）。核心 idea：将采样 token 的角色一分为二，在树构建与剪枝阶段赋予其恒定的代理概率（Proxy Probability）以锚定拓扑排位，在验证阶段回归真实的条件采样概率，从而在解耦中首次实现上下文动态树与随机采样的无损统一。
方法详解¶
整体框架¶
RheoSampling 的整体运行流程分为三个阶段：轻量草稿截断、双重身份动态树构建与基于最优传输的逐层验证。在每一轮推测中，草稿模型先对当前候选节点生成预测 logits，经过稀疏截断后进入局部扩展；在每个父节点的扩展槽位中，确定性地保留 top-\(m\) 个最高概率 token，并在剩余长尾中随机抽取 1 个探索 token，同时赋予其不依赖具体 token 内容的代理概率；整棵候选树沿路径累乘代理分数并经受全局剪枝；最后，目标模型在单次前向中并行验证剪枝后的动态树，通过最优传输验证策略优先检验随机 token 并递归处理确定性候选。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["草稿模型 Logits"] --> B["稀疏投机分布截断<br/>保留 Top-128 截断尾部"]
B --> C["双重身份解耦采样<br/>Top-m 确定性 + 1 个带代理概率的采样 Token"]
C --> D["全局 Top-R 剪枝<br/>基于代理路径得分确定树拓扑"]
D --> E["最优传输验证机制<br/>优先验证随机探针，残差递归分配"]
E --> F["无损输出 Token 序列"]关键设计¶
1. 稀疏投机分布截断：大幅削减大词表采样计算开销 标准大语言模型词表通常高达数万至十余万，在每个动态树扩展节点上直接对全词表执行 softmax 并进行无放回多项分布采样，会引入不可忽视的内核启动与显存搬运延迟。为了将理论算法转化为实际的端到端墙钟加速，RheoSampling 在 softmax 之前对原始 logits 进行 top-128 硬截断，将其余位置强制置为 \(-\infty\)。由于大模型预测在大部分场景下集中度极高，top-128 能够覆盖原始草稿分布 91% 至 96% 以上的概率质量；更为关键的是，只要目标模型验证阶段采用完全一致的截断支持集，根据重要性采样原理，该截断操作严格保持数学无损性，同时将草稿端采样延迟完全抹平。
2. 双重身份解耦采样：破除剪枝对采样分布的条件扭曲 为了打破“直接随机采样会导致剪枝存活偏置、进而破坏无损性”的技术死锁，RheoSampling 在每个含有 \(K\) 个槽位的候选池中推行混合分配机制。首先确定性选取草稿概率最高的 \(m\) 个 token 作为骨干 \(\{x_1, \dots, x_m\}\)；随后计算剩余尾部概率总和 \(z = 1 - \sum_{i=1}^m q(x_i)\)，并从重归一化后的残差分布 \(\tilde{q}(Y) = q(Y)/z\) 中随机抽取单个探针 token \(x_s\)；剩余的 \(K - m - 1\) 个空位由余下词表中概率最高的填充 token 补齐。
该设计的核心突破在于赋予 \(x_s\) 两个截然分开的概率身份。在树的自适应生长、路径累乘与全局 top-\(R\) 剪枝阶段，\(x_s\) 携带固定的代理概率： $\(q_{\text{proxy}}(x_s) = \min\{q(x_m), z\}\)$ 当 \(m=0\) 时代理概率设为 \(q(x_1) + \epsilon\)。根据引理（Local Ranking Preservation），代理概率必然满足 \(q_{\text{proxy}}(x_s) \ge q(x_{m+1})\)。这意味着 \(x_s\) 在同级兄弟节点中的排位恒定紧跟在第 \(m\) 个确定性节点之后、且严格高于所有填充节点，其在全局剪枝中的去留完全由固定代理分数决定，与具体采样出哪一个 token 彻底解耦。而在后续验证时，该节点立即恢复其真实的条件采样概率 \(\tilde{q}(x_s)\)，使得条件后验不发生任何畸变。参数 \(m\) 如同“变阻器”（Rheostat），\(m\) 较小时代理概率更大、\(x_s\) 存活率更高但确定性主干受抑，\(m\) 较大时拓扑质量更高但探针易被剪枝，\(m=1\) 在标准场景下构成了最佳性能平衡点。
3. 最优传输验证机制：先验随机候选以最大化长尾重叠 在剪枝后的候选树逐层验证中，若该层保留了随机采样的探针 token \(u_s\)（对应索引 \(s \ne -1\)），传统的无放回递归拒绝采样（RRSw）会按照固定顺序依次验证确定性 token，最后才验证随机 token，导致随机 token 的分布重叠优势被确定性前缀的拒绝概率 \((1 - \sum_{v \in \text{Top}_m} p(v))\) 严重缩减。RheoSampling 从最优传输（Optimal Transport）视角重新设计了验证流程：优先独立验证随机 token \(u_s\)： $\(\alpha(u_s) = \min\left(1, \frac{p(u_s)}{\tilde{q}(u_s)}\right)\)$ 若 \(u_s\) 被接受，直接沿该分支深入；若被拒绝，则构建目标分布与草稿尾部的残差分布 \(r = \text{norm}(\max(0, p - \tilde{q}))\)，再在 \(r\) 之下对其余候选执行递归拒绝检验；若候选全部拒绝，则从最终残差 \(r\) 中采样输出。理论分析与严格证明表明，这种“随机探针先行”的调度策略完整保留了全量长尾重叠 \(\min(p(v), \tilde{q}(v))\)，在数学上严格支配了 RRSw，显著提高了单步期望接受长度。
损失函数 / 训练策略¶
RheoSampling 属于纯测试时推测解码策略（Training-free / Inference-time Method），完全不需要对草稿模型或目标模型进行任何参数微调或重训。在实验评估中，草稿模型直接加载官方发布的预训练 EAGLE-3 权重，草稿树的最大容量设定为 \(R=60\)，最大展开深度为 8。值得强调的策略细节是：动态树的构建与剪枝过程始终基于原始未带温度的草稿概率，仅在抽取随机探针与目标模型验证时引入温度缩放 \(T>0\)，这一策略有效防止了低温度下排序信息的过度退化。
实验关键数据¶
主实验¶
在 NVIDIA A6000 硬件环境下，以默认温度 \(T=1.0\)、树大小 60、深度 8 为基准设置，涵盖 Alpaca、GSM8K（Math）、HumanEval（Code）、MT-bench（MT）、Natural Questions（QA）和 CNN/DailyMail（Sum）六大评测基准，比较 RheoSampling 与原始确定性 Top-\(K\)（EAGLE-3）的平均接受长度 \(\tau\) 及相对于自回归基线的端到端墙钟加速比。
| 目标模型 | 解码策略 | Alpaca (\(\tau\)) | Math (\(\tau\)) | Code (\(\tau\)) | MT (\(\tau\)) | QA (\(\tau\)) | Sum (\(\tau\)) | 平均 \(\tau\) | 加速比 |
|---|---|---|---|---|---|---|---|---|---|
| Vicuna-13B | Top-\(K\) | 5.59±0.09 | 5.85±0.09 | 6.66±0.13 | 5.69±0.06 | 4.91±0.08 | 5.80±0.03 | 5.75±0.03 | 3.37× |
| Vicuna-13B | RheoSampling | 5.92±0.05 | 5.87±0.05 | 6.72±0.06 | 5.88±0.09 | 5.00±0.04 | 5.94±0.07 | 5.89±0.03 | 3.43× |
| Llama-3.1-8B | Top-\(K\) | 5.68±0.02 | 5.50±0.03 | 6.11±0.09 | 4.60±0.06 | 3.89±0.04 | 4.49±0.07 | 5.04±0.02 | 2.84× |
| Llama-3.1-8B | RheoSampling | 6.03±0.12 | 5.60±0.08 | 6.23±0.03 | 4.84±0.14 | 4.14±0.04 | 4.66±0.07 | 5.25±0.03 | 2.93× |
| DeepSeek-R1-Distill-8B | Top-\(K\) | 4.48±0.04 | 6.70±0.05 | 5.43±0.09 | 4.93±0.02 | 4.08±0.05 | 4.32±0.06 | 4.99±0.01 | 2.89× |
| DeepSeek-R1-Distill-8B | RheoSampling | 4.79±0.06 | 6.71±0.07 | 5.79±0.05 | 5.16±0.05 | 4.35±0.05 | 4.47±0.06 | 5.21±0.02 | 2.99× |
消融实验¶
消融实验包含两部分：一是在 MT-bench 上针对稀疏投机分布截断大小（Support Size）对覆盖率、接受长度 \(\tau\) 与单步草稿延迟的影响；二是针对单层受控实验下不同验证策略的接受率对比（验证 OT 理论增益）。
| 实验配置 | 评测参数 / 验证策略 | 概率覆盖率 / 关键指标 | 接受性能 (\(\tau\) / 接受率) | 草稿阶段耗时 / 说明 |
|---|---|---|---|---|
| 稀疏支持集消融 (L31-8B) | Support Size = 128 | 91.0% | \(\tau = 4.84\) | 13.1 ms (推荐默认值) |
| 稀疏支持集消融 (L31-8B) | Support Size = 256 | 91.8% | \(\tau = 4.85\) | 13.2 ms |
| 稀疏支持集消融 (L31-8B) | Support Size = 512 | 92.2% | \(\tau = 4.78\) | 13.3 ms |
| 稀疏支持集消融 (L31-8B) | Support Size = 1024 | 94.1% | \(\tau = 4.83\) | 13.3 ms |
| 全词表基准 (L31-8B) | Full Dense (全词表) | 100.0% | \(\tau = 4.89\) | 18.8 ms (未优化多项采样) |
| 验证算法对比 (深度1受控) | Top-\(K\) Vanilla | 基线单点质量 | 80.0% 接受率 | 纯确定性点质量截断 |
| 验证算法对比 (深度1受控) | Rheo + RRSw | 引入随机但顺序检验 | 83.5% 接受率 | 混合采样自身带来 +3.5% |
| 验证算法对比 (深度1受控) | Rheo + RheoVerification | OT 视角最优重分配 | 85.4% 接受率 | 较 RRSw 进一步带来 +1.9% |
关键发现¶
- 变阻器参数 \(m\) 的权衡规律：\(m=1\) 在绝大部分模型与任务上均取得最优或接近最优的接受长度。当 \(m=0\) 时，若采样的 token 碰巧命中 top-1 则结构与原树相同，否则会导致不可预期的拓扑漂移；而当 \(m \ge 3\) 时，代理概率随 \(z\) 减小而大幅收缩，随机探针在全局剪枝中几乎必定被裁剪，增益迅速衰减为 0。
- 端到端加速与接受长度的转化比：平均接受长度提升（如 Llama-3.1-8B 从 5.04 提升至 5.25，+4.2%）能够稳定转化为净墙钟加速比增长（2.84× \(\to\) 2.93×，+3.2%）。采样探针与维护残差分布带来的计算开销极低，完全被长尾命中带来的前向跳步收益所覆盖。
- 温度趋零时的退化不对称性：在 \(T \to 0\) 极限下，所有概率采样均收敛至 top-1 确定性选择。\(m=0\) 会展现“采样与结构”的双重退化，完全回归确定性 Top-\(K\)；而 \(m=1\) 即使采样收敛到第二名，其代理概率依然与原始概率存在微小数值偏离，因而在 \(T=0\) 处保留了与基线的微小结构差异。
亮点与洞察¶
- 拓扑排序与分布验证的双重身份解耦：论文最亮眼的 insight 在于指出动态树无损性破损的本质是“token 身份决定了其在剪枝中的存活概率”，通过人工赋予固定代理概率，巧妙绕开了剪枝引起的条件后验分布畸变。
- 等价类分析框架建立无损理论基石：针对动态树中拓扑与 token 取值强耦合构成的复杂随机过程，作者基于等价类划分与数学归纳法证明了任意全局剪枝预算 \(R\) 下的无损性，为后续分析各类随机动态树结构提供了通用的理论模板。
- 实用主义与理论优雅的结合：并没有停留在纯数学证明，而是配合 top-128 稀疏截断与 OT 验证策略，使原本复杂的多项式采样在 GPU 上几乎无感知运行，确保理论提升切实转化为工程加速。
局限与展望¶
- 单步探针数量受限：当前实现中每个父节点仅注入 1 个随机探针，在分支因子较大或多峰分布明显的生成场景下，单一探针可能仍无法充分覆盖高方差尾部。
- 对极低温度场景的额外收益收敛：随着温度降至 0.3 以下，生成分布本身已高度尖锐，随机采样带来的探索增益逐渐趋近于 0，此时维持随机探针的拓扑微扰对吞吐的增益较为有限。
- 未来方向：探索多探针动态自适应分配机制（根据分布熵动态决定 \(m\) 与随机探针个数），以及将双重身份解耦思路拓展到多模态大模型与基于强化学习的投机搜索树中。
相关工作与启发¶
- vs EAGLE-2 / EAGLE-3: EAGLE 系列开创了基于 expand-then-rerank 的动态树推测解码，但在随机解码（\(T>0\)）下直接将草稿概率当做 1.0 处理，退化为确定性 one-hot 剪枝；RheoSampling 继承其高效动态树架构，通过代理概率首次将随机采样无损引入动态树。
- vs Sequoia / SpecTr: Sequoia 与 SpecTr 探索了静态候选树下的最优传输验证算法，但依赖于固定的硬件拓扑；RheoSampling 解决了动态拓扑自适应生成与 OT 验证之间的相容性难题。
- vs RRSw (Recursive Rejection Sampling without replacement): RRSw 采用串行拒绝顺序，导致后置的随机 token 受到前缀拒绝权重的压缩；RheoVerification 优先验证随机探针，最大化保留长尾重叠概率。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次提出解耦树构建代理概率与验证采样概率，化解了动态树随机解码的 one-hot 困境与无损性冲突。
- 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 3 种主流大模型、6 项典型下游基准、详尽的变阻参数消融与等价类理论证明。
- 写作质量: ⭐⭐⭐⭐⭐ 结构清晰，对数学困境的剖析入木三分，算法与理论推导严丝合缝。
- 价值: ⭐⭐⭐⭐⭐ 为大模型推测解码在随机采样场景下的落地扫清了关键理论与工程障碍。