Research map / Image generation / Fast diffusion sampling

BASTION: Budget-Aware Speculative Decoding with Tree-structured Block Diffusion Drafting

NeurIPS 2026 · Image generation › Fast diffusion sampling · approach: Masked diffusion language model decoding · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes BASTION, a training-free framework that builds query-dependent prefix trees from block-diffusion drafter logits via best-first expansion and stops growth with a hardware-calibrated latency model, reaching up to 6.61x speedup over autoregressive decoding.

Problem

Single-path greedy drafting from block-diffusion drafters loses target-preferred continuations, and static tree topologies ignore query and hardware-dependent verification costs in speculative decoding.

Approach

Uses the sum of drafter path scores as a surrogate for expected accepted length, builds the tree by best-first expansion (provably optimal for that surrogate), and picks the tree size online by maximizing estimated speedup with a roofline-based, linearly calibrated verification-cost model.

Key results

• On A100, BASTION averages 6.61x speedup on Qwen3-8B at T=0 across eight benchmarks, versus 4.76x for DFlash and 2.45x for EAGLE-3 (per-model averages in the Qwen3-4B row are 6.93x).
• Best-first expansion beats beam search under matched budgets, improving Qwen3-4B and Qwen3-8B speedups by +7.0% and +6.1% on A6000.
• Calibrated roofline latency reduces verification-latency RMSE by 87-92% (e.g., 30.6 to 3.1 ms on Qwen3-8B).

Surprising findings

• The drafter-only path-sum surrogate correlates with target-verified accepted length (Pearson 0.79-0.82), despite using the drafter in place of the target.
• Online EMA-only latency tracking is fragile and selects oversized trees on Llama-3.1-8B, losing 12.7% speedup.

Limitations

• Evaluated only at batch size 1; larger batch regimes change verification cost profile.
• Default offline calibration assumes a stable runtime profile; runtime drift is not handled well.
• Surrogate relies on drafter-target alignment and is a proxy for target-verified acceptance.

Experimental setup

Datasets: GSM8K MATH500 AIME25 HumanEval MBPP LiveCodeBench MT-Bench AlpacaEval

Benchmarks: GSM8K MATH500 AIME25 HumanEval MBPP LCB MT-Bench Alpaca

Metrics: speedup average acceptance length (tau)

Baselines: DFlash EAGLE-3 beam search tree DDTree

Builds on: DFlash EAGLE-3 Sequoia SpecInfer DDTree

Closest papers by meaning

• SpecBlock: Block-Iterative Speculative Decoding with Dynamic Tree Drafting (NeurIPS 2026 · similarity 0.93)
• Block Verification Accelerates Speculative Decoding (ICLR 2025 · similarity 0.92)
• DFlash: Block Diffusion for Flash Speculative Decoding (ICML 2026 · similarity 0.91)
• D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting (NeurIPS 2026 · similarity 0.91)
• TreeGraft: Adaptive Multi-Drafter Grafting for Tree-Based Speculative Decoding (NeurIPS 2026 · similarity 0.91)
• UniVer: A Unified Perspective for Multi-step and Multi-draft Speculative Decoding (NeurIPS 2026 · similarity 0.91)