PaperlayerView exact revision on arXiv

Source-linked AI summary

BlockPilot: Instance-Adaptive Policy Learning for Diffusion-based Speculative Decoding

Hao Zhang, Yiming Hu, Yong Wang, Mingqiao Mo, Xin Xiao, Xiangxiang Chu

arXiv:2606.31315v1cs.CL

TL;DR

Fixed block sizes can be suboptimal because the best decoding choice varies by input. BlockPilot predicts each sample’s block size from the prefilling state, yielding consistent speedup gains with minimal overhead.

  • Problem

    Diffusion-based speculative decoding typically uses one fixed inference block size, although the optimal size can vary across input samples.

  • Method

    BlockPilot uses a lightweight policy predictor to select a sample-specific block size from the target model’s prefilling representation.

  • Results

    BlockPilot achieves consistent speedup gains with minimal overhead in diffusion-based speculative decoding.

  • Takeaways & Limitations

    Decoding policy is a critical factor in inference efficiency beyond model architecture alone.

  • Takeaways & Limitations

    Training-data construction can be computationally burdensome for very large models because it evaluates multiple candidate block sizes per sample.

Abstract

from arXiv · show

Speculative decoding accelerates inference by using a lightweight draft model to generate candidate tokens in parallel, and are then verified by the target model, enabling lossless acceleration. Recently, diffusion-based speculative decoding further improves parallelism by generating multiple tokens per forward pass via block-level diffusion, achieving state-of-the-art (SOTA) performance. However, existing methods adopt a fixed inference block size and assume a uniform optimal decoding strategy across all inputs. In this paper, we show that this assumption is suboptimal, as the optimal block size varies across samples and plays a critical role in speculative decoding performance. Moreover, these values exhibit a clear local structure, concentrating around the training block size, which reduces the problem to a low-dimensional and structured decision space. Based on these insights, we propose BlockPilot, a sample-adaptive policy that predicts the optimal block size from the prefilling representation. Specifically, we formulate block size selection as a lightweight policy learning problem and propose an instance-adaptive decision mechanism that predicts the optimal block size based on the representation of the prefilling stage. The prediction is performed only once after prefilling, allowing for seamless integration. Extensive experiments demonstrate that our method is plug-and-play, introduces minimal overhead, and consistently improves efficiency, achieving an acceptance length of 5.92 and a 4.20$\times$ speedup on Qwen3-4B under temperature $T=1$.

1 Introduction

BlockPilot reframes diffusion-based speculative decoding as an instance-adaptive policy problem because fixed block sizes are suboptimal across inputs. It exploits the local structure of optimal block sizes to predict a suitable block size from the prefilling state with lightweight overhead.

  • Motivation: Diffusion-based speculative decoding generates multiple tokens in parallel with a draft model and verifies the proposed block using the target model.This preserves speculative decoding’s lossless acceleration while reducing decoding latency through block-wise parallel generation.
  • Problem: Fixed inference block sizes assume one decoding strategy is optimal for every input, but this assumption is suboptimal.The optimal degree of parallelism depends on the input, making block size a critical decoding-policy variable.
  • Key observation: Optimal block sizes vary across samples but usually concentrate in a narrow region around the training configuration.This locality converts sample-adaptive selection into a small, discrete, structured classification problem rather than an expensive online search.
  • Method: BlockPilot predicts the block size from the prefilling state using a lightweight instance-adaptive policy.The method uses the predictive distribution of the final prefilling token as a representation of the current decoding state.
  • Contribution: The proposed framework treats decoding policy as learnable and achieves consistent speedup gains with minimal overhead.Its instance-adaptive prediction is designed for efficient integration into diffusion-based speculative decoding.

2 Methodology

BlockPilot frames diffusion-based speculative decoding as sample-adaptive block-size selection: optimal sizes vary across inputs but remain localized near the training size. It therefore learns a lightweight classifier from prefilling representations, using offline enumeration to supervise efficient candidate selection.

  • Speculative Decoding Framework: Diffusion-based speculative decoding drafts B tokens in parallel with a lightweight diffusion model, then verifies them in parallel using the target autoregressive model.Block-parallel drafting and verification alleviate autoregressive decoding’s sequential bottleneck.
  • Speculative Decoding Framework: Acceptance length τ(B) primarily governs efficiency because draft and verification costs typically grow sublinearly with block size B.The optimal inference block size B∗ therefore maximizes τ(B) and improves end-to-end acceleration.
  • Instance-wise Variability: The optimal block size varies significantly across samples and does not necessarily match the fixed training block size B.This variation reflects differences in context structure and predictability, which affect accepted lengths as block size changes.
  • Local Interval Property: The optimal block size is strongly localized around B: the interval [B − 3, B + 3] covers the optimum for nearly all samples.Outside this narrow range, acceptance length drops sharply and such choices are rarely optimal.
  • Classification Formulation: Block-size selection is formulated as structured classification over the local candidate set {B−k, . . . , B+k}, conditioned on each input sample.This reduces the search space and enables learning-based selection instead of global search.
  • Policy Learning: Supervised labels are constructed by enumerating candidate sizes after extracting the target model’s last-position prefilling distribution p(x), while a lightweight n-layer MLP predicts candidate logits.The compact predictor avoids unnecessary sequence modeling and keeps training and inference computationally lightweight.

3 Experiments

Experiments across four LLMs and Math, Code, and Chat benchmarks show that BlockPilot consistently outperforms autoregressive, EAGLE-3, and fixed-block DFlash baselines with negligible latency. Ablations indicate that moderate predictor capacity, k = 2, and raw prefilling distributions support effective block-size selection.

  • Models and Benchmarks: Experiments cover four LLMs and Math, Code, and Chat benchmarks spanning mathematical reasoning, code generation, software engineering, and conversational generation.Models include Qwen3-4B, Qwen3-8B, Llama-3.1-8B-Instruct, and Qwen3-Coder-30B-A3B.
  • Baselines and Metrics: BlockPilot is compared with standard autoregressive decoding, EAGLE-3, and fixed-block DFlash variants, focusing on speedup ratio and average acceptance length τ.Because speculative decoding preserves the target model’s exact output distribution, the experiments focus on efficiency rather than generation quality.
  • Main Results: 4.20× average speedup is achieved on Qwen3-4B at temperature = 1, while BlockPilot consistently performs best across models, temperatures, and benchmark categories.The method modifies neither draft nor target models and adds only negligible latency.
  • Main Results: 0.18× speedup and 0.28 acceptance-length gains are obtained over DFlash(16) on Qwen3-4B at temperature = 0, improving speedup from 3.99× to 4.17× and τ from 6.31 to 6.59.On Qwen3-8B, speedup rises from 4.42× to 4.66× and τ improves from 6.13 to 6.46.
  • Main Results: Performance remains effective across Math, Code, and Chat benchmarks and under temperature = 1, where stochastic decoding makes draft-token acceptance more difficult.The results suggest that the last-token predictive distribution after prefilling provides a useful adaptive block-size signal.
  • Ablation Studies: Increasing predictor width from D = 1024 to D = 2048 improves speedup and τ, while D = 4096 yields negligible gains; k = 2 balances coverage and prediction difficulty.Normalization and softmax preprocessing reduce speedup and τ relative to using the raw prefilling distribution.

4 Related Work

Prior work has advanced speculative decoding through lightweight draft models, multi-head prediction, and diffusion-based parallel generation. Diffusion-based speculative decoding further explores diffusion drafting, but existing approaches face quality or effectiveness limitations.

  • Speculative Decoding: Speculative decoding uses a lightweight draft model to propose candidate sequences that a larger target model verifies in parallel.Medusa removes the external draft model by adding multiple prediction heads to the base LLM.
  • Diffusion Language Models: Diffusion language models enable nonautoregressive generation through parallel masked-token prediction, while block diffusion combines parallelism with autoregressive structure.LLaDA scales diffusion language models to the billion-parameter level, but fully parallel models face fixed-length generation and inefficient KV-cache usage.
  • Diffusion-based Speculative Decoding: Diffusion-based speculative decoding generates drafts in parallel, but existing methods remain limited by quality or effectiveness constraints.TiDAR does not achieve lossless generation quality, while LoRA-based parallel drafting has limited effectiveness; DiffuSpec and SpecDiff-2 are also described as diffusion-drafting approaches.

5 Conclusion

BlockPilot identifies block size as a key determinant of diffusion-based speculative decoding efficiency, with optimal values varying by sample but remaining local around the training configuration. It uses prefilling-stage predictions to select a local candidate once per sample, integrating with existing frameworks.

  • 5 Conclusion: Block size is a key factor in diffusion-based speculative decoding inference efficiency, and its optimum varies across samples.The optimal values exhibit strong locality around the training configuration.
  • 5 Conclusion: BlockPilot formulates block-size selection as a small structured decision problem using locality around the training configuration.This structure enables selection from a local candidate set.
  • 5 Conclusion: BlockPilot uses the prefilling-stage predictive distribution to select block size from a local candidate set once per sample.The method integrates seamlessly into existing diffusion-based speculative decoding frameworks.
  • 5 Conclusion: The decoding policy, rather than model architecture alone, plays a critical role in inference efficiency.This conclusion follows from the observed sample-wise variation in optimal block size and BlockPilot’s adaptive policy.

A Theoretical Analysis of Sample-Adaptive Block Size Selection · A.1 Acceptance Length as a Prefix-Survival Process

The analysis expresses expected acceptance length as a truncated prefix-survival process, where each position contributes according to the probability that the verified prefix survives to that position. Block size affects both the available acceptance horizon and the conditional survival probabilities, creating competing effects on acceptance length.

  • A.1 Acceptance Length as a Prefix-Survival Process: Expected acceptance length decomposes into a sum of prefix-survival probabilities across positions up to block size b.Because L_b(x) is non-negative and bounded by b, its expectation admits a survival-probability decomposition.
  • A.1 Acceptance Length as a Prefix-Survival Process: Each term represents the probability that the verified prefix survives through at least position i.This gives acceptance length a prefix-level interpretation rather than treating it solely as an empirical statistic.
  • A.1 Acceptance Length as a Prefix-Survival Process: Accepting at least i draft tokens requires all first i drafted tokens to be accepted, defining the prefix-consistency event A_i.The target model accepts only the longest consistent prefix of the drafted block.
  • A.1 Acceptance Length as a Prefix-Survival Process: The probability that the first i drafted tokens survive verification factors through conditional acceptance probabilities q_j(x, b).The factorization follows from defining q_j(x, b) conditioned on survival through the preceding position and applying the chain rule.
  • A.1 Acceptance Length as a Prefix-Survival Process: Acceptance length is the expected stopping time of a truncated prefix-survival process whose horizon is set by b.The contribution of each additional drafted position is governed by its corresponding prefix survival probability.
  • A.1 Acceptance Length as a Prefix-Survival Process: Increasing block size raises the maximum accepted length but can make prefixes harder to verify because block-level diffusion tokens are mutually dependent.The proposal distribution must remain coherent over a longer span as block size increases.
  • A.1 Acceptance Length as a Prefix-Survival Process: A larger b extends the summation horizon while potentially changing the conditional acceptance probabilities within each multiplicative survival term.Because prefix survival is multiplicative, mild degradation in q_j(x, b) can amplify over longer prefixes.
  • A.1 Acceptance Length as a Prefix-Survival Process: Expected acceptance length is governed by the interaction between block length and prefix survival, not by the number of drafted tokens alone.This interaction explains why increasing the block size does not guarantee higher acceptance length.

A.2 Locality Induced by Predictability and Block-Size Retention

The analysis decomposes acceptance into sample-dependent predictability and block-size retention, showing why deviations from the training block size degrade longer verified prefixes. This locality motivates restricting adaptive selection to a candidate interval near the training block size, turning global search into a structured local decision problem.

  • Predictability–retention decomposition: Draft proposals are best calibrated near the training block size B, while substantial inference deviations can reduce proposal quality and acceptance probabilities.The retention model separates sample-dependent predictability from block-size-induced proposal degradation.
  • Predictability–retention decomposition: Deviations from B reduce prefix survival through r(b; B)^i, with stronger effects for longer prefixes because retention is exponentiated by prefix length.The retention function is an analytical explanation of locality, not an additional algorithmic component.
  • Geometric approximation: Sample predictability controls survival: structured inputs sustain longer verified prefixes, whereas uncertain inputs experience faster survival decay.The effective survival factor combines sample-level predictability γx with block-size-induced proposal degradation r(b; B).
  • Local candidate interval: The local interval Bloc = {b ∈B : |b −B| ≤k} captures the optimal block size for most samples and reduces selection from global search to a structured local decision problem.The best block size remains sample-dependent but is expected to concentrate near the diffusion draft model’s training block size.

A.3 Regret of Local Block-Size Prediction

The section analyzes acceptance-length regret caused by predicting a block size different from the optimum. Under a discrete Lipschitz assumption, regret is controlled by prediction distance, motivating local candidate sets and lightweight prediction.

  • Regret definition: Acceptance-length regret compares the predicted block size ˆB(x) with the optimal block size B∗(x) over the candidate set B.The optimal and predicted block sizes are defined for each input x.
  • Lipschitz assumption: The acceptance-length function satisfies a discrete Lipschitz condition, with |τ(b1; x) − τ(b2; x)| ≤ Lτ|b1 − b2| for all b1, b2 ∈ B.The condition assumes a constant Lτ > 0.
  • Regret bound: Acceptance-length loss is controlled by the distance between predicted and optimal block sizes, so near-optimal predictions retain bounded regret.Exact prediction is not strictly necessary when the selected block size is close to the optimum.
  • Design implications: The analysis supports restricting prediction to a local candidate set around B and using a lightweight predictor because neighboring block sizes have similar acceptance behavior.These are presented as two design choices motivated by the regret bound.

B Performance Evaluation on Instruction and Code Models

BlockPilot consistently delivers the strongest performance across Llama-3.1-8B-Instruct and Qwen3-Coder-30B-A3B on Math, Code, and Chat benchmarks. Its adaptive block selection improves speedup across decoding temperatures and generalizes across heterogeneous workloads and model scales.

  • Temperature = 0: 3.25× and 4.12× are the highest average speedups achieved by our method on Llama and Qwen, respectively, under temperature = 0.The method also achieves the best or near-best average acceptance length τ.
  • Temperature = 0: 4.12× improves Qwen speedup over 3.86× for DFlash(16) under temperature = 0, with consistent gains in τ.The passage identifies this as a significant advantage over all DFlash variants.
  • Temperature = 1: 2.40× and 3.95× are the average speedups achieved by our method on Llama and Qwen, respectively, under temperature = 1.The passage states that similar performance trends hold at this temperature.
  • Generalization: Improvements remain consistent across Math, Code, and Chat, confirming adaptive block selection’s effectiveness across heterogeneous workloads, medium-scale models, and large-scale LLMs.The results support the method’s robustness across both model scales.
  • Evaluation setting: Table 6 evaluates speedup ratios and average acceptance length τ for Llama-3.1-8B-Instruct and Qwen3-Coder-30B-A3B across Math, Code, and Chat.DFlash(n) denotes DFlash with block size n.

C Limitations and Future Work

BlockPilot’s data construction identifies suitable block sizes for different training samples, but evaluating candidate sizes introduces computational burden, especially for very large models. For a 32B model, one sample under a single block size takes approximately 5 seconds.

  • Limitations: BlockPilot’s data construction enables fine-grained execution preferences but incurs computational burden, particularly for very large models.The method identifies suitable block sizes for different training samples by evaluating candidate block sizes.
  • Limitations: 5 seconds is the approximate execution time for one 32B-model sample under one block size.Constructing one training sample requires evaluating all candidate block sizes in the specified range.
Loading 2606.31315v1…