[控场AI]
· 6 分钟阅读· 3,059 字

TreeSpark:自适应草稿树让推测解码提速8-14%

TreeSpark:自适应草稿树让推测解码提速8-14%

TreeSpark用路径存活率统一驱动草稿树的构建、停止与负载调度,解码提速8-14%且保持无损。

TreeSpark是针对大模型推测解码中草稿树构建方式的优化工作,专为半自回归草稿器设计。现有方法存在两个核心缺陷:用边缘概率排序节点导致上下文信息丢失,以及树尺寸固定无法适应动态负载。TreeSpark的核心思路是从草稿器已有的Markov头中以极低代价读取父节点条件分布,将其校准为边接受估计,再以"路径存活率"这一统一信号驱动最优优先扩展、逐轮动态停止和负载自适应服务策略三项决策。通过无放回采样兄弟节点与递归残差匹配,TreeSpark在任意温度下保持解码无损。实验表明,相比调优的链式方案每轮多接受15-25%草稿token,单请求解码提速8-14%,且在高负载时能平滑退化为链式结构,具备工程实用性。

推测解码的下一站:更聪明的草稿树

大模型推理的成本瓶颈之一在于逐词生成(autoregressive decoding)的串行本质。推测解码(Speculative Decoding)提供了一条加速路径:用一个廉价的"草稿模型"(drafter)先提出若干候选 token,再让目标模型一次性并行验证。这种"先猜后验"的机制在保持输出无损的前提下,显著减少了目标模型的前向计算次数。

近期出现的块状草稿器(block drafter)把草稿生成的成本压得几乎为零——一次骨干网络前向就能吐出一整块草稿 token。而草稿树(Draft Tree)则试图更进一步:在一次目标模型前向中同时验证多条候选延续路径,从而榨取更多加速空间。

arXiv 上的新论文 TreeSpark 正是针对这一方向的优化工作,作者提出了一套"经过校准、负载自适应"的草稿树构建方法,专为半自回归(semi-autoregressive)草稿器设计。

TreeSpark 论文

推测解码的无损性保证来自一个关键的数学性质:若草稿 token 被目标模型以特定概率接受或拒绝,最终输出序列的分布与目标模型单独自回归生成完全等价。这一机制由 Leviathan 等人(2023)正式确立。块状草稿器(block drafter)是推测解码的一种变体,它不再逐 token 自回归生成草稿,而是用一次前向传播并行吐出多个位置的 token,代价是草稿质量略低于自回归草稿器。半自回归(semi-autoregressive)草稿器则介于两者之间:以固定块为单位生成草稿,块内并行、块间串行,在草稿质量与生成速度之间取得折中。草稿树的引入进一步将单条候选链扩展为多条并行候选路径,让目标模型在一次前向中同时验证树上所有节点,理论上可以在不增加目标模型调用次数的前提下大幅提升平均接受长度。

现有草稿树的两个硬伤

论文明确指出了当前草稿树构建方式存在的两个核心问题。

排序信号错位

现有构造方法在对候选节点排序时,依赖的是"每个位置的边缘概率"(per-position marginals),也就是说它忽略了一个候选究竟是从哪个父节点延伸而来的。这在半自回归草稿器上尤其致命:当树变得更宽时,新增的节点大多是被错误排序的"噪声",宽树带来的收益被大量误排节点抵消。

边缘概率(marginal probability)与条件概率的区别在树结构中至关重要。边缘概率描述的是某个位置上某个 token 出现的总体可能性,它把所有可能的前缀路径都平均进去了;而条件概率则明确刻画"在某个具体父节点之后,该 token 出现的可能性"。在链式草稿中两者差异不大,因为路径只有一条。但在树结构中,同一位置的同一 token 可能挂在不同父节点之下,其实际被目标模型接受的概率差异极大。用边缘概率排序相当于忽略了路径上下文,会把一些"孤立看起来高频但在当前路径下几乎不可能被接受"的节点排到树的高优先级位置,浪费了宝贵的验证预算。

固定树尺寸的僵化

另一个问题是树的尺寸被固定死了。但实际情况是,每一轮解码能支撑多少"投机"是动态变化的,不同的服务负载(serving load)下能承受的投机预算也不同。用一棵固定大小的树去应对所有场景,显然无法做到最优。

TreeSpark 的核心思路

TreeSpark 的关键洞察是:把"路径存活率"(path survival)作为统领全局的信号。

它从草稿器已有的 Markov 头(Markov head)中读取一个父节点条件分布(parent-conditioned distribution),这一步几乎不增加额外成本,因为它利用的是草稿器本就存在的结构。随后把这个分布校准成一个边接受估计(edge-acceptance estimate),用于判断每条边被目标模型接受的可能性。

有了这个校准后的存活率估计,TreeSpark 让路径存活率去支配三件事:

  • 最优优先扩展(best-first expansion):优先扩展更可能存活的路径
  • 逐轮停止(per-round stopping):动态决定每一轮何时停止扩展
  • 负载自适应服务策略(load-adaptive serving policy):根据当前负载调整树的规模

保证无损解码

值得关注的是解码无损性的处理。TreeSpark 采用无放回采样兄弟节点(sampling siblings without replacement),并在递归拒绝采样中匹配残差(matching residuals in recursive rejection),从而在任意温度下都保持解码结果与目标模型分布一致。这一点很重要——很多加速方法在高温采样下会改变输出分布,而 TreeSpark 声称能规避这个问题。

无放回采样(sampling without replacement)在树验证中解决的是一个微妙的概率一致性问题。当目标模型拒绝某个 token 时,标准的推测解码会从残差分布(目标分布与草稿分布之差)中重新采样一个替代 token。在树结构中,同一父节点可能有多个兄弟候选节点,若这些兄弟节点是有放回采样的,则不同分支之间会存在重叠,导致整体接受概率的计算出现偏差,最终破坏分布无损性。无放回采样确保兄弟节点彼此互斥,使残差分布的匹配在数学上保持严格正确。这一细节在温度较高时尤为关键,因为高温采样本身就意味着更大的分布方差,任何概率计算的偏差都会被放大,更容易导致输出分布偏离目标模型。

实测收益

根据论文给出的数据,TreeSpark 的表现可以总结为几个层面:

  • 在匹配的固定预算下,自适应树在每一个温度上都优于固定树;
  • 相对于同一草稿器上调优过的"链式"(chain)方案,TreeSpark 每轮多接受 15-25% 的草稿 token;
  • 在单请求的墙钟时间(wall-clock)上,解码速度提升 8-14%;
  • 在负载上升时,它能"优雅地"把树收缩回链式结构,避免高负载下的资源浪费。

最后一点尤其体现了工程实用性——加速方案在真实服务环境中往往面临波动的负载,能够在"追求极致加速"和"退化为保守策略"之间平滑切换,比一个只在理想条件下跑分漂亮的方法更有落地价值。

意义与局限

TreeSpark 的贡献不在于发明全新范式,而在于把推测解码中一个被忽视的信号——父节点条件下的路径存活率——重新拉回到树构建的核心位置,并让它统一驱动扩展、停止和负载调度三个决策。这种"用一个统一信号简化多个启发式"的设计思路,在系统优化中往往比堆砌复杂规则更稳健。

需要留意的是,论文所报告的加速数字(8-14%)是在特定草稿器和对比基线下测得的,实际收益会随模型规模、任务类型和硬件配置而变化。半自回归草稿器本身也还不是主流部署方案,其适用范围有待更广泛的验证。

对于关注 LLM 推理加速的工程师和研究者,TreeSpark 已开源了代码与实验产物,可在 GitHub(PopSoda2002/TreeSpark)上进一步探索其实现细节。

#推测解码#TreeSpark#草稿树#大模型推理加速#半自回归#speculative decoding#LLM推理优化
分享:

相关推荐