DBLAST: Dependent Block Drafting for Stochastic Speculative Decoding

Technical Summary: DBlast – Dependent Block Drafting for Stochastic Speculative Decoding

Problem Statement

Speculative decoding accelerates Large Language Model (LLM) inference by using a lightweight drafter to propose multiple future tokens, which a target model then verifies in parallel. While recent block-based drafters (e.g., DFlash) can predict several positions in a single pass, they typically rely on independent block sampling. This approach assumes that tokens within a draft block are conditionally independent given the prefix.

This assumption holds well for greedy or low-entropy decoding but becomes brittle in stochastic (non-greedy) decoding regimes, which are essential for open-ended tasks like dialogue, creative writing, and reinforcement learning exploration. In stochastic settings, the target distribution is deliberately diverse, and multiple continuations are plausible. The paper identifies a fundamental mismatch: while the drafter samples tokens independently to match marginal target distributions, the target model verifies them sequentially, conditioning each acceptance on previously accepted tokens. Consequently, a block with accurate per-position marginals may fail to form a coherent conditional trajectory, leading to a significant degradation in accepted draft length as the entropy of the target sampling distribution increases.

Methodology

The authors propose DBlast (Dependent Block Drafting), a framework that injects dependencies into block diffusion drafters without sacrificing the efficiency of parallel prediction. The method consists of two core components:

1. Dependent Block Sampling via Latent Mixture:
   Instead of factorizing the proposal distribution across positions ($q(x_{1:b}|y) = \prod q_i(x_i|y)$), DBlast introduces a categorical latent variable $z$ (a "branch" or "mode") to induce correlations within the block. The joint distribution is modeled as a low-rank canonical polyadic (CP) mixture:
   $$q_{dep}(x_{1:b} | y) = \sum_{z=1}^{K} q(z | y) \prod_{i=1}^{b} q_i(x_i | y, z)$$

   • Architecture: The method modifies a DFlash-style drafter by adding a lightweight "category hidden expander" that maps the transformer hidden states to $K$ category-specific offsets, and a "category-prior head" to predict $q(z|y)$.
   • Inference: During inference, a latent category is sampled (optionally with a temperature $Z_T$), and the corresponding branch is decoded greedily in a single parallel pass. This preserves the one-pass efficiency of block drafting while allowing the model to represent coherent, alternative block-level continuations.

2. Acceptance-Oriented Training Objective:
   Standard training minimizes the Negative Log-Likelihood (NLL) of target blocks, which optimizes the likelihood of the proposal but ignores the sequential nature of speculative verification. DBlast introduces a surrogate loss based on Expected Accepted Length (AL).

   • The objective is derived from the expectation of the accepted prefix length, decomposed into prefix-acceptance events.
   • To handle the instability of importance weighting when the drafter and target distributions have low overlap (early in training), the authors propose a threshold-truncated log-domain surrogate. This loss sums the log-ratio of probabilities and the log of the conditional accepted length over retained prefixes where the importance ratio exceeds a threshold $\tau$.
   • This approach directly targets the verifier's sequential acceptance behavior, complementing the dependent modeling when block likelihood alone is insufficient.

Key Contributions

1. Diagnosis of the Mismatch: The authors empirically demonstrate that independent block sampling fails increasingly as target decoding becomes less deterministic. They show that accepted length degrades significantly in high-entropy regimes because independent proposals lack the coherence required for sequential verification.
2. Dependent Block Sampling: They introduce the first application of CP-style latent mixture proposals to subword block-diffusion drafters for non-greedy speculative sampling. This method induces correlations among block positions while maintaining parallel token prediction.
3. Acceptance-Oriented Loss: They propose a novel training surrogate motivated by expected accepted length. This objective, combined with threshold truncation, provides a stable training signal that aligns the drafter with the sequential verification process.

Experimental Results

Experiments were conducted using Qwen3-4B and Qwen3-8B target models across four benchmarks: GSM8K (math), MT-Bench (chat), HumanEval (coding), and a creative-writing benchmark. The evaluation covered three target-sampling regimes ranging from high determinism (Temperature 0.7) to low determinism (Temperature 1.5).

• Performance Gains: DBlast consistently outperformed independent block sampling (DFlash) and baselines trained with NLL.
  • In the high-entropy setting (Target Temp 1.5), DBlast achieved a macro-average gain of 12.1% in accepted length for Qwen3-8B compared to the independent baseline.
  • The gains were most pronounced in low-determinism regions, confirming that dependency modeling is most valuable when multiple continuations are plausible.
• Ablation Studies:
  • Latent Categories ($K$): Increasing $K$ improved performance, with diminishing returns observed after $K=4$.
  • Training Objective: The acceptance-oriented (AL) loss consistently outperformed NLL training, even for independent ($K=1$) drafters, though the combination of dependency and AL yielded the best results.
  • Inference Temperature: Calibrating the category sampling temperature ($Z_T$) allowed the drafter to better match the stochasticity of the target distribution, further improving accepted length.

Significance and Claims

The paper claims that DBlast addresses a critical gap in speculative decoding: the inability of current block drafters to handle stochastic verification efficiently. By demonstrating that within-block dependency modeling and acceptance-aligned training are complementary, the authors show that it is possible to achieve high acceptance rates in diverse, non-greedy decoding scenarios without sacrificing the parallel efficiency that makes speculative decoding attractive.

The authors modestly note that while their surrogate loss is not a formal unbiased estimator or lower bound, it serves as a practical and effective training signal. They also acknowledge that the relationship between the soft latent-mixture used during training and the greedy-branch decoding used at inference remains an open area for formal characterization, though their empirical results demonstrate effective transfer. The work suggests that future speculative decoding systems for open-ended generation must move beyond independent token prediction to capture the structural coherence of future sequences.