xPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding

2026-08-03

AI Summary

Existing diffusion-based draft generators, such as dFlash, generate entire token blocks in a single forward pass through parallel sampling; however, their independence assumption neglects causal dependencies among tokens, resulting in low sequence joint probability and frequent early rejection, which limits accepted length. This work proposes xPress, the first diffusion draft generator to incorporate a lightweight parallel causal refinement mechanism that simultaneously preserves generation parallelism and restores inter-token causal dependencies in one step, effectively approximating the true joint distribution. Experiments on Qwen3-8B demonstrate that xPress improves average accepted length by 30% (up to +56%) across seven benchmarks spanning mathematical reasoning, code generation, and dialogue, while achieving an average end-to-end decoding throughput speedup of 1.3× (up to 1.7×).

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.

Problem

• speculative decoding
• diffusion drafters
• causal dependencies
• token acceptance
• parallel refinement

Innovation

• speculative decoding
• diffusion drafter
• causal refinement
• parallel refinement
• token acceptance