xPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding

2026-08-03Artificial Intelligence

Artificial Intelligence
AI summary

The authors studied a method called block-diffusion drafting, which produces many draft tokens at once but picks each token independently, losing the natural order of language. This independence makes the drafts less likely to be accepted because tokens don't fit well together. To fix this, they proposed xPress, a lightweight tool that refines the entire draft in parallel to add causal dependencies between tokens without slowing down the process. Their tests showed that xPress improves how long accepted sequences can be and speeds up decoding compared to the original method.

block-diffusion draftingdFlashlogit distributionspeculative decodingconditional independencetoken marginalscausal dependenciesparallel refinementQwen3-8Bdecoding throughput
Authors
Zheng Wang, Davis Wertheimer, Yu Chin Fabian Lim, Mudhakar Srivatsa, Raghu K. Ganti, Minjia Zhang, Naigang Wang
Abstract
Block-diffusion drafters like dFlash generate an entire block of draft tokens in a single forward pass, drastically reducing the overhead of multiple-token drafting in speculative decoding. The crucial final step of the single-pass discrete denoising process involves using the logit distribution at each position to sample conditionally independent tokens. The resulting draft is thus a set of per-position marginals, rather than a joint distribution: no draft token is guaranteed to depend on its predecessors. Such independently sampled marginals tend to produce sequences with tokens that are individually likely, but jointly improbable under the target model's distribution, which verifies each token conditionally. This can cause early rejection and limits acceptance length. To address this, we propose xPress as a means to restore the missing causality in diffusion drafters. xPress is a lightweight causal refiner that reconciles the whole diffusion block at once through parallel refinement, restoring and propagating causal dependencies across the draft without a token-by-token loop. On Qwen3-8B, across seven math, code, and chat benchmarks, xPress raises acceptance length by about 30% on average (up to +56%) and its end-to-end decoding throughput by about 1.3 on average (up to 1.7) compared to the original dFlash diffusion drafter.