D3Tree:学习剪枝块扩散草稿树以加速投机解码验证
D3Tree: Learning to Prune Block-Diffusion Draft Trees for Efficient Verification
投机解码在高并发下常被草稿树拖慢,这篇论文给出一个只学 2H+2 个参数的小头来判断哪些节点值得验证,64 路并发还能比自回归快 1.3-1.8 倍。做推理服务、部署块扩散模型的人应该看看,72 个负载设置里 70 个拿了最快。
论文提出 D3Tree,一种针对块扩散草稿器的草稿树剪枝方法。块扩散草稿器一次前向传播生成整块 token,由位置边缘概率构建的草稿树能提高投机解码的期望接受长度,但每个树节点都占用一次目标模型验证的查询行,在并发服务下低产出分支会与其他请求争夺验证容量。
现有树方法按草稿概率在固定或按成本选定的单请求预算内选择验证节点,总是保留草稿器排序的前缀。D3Tree 则学习哪些节点值得验证:一个仅含 2H+2 参数的树规模头,从草稿状态、候选嵌入和草稿对数概率预测每个候选的条件存活分数,沿父链对分数乘积做阈值化,得到祖先封闭且大小随请求和解码轮次自适应的子树。训练在冻结草稿器上用全候选树的条件标签完成,可选联合模式还能引导草稿器偏向保留的参考路径。
服务时轻负载下验证整棵树,重负载下按预测存活率分配整批节点预算并对齐目标模型的矩阵乘法分块;树的构建、剪枝与调度在 GPU 上运行并与主机重叠。在 6 个视觉语言验证器和 17 个文本与多模态基准上,D3Tree 在 72 个负载设置中的 70 个最快:单请求相对自回归解码提速 3.2-4.0 倍,64 并发请求下仍保持 1.3-1.8 倍,而固定预算树此时已低于自回归速度;在时变到达轨迹下每个到达阶段的平均端到端延迟均最低。
D3Tree: Learning to Prune Block-Diffusion Draft Trees for Efficient Verification
Block-diffusion drafters propose a whole block of tokens in one forward pass, and draft trees built from their position-wise marginals increase the expected accepted length of speculative decoding. Every tree node occupies a query row in target verification, so under concurrent serving low-yield branches compete with other requests for verification capacity. Existing tree methods select verified nodes by draft probability under a fixed or cost-selected per-request budget, and thus always retain a prefix of the drafter's ranking. We present D3Tree, which learns which nodes of a block-diffusion draft tree are worth verifying. A tree-size head with 2H+2 parameters predicts each candidate's conditional survival from the draft state, the candidate embedding, and its draft log-probability. Thresholding the product of these scores along parent chains yields an ancestor-closed subtree whose size adapts to each request and decoding round. The head is trained on a frozen drafter with conditional labels on the full candidate tree; an optional joint mode also steers the drafter toward the retained reference path. At serving time, D3Tree verifies whole trees under light load and, under heavy load, allocates one batch-wide node budget by predicted survival, sized to the target's matrix-multiplication tiles. Retained trees share one ancestor-masked target pass, while construction, pruning, and scheduling run on the GPU and overlap with the host. Across six vision-language verifiers and 17 text and multimodal benchmarks, D3Tree is fastest in 70 of 72 load settings, with 3.2-4.0x speedup over autoregressive decoding for one request and 1.3-1.8x at 64 concurrent requests, where fixed-budget trees fall below autoregressive speed. Under a time-varying arrival trace, it has the lowest mean end-to-end latency in every arrival phase.
This report is categorized under AITOP's "Papers" desk. Within our cross-regional analytical framework, it directly maps to AI developer toolchains, coding agents, and engineering workflows. Key entities and technical topics under continuous tracking include D3Tree, speculative decoding, 扩散模型, 推理加速.
Source attribution: curated from arXiv cs.AI (@Hong Nie, Zhangchi Gao, Jun Yu, Bo Wang, Jianqing Sun, Shengping Liu).