用动态树结构扩散采样,让大模型生成更快更准。
Bastion: Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting

- 基于动态树结构的扩散采样,自适应生成候选路径。
- 在多种GPU上实现最高6.61倍加速,优于当前最优方法39%。
- 无需训练调参,适配不同硬件与模型部署场景。
块扩散草案生成器近年来成为一种强大的推测解码替代方案,可在单次并行步骤中预测多个未来标记分布。然而,由于这些并行预测来自位置边际分布而非完整条件序列,仅选择单一贪婪路径往往无法捕捉目标模型偏好的生成轨迹。为此,我们提出BASTION——一种预算感知的推测解码框架,采用树状扩散草案生成机制。不同于依赖静态树结构的现有方法,BASTION通过平衡草案质量与硬件约束,动态构建查询相关的树结构。该框架集成三个协同组件:(1) 接受率代理,通过路径置信度估计预期接受长度;(2) 在线延迟估算器,校准硬件感知的性能上限模型;(3) 自适应最佳优先扩展策略,持续扩展树结构直至边际收益不再覆盖验证成本。BASTION无需训练,保持目标模型分布不变,且无需针对每种配置进行调优。在多种基准和GPU架构上,其相对于标准自回归解码最高实现6.61倍加速,优于当前最先进的块扩散基线39%。
原文摘要 · Abstract (English)
Block-diffusion drafters have recently emerged as a powerful alternative for speculative decoding by predicting multiple future-token distributions in a single parallel step. However, since these parallel predictions are sampled from position-wise marginals rather than fully conditioned sequences, committing to a single greedy path often fails to capture the target model's preferred trajectory. To address this, we propose BASTION, a budget-aware speculative decoding framework with tree-based diffusion drafting. Unlike existing methods that rely on static tree topologies, BASTION dynamically constructs query-dependent trees by balancing draft quality against hardware constraints. Our framework integrates three synergistic components: (1) an acceptance surrogate that estimates expected accepted length via path confidence, (2) an online latency estimator that calibrates a hardware-aware roofline model, and (3) an adaptive best-first expansion that grows the tree until marginal gains no longer justify incremental verification costs. BASTION is training-free, preserves the target model's distribution, and requires no per-setting tuning. Across diverse benchmarks and GPU architectures, BASTION achieves up to a 6.61x speedup over standard autoregressive decoding, outperforming state-of-the-art block-diffusion baselines by 39%.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。