XPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding XPress, a lightweight causal refiner for diffusion drafters in speculative decoding, raises acceptance length by ~30% on average (up to +56%) and decoding throughput by ~1.3× on average (up to 1.7×) compared to the dFlash diffusion drafter on Qwen3-8B across seven math, code, and chat benchmarks. The method restores missing causal dependencies in block-diffusion drafters through parallel refinement, addressing the issue of independently sampled marginals that cause early rejection. XPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding August 2, 2026 12 minutes Case study. A real GSM8K prompt decoded three ways under the same timing setup: autoregressive, the dFlash drafter alone, and XPress ours . Each pane advances by the tokens it accepts per target-verification step, so XPress finishes first. Once the dFlash drafter finishes, the autoregressive pane is fast-forwarded »» so you are not left watching it crawl. TL;DR. 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 ~30% on average up to +56% and its decoding throughput by ~1.3× on average up to 1.7× compared to the dFlash diffusion drafter. 1. Diffusion drafters and the problem with parallel prediction Speculative decoding SD Leviathan et al. 2023 https://arxiv.org/abs/2211.17192 accelerates autoregressive generation by using a lightweight draft model to propose future tokens, which the larger target model verifies in one parallel forward pass. A single multi-token verification pass costs about the same as a standard single-token target-model decoding step in low concurrency scenarios, so every draft token that matches target model outputs is another token produced at no additional target cost. The speedup achieved by SD is mainly governed by two factors: the acceptance length τ the number of drafted tokens the target accepts per verification step and the cost of the drafting process itself, T draft. A larger acceptance length amortizes each target verification pass over more generated tokens, while a cheaper drafter reduces the overhead paid to produce them. This forms a natural trade-off: we want to maximize the expected acceptance length of the drafter, while avoiding a proportional increase to drafting overhead. Autoregressive AR generation has long been the default approach to drafting models, and the EAGLE series Li et al. 2024a https://arxiv.org/abs/2401.15077 2024b https://arxiv.org/abs/2406.16858 2025 https://arxiv.org/abs/2503.01840 is one of the most representative AR-drafting SD methods. EAGLE’s drafter is remarkably lightweight, as small as a single layer, yet it yields high-quality drafts. However, because drafting is autoregressive, generating n draft tokens involves n sequential forward passes of the draft model. As n becomes larger, the drafting overhead becomes increasingly pronounced, but acceptance length, which requires an unbroken chain of accepted verifications, does not. dFlash Chen et al. 2026 https://arxiv.org/abs/2602.06036 resolves this via a block-diffusion model that proposes an entire block of draft tokens in a single forward pass. By turning n serial steps into a single parallel one, dFlash enables longer drafts and downstream speedups for the target model at near-constant overhead. But the parallelism brought by the diffusion drafter carries an inherent limitation in accuracy. Unlike an AR drafter, where each position is conditioned on the preceding tokens, positions in a diffusion drafter are decoded from their marginal distributions. The token k is drawn without seeing what the token k −1 turned out to be, so the block is a set of individually plausible tokens with no guarantee of causality. This can violate natural linguistic dependencies across positions, even when every token is locally high-probability. For example, a drafter predicting each position independently can put a plural verb after a singular subject and produce “she are”, where each word is fine on its own but the verb contradicts the subject. At verification, these locally reasonable but jointly incoherent samples are rejected early by the left-to-right target model, limiting the achievable acceptance length. One line of prior work tries to address this limitation by constructing a draft token tree. Tree-based drafting proposes a tree of candidate continuations and verifies the whole tree in a single target pass, so that the longest accepted path through the tree can be selected. Recent works, like PRESTO Wang et al. 2026 https://arxiv.org/abs/2607.22634 and DDTree Ringel & Romano 2026 https://arxiv.org/abs/2604.12989 , have demonstrated that tree drafting can effectively enhance the achievable acceptance length of diffusion-drafter-based SD methods. Nevertheless, tree drafting has real limitations. Fundamentally, it hedges around the non-causality of the diffusion drafter by targeting recall rather than accuracy. On top of that, the required sparse, irregular tree attention is expensive and complex to serve. Moreover, because a candidate tree typically spans tens to hundreds of tokens, its gains fade quickly at large batch sizes, where the target pass is already compute-bound and those extra tokens are no longer free. The more direct fix is to see whether we can cheaply and effectively restore the lost causality for the diffusion drafter, which is the core research problem motivating XPress. 2. Causal refiner design in XPress Rather than redesigning the diffusion drafter from scratch to restore causality, we instead ask if a small correction to existing outputs will suffice. What makes this plausible is a property already established for diffusion drafters: the correct token is frequently among the drafter’s top- k candidates at each position Wang et al. 2026 https://arxiv.org/abs/2607.22634 . The drafter’s block-level marginals already narrow each position to a small candidate set, it just lacks the causal information needed to identify the right token within that set. The problem of correction therefore reduces to picking the right token from a narrow preexisting set, which should be feasible for a lightweight causal refiner. We thus formulate the correction process as a causal refiner with four properties: Lightweight. The refiner has little room for new parameters or architectural complexity, as the drafter’s single parallel forward pass is already highly streamlined and performant, and this existing capability should be preserved. Causal. Within that small resource budget, the refiner should inject real causal information, conditioning each token on its discretely sampled predecessors, rather than merely smoothing the drafter’s marginals locally. Drafter-grounded. The refiner should make good use of what the diffusion drafter already computes, like its hidden states, which typically carry rich information about the block compared to the pure token id S. L. Wang et al. 2026 https://arxiv.org/abs/2605.09969 . Low overhead. Refinement cost at inference should stay a small fraction of the total drafting time, adding as little latency as possible. In particular, we must avoid reintroducing a fully serial, left-to-right pass. We propose XPress , a lightweight causal refiner instantiating all four properties above. It is lightweight , adding only ~80M parameters 161 MB in bf16 on top of the diffusion drafter; it injects causal information, conditioning each token on its predecessors rather than smoothing marginals locally; it is drafter-grounded , reusing the diffusion drafter’s hidden states rather than a bare token id; and it keeps overhead low , resolving the block in a few parallel iterations rather than a serial left-to-right pass. Figure 1 a shows the full pipeline of XPress: the diffusion drafter proposes the initial block in one pass, the target LM head reads out the base logits s k, and the refiner adds a learned correction on top. Consistent with the aforementioned top- k observation, the refiner does not score the vocabulary from scratch. It adds a small per-position logit bias δ k to the drafter’s own logits, which re-ranks the handful of candidates the drafter already favours. The proposed causal refiner is shown in Figure 1 b . Let V denote the vocabulary size, H the drafter’s hidden width, B the block length, and r =256 the low-rank dimension. All learned linear projections are named W with a descriptive subscript. The position- k logit bias is built in five steps: a k= W in h k W h‖ g W g‖ W e x k−1 ii mix: c k= a k+ Σ j≤k L k,j⊙ a j iii MLP: z k= c k+ MLP c k iv readout: δ k= z k W r∈ ℝ V v correct: ℓ k= s k+ δ k The three inputs to the refiner are the drafter’s per-position hidden state h k, the previous-token id x k−1, and a block-global summary g = 1/ B Σj h j, obtained by mean-pooling the drafter’s hidden states over the block. On the output side, s k is the drafter’s own base logit vector for position k , δ k is the learned correction, and ℓ k is the corrected logit the block is re-decoded from. Among the learned maps, W e ∈ ℝV×r is the token embedding, W h, W g ∈ ℝH×r the down-projections for h k and g , W in ∈ ℝ3r×r the input fusion projection, and W r ∈ ℝr×V the readout LM head. L is a per-channel lower-triangular mixer, so L k,j is nonzero only for j ≤ k . Each design choice earns back one of the four properties: Lightweight. Everything except the two vocabulary matrices W e, W rlives in r -space. W eand W rare the embedding and prediction head required by any logit-bias model. On top of them the refiner adds only small r -space components, so it stays a correction rather than overwriting the base drafter outputs. Drafter-grounded. Feeding the per-position hidden h kand the block-global summary g into the correction gives it strictly more signal than a bare token id S. L. Wang et al. 2026 https://arxiv.org/abs/2605.09969 , and these are obtained for free from the drafter. Causal. Positions exchange information in r -space, via lower-triangular mixing so that position k sees its whole prefix j ≤ k . Because the mix step ii happens after the prior sampled tokens are incorporated in the fuse step i , this enables real causal conditioning. The mixer is also lighter than a conventional attention layer, since it is a single fixed triangular combination rather than a computed attention score, yet still expressive, encompassing the full set of learnable conv1d patterns. Low overhead. Mixing in r -space is lightweight, and a per-position r →2 r → r residual MLP adds cheap nonlinear expressiveness. Steps i and ii stay linear, so the fuse-and-mix path folds into a single matrix at inference. For a Qwen3-8B target and a dFlash drafter, XPress adds 80.5M parameters, of which 96% are the two vocabulary maps that any logit-bias head needs. The causal refiner itself, meaning the hidden inputs, the mixer, and the MLP, is only 2.8M 3.4% , and none of it scales with the 152k-token vocabulary. Relative to the 1.05B-parameter dFlash drafter it attaches to 5 transformer layers plus the input projection; the drafter reuses the target’s embedding and LM head , the whole refiner is a 7.7% add-on, and the causal core alone is 0.26% . | Module | Shape | Params | |---|---|---| We · embed | V × r | 38.9M | Wr · readout | r × V | 38.9M | Wh + Wg | 2·H × r | 2.10M | Win | 3r × r | 0.20M | L · causal mixer | r × B × B | 0.07M | | MLP SwiGLU | r → 2r → r | 0.39M | Total | 80.5M | 3. Parallel refinement via Jacobi decoding The causal mixer restricts the visibility of every position to its own prefix j ≤ k , so the refiner naturally supports autoregressive generation. However, this left-to-right generation process pays an additional B −1-step loop cost over a block of B draft tokens. Instead of finalizing one position before moving to the next, XPress updates all positions at once and repeats this a few times, correcting any prior mistakes, an approach known as Jacobi decoding Song et al. 2021 https://arxiv.org/abs/2002.03629 . The block converges in far fewer than B −1 iterations, which keeps the correction a small fraction of the draft step. Specifically, XPress seeds all positions from the diffusion drafter’s one-shot predictions, then updates them jointly: Seed in parallel. Take the latest set of discrete token predictions. Refine everything at once. One forward of the refiner re-corrects all positions, accounting for any changes in prior positions that surfaced in the last step. Draw new tokens and repeat. Tokens converge to the ground truth AR output from left to right, in at most B −1 steps. Formally, sequential greedy decoding computes y k = argmaxv p k v | y