用目标模型指导的树形草稿生成,让扩散语言模型更快更准。
GRAFT: Adaptive DLM-Based Draft Tree Construction with Target-Distilled Edge Scoring

- 基于目标模型提取父子令牌兼容性,动态选择最优路径。
- 自适应调整每轮树大小,实现2.13到6.36倍端到端加速。
- 适合追求高吞吐、低延迟的生成式推理场景。
基于树的推测解码通过验证多条草稿路径提升了标准推测解码的平均采纳词元数。现有树构建方法通常采用父路径条件扩展,每个子词元均基于父路径生成,这与扩散语言模型(DLM)如DFlash不兼容,后者可在单次前向传播中生成所有未来位置的概率分布。DDTree通过将各未来位置的高概率词元作为候选节点,并在固定节点预算下选择连续位置间的边来填补这一空白。然而,其边选择仅依赖词元概率,未建模父子兼容性,导致目标匹配的词元可能连接错误父节点;且固定预算忽略了最优树大小随解码状态变化的特性。本文提出GRAFT框架,针对基于DLM的推测解码设计草稿树构建方法。GRAFT引入目标蒸馏边评分(TDES),从目标模型轨迹中蒸馏父子偏好以选择目标兼容的边;并提出状态感知预算分配(SABA),根据预期草稿收益与验证成本平衡每轮树预算。在多个模型和任务上,GRAFT相比自回归解码实现2.13×至6.36×的端到端加速,每轮额外开销低于0.5毫秒,约为目标模型验证延迟的1.4%。
原文摘要 · Abstract (English)
Tree-based speculative decoding raises the mean accepted tokens of standard speculative decoding by verifying multiple draft paths, and existing tree builders typically construct these paths through parent-conditioned expansion, where each child token is generated conditioned on its parent path. This construction is incompatible with diffusion language model (DLM) drafters such as DFlash, which produces all future-position distributions in a single forward pass. DDTree bridges this gap by treating high-probability tokens from each future-position distribution as candidate nodes and selecting edges between consecutive positions under a fixed node budget. However, its edge selection relies on token probability alone without modeling parent--child compatibility, so target-compatible tokens can be attached to wrong parents; moreover, its fixed budget ignores that the throughput-optimal tree size varies with the decoding state. We propose GRAFT, a draft-tree construction framework for DLM-based speculative decoding. GRAFT introduces Target-Distilled Edge Scoring (TDES), which distills parent--child preferences from target-model traces to select target-compatible edges, and State-Aware Budget Allocation (SABA), which sets the per-round tree budget by balancing expected draft gain against verification cost. Across multiple models and tasks, GRAFT achieves $2.13\times$--$6.36\times$ end-to-end speedup over autoregressive decoding while adding less than $0.5$\,ms of overhead per round, approximately $1.4\%$ of the target-model verification latency.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。