On the SPOT: Enhancing Speculative Decoding for Diffusion Language Models via Commit-Order Prediction

Paper

Diffusion language models must decide ‘what’ and ‘where’ at the same time

Diffusion language models generate sentences by repeatedly filling masked tokens. To decode exactly, each model forward must commit only the single most certain masked position as a token, which is slow; committing multiple positions at once to increase speed degrades generation quality. This paper moves speculative decoding, a lossless acceleration technique validated on autoregressive models, to diffusion language models, but finds the solution by predicting not tokens but the order of which position to fill next.

The authors first show one observation experimentally. Looking only at the probability distribution left by the target model in the previous forward, the token that belongs at each position can mostly be predicted in advance (92% for the first block commit and over 90% thereafter), but which position will be filled next is rarely predicted. So they reuse tokens from the old distribution as-is, and a small 0.62M-parameter model (SPOT) predicts only the commit order to build a tree of candidate states. The target verifies that tree with a single batched forward. On SDAR-8B, this reduces target forwards by up to 6.8× and raises throughput up to 3.4× versus standard decoding (§1, §5.2).

Background: decoding that fills a block one slot at a time

The target of this paper is a block-causal diffusion language model. An active block has positions of length $L$ (the paper mainly uses b32, i.e. $L{=}32$), and each position is either a token or a mask. A single forward returns a full-vocabulary probability distribution for every masked position. The target then selects only the position whose top probability is highest and fills it with that position’s argmax token. This is called commit, and the paper counts the number of target forwards as NFE (number of function evaluations) (§3).

The issue is that attention within a block is bidirectional. The moment one position is filled, the representations and probability distributions of all remaining masked positions change. So, unlike an autoregressive model, “the next position is fixed and only its token needs to be predicted” does not hold. Speculative decoding for diffusion language models must determine both which token to use and which position to use (§1, §2).

Illustrative example (not experimental numbers)

Let us reduce the block width to 8. Three tokens, “The”, “cat”, and “sat”, have already been filled, and the remaining five slots are masks. Suppose the highest-probability candidates (stale argmax) for the masked positions returned by one forward are as shown below.

In this forward, the target commits only the single most certain position. Position 7’s “.” has the highest probability at 0.95, so it fills position 7 with “.” and examines the remaining four slots again in the next forward. The state retained here consists of the already committed tokens (“The cat sat …”), and what changes from forward to forward is the single filled position and the entire probability distribution over the remaining masked positions.

The key point is that two things diverge in this example. The token that will ultimately go into each position largely matches the stale argmax in this table (“on”, “the”, and “mat” are mostly right). But “which position to commit first” cannot be known from probability magnitude alone. The moment position 7 is filled, the distributions of the remaining slots all shift, so this table becomes the stale (old) distribution from the next moment onward.

Observation: the tokens are already known, but the order is not

Before designing the method, the authors quantified what the drafter should predict for the gain to materialize. To do this, they use a live oracle that peeks at the target’s actual future. They compare three kinds: a position oracle that reveals the true next position, a token oracle that reveals the true token, and a joint oracle that gives both (§3.1).

Original Figure 2 (live oracle, SDAR-8B b32). The left panel shows the throughput ratio versus draft-free by oracle type and budget, and the right panel shows the probability that committed tokens in the next $K$ steps fall within the stale top-$t$ candidates. On the left, read that the position oracle lies above the token oracle on every task; on the right, read that even a few candidates cover almost all future tokens.

The results are clear. Knowing the position is always more beneficial than knowing the token. Going from D8 to D20, the position oracle’s NFE reduction grows from 5.4×→7.2× on GSM8K and from 5.7×→8.2× on MBPP, whereas the token oracle reaches only 3.4–6.5× at D20. A perfect drafter (joint oracle) reaches a throughput of 6.2–7.0× (§3.1, Figure 2).

The reason is already hidden in the stale distribution. The probability that the stale argmax matches the actually committed token is 92% at the first block commit and remains above 90% thereafter. But “the probability that the position the model was most confident about in the last round will be committed next” falls to 3% by block 9. The tokens are almost already known; what is truly unknown is the order. A model trained on this order predicts the next commit position with 75% accuracy, whereas reading the order from stale confidence reaches only 22% (§3.1, Appendix).

Method: a drafter that predicts only order, and coupled-noise verification

SPOT (Speculating Positions Order Trajectory) is split into two parts: what to draft (the order model) and how to verify it (coupled noise).

Drafter: commit-order model

The drafter is a 3-layer Transformer encoder ($d_{model}{=}128$, 8 heads, 0.62M parameters). It reads the stale top-32 probabilities (raw and log) for each masked position, summary statistics, and path features, and predicts the distribution over the next position to be committed $q_\phi(s|y)$. It uses no target hidden states, token embeddings, or text at all (§4.2).

It does not predict tokens. Candidate tokens for each position come from the stale distribution. At $T{=}0$, it uses the stale argmax as the primary candidate, and if the stale top-1 probability is below 0.6, it adds one more second-ranked candidate (token hedge). This exception reduces cases where, even with the right order, the wrong token is rejected during verification (§4.2).

The key is the prefix-conditional approach. When the drafter predicts the next position in some virtual state, it then predicts the next position again in the new virtual state after filling that position, and repeats this recursively to grow a tree of candidate states. The criterion for expanding the tree is prefix mass, defined as child mass = parent mass × (probability that this position will be committed) × (probability that this token will be filled), and since all factors are in $[0,1]$, a child cannot have greater mass than its parent. Thanks to this monotonicity, a search that selects the top-$D$ highest-mass candidates level by level (DDTree) produces the same result as sequential best-first search (§4.2, Appendix).

  flowchart LR
  A["Current block state<br/>some slots committed, some masked"] --> B["Read stale distribution<br/>probabilities for each masked position"]
  B --> C["Order model: predict next commit position"]
  C --> D["Expand candidate tree<br/>child = position + token committed"]
  D -->|"Recurse until budget D"| C
  D --> E["One batched target forward<br/>verify root + all drafts"]
  E --> F["Commit only the longest matching path"]
  F --> A

Original Figure 1 (SPOT overview). (a) The target commits the single most certain position per forward, but the tokens are largely known in advance and only the order is unknown. (b) SPOT reads the stale distribution, predicts the next position with the 0.62M order model, grows a state tree by prefix mass (edge thickness), verifies all of it with a single batched forward, and commits three positions.

Verification uses state equality. Each draft state is appended as a block copy after the cached prefix, and forwarding $(1{+}D)L$ positions at once yields the target’s next commit for the root and every draft state in one pass. A draft is accepted only when its next commit matches the commit the target actually produces at that parent. Since the target’s next commit is unique, what is accepted is always a single path. Therefore, the output never changes no matter how wrong the drafts are — the drafter determines only speed and does not affect correctness (§4.1, §4.2).

Verification: coupled noise that preserves correctness even under sampling

At $T{=}0$, commits are deterministic, so state-equality verification holds as is (Proposition 1). The problem is sampling at $T{>}0$. Drawing samples independently from the same distribution can produce different tokens, so even a correct draft is rejected because its state differs.

SPOT’s solution is for the drafter and the target to share the same Gumbel random numbers. Drawing a token at temperature $T$ is mathematically equivalent to drawing a Gumbel random number $G_{a,v}$ for each position-token pair and choosing $\argmax_v[(1/T)\log p + G_{a,v}]$ (the Gumbel-argmax identity). If both sides share the same table of random numbers on an absolute-position basis, then once the random numbers are fixed, the target’s draw becomes deterministic. Sampling verification then reduces once again to state equality. It is exact conditional on the random numbers $G$, and integrating over $G$ recovers the original sampling distribution exactly (Proposition 2, §4.3).

This coupled verification matters more as temperature rises. At $T{\le}1$, the drawn token is the argmax 99–99.8% of the time, so the difference between the two methods is within 2%, but at $T{=}2$, independent-sample verification accepts almost nothing (257.5 NFE versus 267.1 draft-free), whereas coupled verification keeps drafting (187.2 NFE). The acceptance probability of coupled verification is lower-bounded by the total variation distance $\delta$ between the stale and fresh distributions: at least $(1-\delta)/(1+\delta)$ (§3.2, §4.3).

SPOT-Sets: on decoding that commits several slots at once

So far, this has been exact decoding that commits one position per forward. But faster decoders commit all positions whose probability exceeds a threshold $\tau$ in a single forward (dynamic-threshold decoding). In that case, the unit of drafting becomes a commit set rather than a single position, and the drafter must also determine the tokens. That is because the stale argmax is 83% correct one step ahead but only 44% correct three steps ahead (§4.4).

SPOT-Sets uses a 5.1M-parameter drafter for this. It additionally reads embeddings of the stale top-16 candidate tokens and the target’s last hidden state, and predicts whether each position will exceed $\tau$ at the next step and which candidate token to choose. A node in the tree is not a single position but an entire target step (§4.4).

A small example (illustrative example): in threshold decoding, one forward commits several positions at once, such as {3,4,6}. SPOT-Sets drafts this entire set, and a draft that fills the whole block continues proposing into the next block.

Original Figure 3 (max-commit). (a) The threshold decoder commits the positions $S_\tau$ that exceed $\tau$. (b) Exact acceptance accepts only when the draft step equals $S_\tau$, subset/cut accept a subset of it, and product accepts any step whose geometric mean confidence is at least $\tau$. (c) Max-commit follows the path with the most commits: 7 commits under exact and 9 under product, and product’s output differs from the threshold decoder’s own output.

Here the trade-off between correctness and speed appears again. Exact acceptance keeps the output identical to the threshold decoder, but the authors’ main results use the looser product rule. Product accepts up to the step where the geometric mean confidence is at least $\tau$, committing more per forward, but it no longer reproduces the threshold decoder’s output. The paper therefore also reports an exact variant. To reduce the cost incurred every few forwards at block boundaries, it adds two block-boundary optimizations: cross-block drafting and draft refresh (§4.4, Appendix).

Results: read forward count and throughput separately

When looking at the results, do not conflate the two metrics. NFE is the number of target forwards (the drafter’s own cost is excluded), and the throughput ratio is based on generated tokens per second relative to a draft-free run on the same GPU, same task, and same temperature. One draft round forwards $(1{+}D)L$ positions regardless of acceptance, so the NFE reduction and the throughput gain are not always the same — a recurring point in this paper (§3.1, Appendix).

The key question is this. With the same draft budget, does SPOT use fewer forwards and also achieve higher throughput than the baselines? Below are NFE (lower is better) and throughput ratio (higher is better) for SDAR-8B b32, $T{=}0$, budget D8 (§5.2, main results table).

SPOT uses 16–36% fewer forwards than Spiffy across all tasks and budgets and has 1.04–1.41× higher throughput. Notably, the baselines’ throughput collapses as the budget grows. Going from D8 to D20, the SSD chain falls from 2.3–2.6× to 1.4–1.7× and the BlockSpec tree falls from 1.5–1.7× to 0.86–1.10×, becoming slower than draft-free on MBPP. This is because wide drafts consume $(1{+}D)L$ positions regardless of acceptance, yet accept only 1.1–2.9 tokens per round. SPOT accepts 2.5–7.5 tokens per round (§5.2).

Here SPOT-$\gamma$, with the mass floor $\gamma$ enabled, almost eliminates budget dependence. It discards drafts whose mass is below $\gamma$, matching the actual verification budget to drafter confidence, and the throughput ratio stays nearly flat, from 2.5–3.5× at D8 to 2.3–3.4× at D20. The same ordering holds under $T{=}1$ sampling (16–35% lower NFE and 1.08–1.37× throughput versus Spiffy), and the T=1 results are distribution-preserving, with correctness within 0.04 of draft-free (§5.2, §5.3).

7.9× on threshold decoding

Stacking SPOT-Sets on faster threshold decoding (which by itself is already 3.3–6.1× versus standard, 5.0× on average) raises SDAR-8B to 5.2–9.9×, 7.9× on average. Relative to the threshold decoder, it cuts forwards by 47–48%. The exact variant that preserves correctness also reaches 7.6× (§5.2).

Original Figure 4 (ladder). Throughput versus standard decoding, averaged over four tasks, as each part of SPOT-Sets is enabled in turn (labels show the gain at each step). The set drafter and tree alone reach 6.7×; the mass floor raises it from 6.7→6.9×; the block-boundary optimizations cross-block drafting and draft refresh add the largest gain, from 6.9→7.5×; and product max-commit finishes at 7.9× (§5.5, Figure 4).

Transferability is also central to the claim. Even when a drafter trained with the same recipe is attached to another model without retraining or recalibration, moving the SDAR-8B checkpoint to LLaDA2.0-mini (an MoE model with a different tokenizer) loses only 3–11% NFE and still beats all baselines on GSM8K and MATH500. On LLaDA-8B and Dream-7B, which it has never seen, it reduces 256 forwards to 55–88 and 51–80, respectively. However, these two apply attention over the entire sequence and lack an exact prefix cache, so they use an approximate cache; these rows therefore show “reach” rather than losslessness (§5.4, Appendix).

Limitations and deployment judgment

The authors explicitly acknowledge several limitations (§6).

1. The gain shrinks when verification is expensive or the acceptance rate is low. With a wide block, verification becomes the bottleneck, and the problem of packing the tree into exactly one sequence under bidirectional attention remains unsolved.
2. The gain narrows on tasks with low acceptance rates. On the smallest model, SDAR-1.7B, Spiffy takes one throughput win from SPOT on MATH500 D20. This is because forwards are so cheap that the cost of draft verification becomes a larger share. SPOT-γ discards low-mass drafts to reduce verification volume, taking that win back (Appendix).
3. At high temperature, an order model trained at $T{=}0$ is overtaken by noise-aware ranking. This is true at $T{\ge}2$, and in fact from $T{\ge}1.5$ the target’s own quality collapses (correctness 0.11, 0.00 at $T{\ge}2$), leaving little correctness to preserve.
4. SPOT-Sets requires brief calibration per model. Because it reads hidden states and token embeddings, unlike SPOT it must re-fit the input layer for each new target. The 5.1M-parameter drafter becomes a greater burden on models where target forwards are cheap.
5. The main results’ max-commit (product rule) changes the output. To preserve correctness, one must use the exact variant, which somewhat reduces the speed gain.

From a reviewer’s perspective, the points to note are as follows. The 6.8× NFE reduction shown in the paper is on a “forward count” basis, while actual throughput is 2.5–3.4× (SPOT) or 7.9× on threshold decoding; the two numbers must not be confused. Also, the SDAR-8B drafter is trained only on code data (CodeAlpaca), yet it works well even though GSM8K and MATH500 are out of domain, which supports the generality of order knowledge; conversely, the SPOT-Sets drafter trained on a distribution like the evaluation tasks is faster on SDAR-8B, at 8.2× on average (7.9× versus a general chat drafter), so the distribution of the deployment environment affects the results (§5.2, Appendix). The LLaDA-8B and Dream-7B results use an approximate cache and thus do not guarantee losslessness; this must also be distinguished from the lossless results.

In summary, SPOT fits best when lossless acceleration is needed on block-causal diffusion models that have an exact prefix KV cache and do not have excessively wide blocks. Coupled-noise verification preserves the distribution at any temperature (Proposition 2), and in the typical generation setting of $T{\le}1$ its results are within 2% of independent-sample verification, making it the safest deployment configuration. On fast paths that already accept some loss, such as threshold decoding, SPOT-Sets’ product rule provides additional gain; according to the paper, the correctness difference versus exact is within 0.04 per task, but the output itself changes (example rule: exact leaves correctness unchanged, while cut only makes the decoder stricter), so quality must be verified separately when deploying.

Figures in this post are taken from the original arXiv:2610.22543 (CC BY 4.0). Only size and format were changed.