Technical Summary: xPress: Parallel Refinement for Diffusion Drafters in Speculative Decoding

1. Problem Statement

Speculative decoding (SD) accelerates autoregressive generation by using a lightweight drafter to propose tokens for verification by a larger target model. While autoregressive (AR) drafters (e.g., EAGLE) are effective, they suffer from sequential overhead when generating long blocks. Block-diffusion drafters (e.g., dFlash) address this by generating an entire block of tokens in a single parallel forward pass, drastically reducing drafting latency.

However, diffusion drafters possess a fundamental accuracy limitation: they predict tokens based on marginal distributions rather than joint distributions. In a single-pass diffusion process, token $k$ is sampled without conditioning on the realized values of tokens $1$ through $k-1$. This lack of intra-block causality leads to sequences where individual tokens are locally probable but jointly incoherent (e.g., subject-verb disagreement). During verification, the target model rejects these incoherent sequences early, severely limiting the acceptance length ($\tau$) and negating the potential speedup of parallel drafting.

Existing solutions, such as tree-based drafting (e.g., PRESTO, DDTree), attempt to mitigate this by expanding the candidate space into a tree. However, these methods are computationally expensive, complex to serve, and often hedge around the problem by targeting recall rather than fixing the underlying causal incoherence.

2. Methodology: xPress

The authors propose xPress, a lightweight causal refiner designed to restore missing causal dependencies in diffusion drafters without reverting to a slow, token-by-token autoregressive loop.

2.1 Core Concept

xPress operates as a "causal refiner" that takes the diffusion drafter's initial block proposal and reconciles it through parallel refinement. Instead of re-generating tokens sequentially, xPress injects causal information into the diffusion drafter's logits to correct the marginal distributions.

2.2 Architecture

The refiner is a lightweight module that adds a learned logit bias ($\delta_k$) to the drafter's base logits ($s_k$) at each position $k$. The architecture consists of four stages:

1. Fuse: Projects three inputs into a low-rank space ($r=256$):
   • The drafter's per-position hidden state ($h_k$).
   • A block-global summary ($g$) obtained via mean-pooling.
   • The embedding of the previous token guess ($x_{k-1}$).
2. Mix: Applies a strictly lower-triangular linear mixer ($L$) across the block positions. This allows position $k$ to aggregate information from all predecessors ($j \le k$) in a single operation, enforcing causal structure without sequential computation.
3. Trans: A residual MLP in the low-rank space.
4. Read: Projects the result back to the vocabulary size to produce the bias $\delta_k$.

The final refined logit is $\ell_k = s_k + \delta_k$.

2.3 Parallel Refinement via Jacobi Decoding

A straightforward causal refiner would require a sequential left-to-right pass, reintroducing the latency xPress aims to avoid. To solve this, xPress employs Jacobi decoding:

• Initialization: The block is seeded with the drafter's initial greedy predictions.
• Iteration: All positions are updated simultaneously in parallel using the current state of the block as the prefix for the next iteration.
• Convergence: The process repeats for $K$ iterations. Theoretically, for a block of size $B$, the sequential fixed point is reached in at most $B$ iterations. Empirically, the authors find that $K \approx 6$ iterations are sufficient to match the acceptance length of a full sequential decode, as the drafter's seed is already a high-quality guess.

2.4 Training Strategy

xPress is co-trained with the diffusion drafter using a composite loss function:

1. Cross-Entropy (Teacher Forcing): Maximizes the probability of the ground-truth token given the ground-truth prefix.
2. Total Variation (TV) Loss: Minimizes the distance between the refined distribution and the target model's distribution, directly optimizing for acceptance rate.
3. Consistency Loss: An auxiliary loss applied to a second forward pass where inputs are the drafter's own predictions (self-conditioned). This aligns the training distribution with the inference-time Jacobi process, preventing the refiner from overfitting to gold prefixes.
4. Drafter Anchor: A loss term applied to the base drafter's output to ensure the underlying diffusion model remains high-quality.

3. Key Contributions

• Causal Dependency Injection: xPress injects causal information into diffusion drafters using a minimal mechanism (low-rank projection and a triangular mixer) that operates on the drafter's hidden states, avoiding the need for large parameter increases.
• Parallel Refinement: By utilizing Jacobi decoding, xPress resolves causal dependencies in parallel. It achieves the effect of $B$ dependent autoregressive steps using only $K \ll B$ parallel iterations, preserving the speed benefits of parallel drafting.
• Consistent End-to-End Gains: The method significantly improves both acceptance length and throughput without modifying the verification process or requiring a larger drafter.

4. Experimental Results

Experiments were conducted on Qwen3-8B using a block-16 dFlash drafter across seven benchmarks (Math: GSM8K, MATH-500, AIME25; Code: HumanEval, MBPP, LiveCodeBench; Chat: MT-Bench).

• Acceptance Length: xPress increases the average acceptance length by ~30% compared to the original dFlash diffusion drafter.
  • Under greedy decoding ($T=0$), gains range from 18% to 56%.
  • Under sampling ($T=1$), gains range from 19% to 58%.
• Throughput Speedup:
  • Greedy ($T=0$): xPress achieves an average 1.3× end-to-end speedup over the dFlash drafter, with peaks up to 1.7× (e.g., on GSM8K).
  • Sampling ($T=1$): The speedup improves to an average of 1.46×, reaching up to 1.6×.
• Comparison with Sequential Refiners: When compared to the Markov head (a sequential causal refiner from DSpark), xPress matches or exceeds its acceptance length while running significantly faster.
  • At $K=4$ Jacobi iterations, xPress matches the Markov head's accuracy but runs 1.6× faster (379 µs vs. 601 µs per block).
  • The Markov head requires a fixed 15-step serial decode, whereas xPress converges in fewer parallel steps.
• Scalability: In large-batch serving scenarios (vLLM, batch sizes 16 and 32), xPress maintains its advantage. While absolute speedups over the AR baseline compress as batches grow (due to compute-bound verification), xPress consistently outperforms the plain dFlash drafter by 1.2–1.4× in throughput, proving the refiner's benefits survive under load.

5. Significance and Claims

The paper claims that xPress addresses the "inherent mismatch" between parallel diffusion drafting and sequential verification by repairing the marginal distribution itself rather than expanding the search space.

• Efficiency: xPress demonstrates that causal information can be injected at near-zero cost relative to the drafting overhead. The refinement cost is a small fraction of the total drafting time, and the parallel nature of the refinement avoids the latency penalties of sequential correction.
• Orthogonality: The method is orthogonal to other improvements; it can be combined with larger drafters or tree-based verification methods, as it improves the "backbone" logits that these methods rely on.
• Generalizability: The authors suggest the philosophy of reconciling mutually dependent predictions via a few parallel iterations may be applicable beyond speculative decoding, wherever models emit interdependent outputs simultaneously.

The paper concludes that xPress successfully restores the causality lost in single-pass diffusion drafting, enabling longer acceptance chains and higher throughput without sacrificing the parallel efficiency that makes diffusion drafters attractive.