Contextual Correction of Parallel Draftsfor Speculative Decoding

Paper

Contextual Correction of Parallel Drafts — DCOCO: How to Recover the Lost Inter-Token Dependencies in Speculative Decoding

TL;DR — Parallel diffusion drafters like DFlash produce a whole block of tokens at once, but within the block they predict each token without knowing the identity of the preceding tokens. DCOCO is a lightweight corrector that re-selects contextually appropriate tokens by scoring the adjacent transitions among already-drafted candidates in a single batched operation, with no extra drafter forward pass. Across 4 backbones and 10 benchmarks it is up to 6.6× faster than vanilla and up to 18% faster than the strongest speculative baseline, and it retains a 6–23% margin even at concurrency of 4–32 streams.

Core Idea

Speculative decoding accelerates LLM inference with the idea that “a lightweight drafter proposes tokens and a heavy target model verifies them in one pass.” The problem is the tension between drafting speed and acceptance rate.

The authors focus on exactly one thing. The currently popular parallel drafters (which predict an entire block at once in a diffusion-based manner, as DFlash does) are fast, but within a block they do not see the actual selection made in the preceding slot. As a result, tokens that look plausible individually combine to simultaneously assign high probability to “New York” and “San Diego,” producing an unnatural continuation like “New Diego” (§2). This lowers the acceptance rate and ultimately erodes the speed gain.

Existing methods such as DSpark and Domino try to solve this problem with a sequential module, but that sequential module sits on the critical path of every decoding cycle. DCOCO (DFlash Contextual Co*rrector) answers “yes” to the question of whether intra-block dependencies can be restored without reintroducing a sequential module (§1).

The key lies in two observations.

1. The drafter already has the answer. For Qwen3-8B, in 59% of slots where top-1 is wrong, the verifier’s actual token is present among the top-16 candidates (§2). Correction can therefore be reduced not to “generating new tokens” but to re-selection among candidates that already exist.
2. Verification stops at the front. Speculative verification accepts only a contiguous prefix and halts at the first rejection. Improving later slots therefore matters only if all earlier slots are accepted, so correction must concentrate on the quality of the early prefix (§2).

Combining these two turns the corrector into the following task: “keep the top $K$ candidates per slot, score all $K^2$ transitions between adjacent slots in parallel, then greedily choose from the front.” The decisive difference from DSpark and Domino is that neural scoring is parallel and only the selection is sequential (a simple walk over the grid).

Background: the Problem They Solved

Structure of Speculative Decoding

The verifier scores the proposed block $\hat y_{1:B}$ in a single forward pass, decides slots from left to right, commits the longest accepted prefix, and then adds one token of its own. Under greedy decoding, slot $j$ is accepted if it matches the verifier’s greedy choice; under stochastic decoding it is accepted by the standard speculative sampling rule $\min\{1, q_j(\hat y_j)/p_j(\hat y_j)\}$ (§3.1).

Writing the acceptance of slot $j$ as $a_j\in\{0,1\}$, the acceptance length is

and each round commits $A(\hat y)+1$ tokens. Throughput is ultimately determined by the expected acceptance length and the wall-clock time per round. And this single structural fact — “verification stops at the first rejection” — governs every design choice in DCOCO, because if slot 1 is rejected, the corrected token in slot 5 is never even examined (§3.1).

Candidate Coverage: the Missed Headroom

The authors recorded DFlash’s candidates during $T{=}0$ speculative decoding and compared them against the actual continuation to measure top-$k$ coverage — “the fraction of draft slots whose top-$k$ candidates contain the token the verifier ultimately commits” (§2).

The results are regular. On Qwen3-8B, at the first slot top-16 is about 95.1% and top-1 (the drafter’s own choice) is 81.4%, a small gap, but at the last slot top-16 is still 60.8% while top-1 falls to 15.4% (Fig. 1, Tab. 7). Averaged over the whole block, top-1 is 40.0% vs top-16 75.5%. That is, at deep slots nearly one in two slots has the correct answer in its candidate pool, yet the drafter increasingly fails to pick it. This “missed correct answer” is exactly the headroom the corrector re-selects from.

The New Approach: DCOCO

DCOCO operates in two stages. (1) Given the realized tokens, it re-assigns a stronger contextual score to the drafter’s per-slot candidates, and (2) it re-selects the block according to those scores (§3).

1) A Learned Corrector over the Block Grid

At each slot $i\in\{1,\dots,B\}$, DFlash leaves a hidden state $h_i\in\mathbb{R}^d$ and the top $K$ candidates. Here $c_{ik}$ is the token id of the $k$-th candidate at slot $i$, and $\ell_i(k)$ is its draft logit. These candidates form a dense first-order grid of $B\times K$ nodes, with every candidate at slot $p$ connected to every candidate at slot $p+1$.

Each transition is assigned a score:

where $\Delta_\theta$ is a shared 2-layer MLP, $z_{ik}$ is a candidate representation, and $s_p$ is a candidate-independent pair state (§3.2). The corrector reads only the hidden states, candidate ids, and logits left by a single DFlash pass.

• Candidate representation $z_{ik}$: it fuses $b_{ik}$, obtained by projecting and normalizing the embedding of $c_{ik}$ from the frozen token embedding table $E$, with $g_i$, obtained by projecting and normalizing the slot hidden state $h_i$, via concatenation. Because the token embedding and hidden state differ in scale by more than two orders of magnitude, it uses a non-affine layer normalization.
• Pair state $s_p$: it concatenates the hidden states of adjacent slots, $s_p = W_{\text{pair}}[h_p; h_{p+1}]$. Since $s_p$ depends only on slots, it is computed once and reused across all $K^2$ transitions between them.

There are two key design tricks. First, the MLP’s output layer is initialized to zero, so initially $\Delta_\theta\equiv 0$ and every transition score reduces to the successor’s draft logit $S_p(a,b)=\ell_{p+1}(b)$, thereby reproducing the drafter’s original proposal exactly at initialization (§3.2). Second, all projection and MLP parameters are shared across slots and transitions, and the $(B-1)K^2$ transition scores are computed in a single batched operation before path selection.

2) Prefix-first Greedy Selection

On the scored grid, simply maximizing the total path score lets high scores at later slots compensate for low scores at earlier ones — but later tokens contribute only if the earlier prefix survives. Because of this asymmetry, DCOCO uses a prefix-first greedy rule:

The first slot keeps the drafter’s top candidate, and each subsequent slot picks its successor only from scores conditioned on the selected predecessor. The absence of lookahead into future transitions is actually an advantage: this walk costs only $O(BK)$ (added on top of the $O(BK^2)$ scoring), and if $\Delta_\theta\equiv 0$ it agrees exactly with the drafter’s argmax (§3.3).

3) Training: Supervise Only on Verifier-Correct Prefixes

The drafter and corrector are optimized jointly. The label $(a^\star_p, b^\star_p)$ for transition $p$ is the pair of verifier argmax tokens at the two slots under a teacher-forced context, and it is activated with mask $m_p=1$ only when it lies within each slot’s top-$K$. The correction loss multiplies a cross-entropy, normalized over the $K^2$ transitions by a joint softmax, by an exponential decay weight:

This softmax exists only during training; at inference time $S_p$ is used without normalization. The exponential weight is a fixed proxy reflecting the prefix-survival structure — “a later transition contributes only if the earlier slot is accepted” (§3.2, App. A). The supervision simply drops rather than suppresses transitions the drafter finds unreachable, since the verifier-correct prefix is the only state in which the corrector’s output is actually consumed (App. A).

4) Deterministic Proposals Even Under Stochastic Decoding

This is an interesting theoretical result. DCOCO selects tokens deterministically even under stochastic decoding, so each slot’s proposal distribution becomes a point mass $p_j(v)=\mathbf{1}[v=\hat y_j]$. The acceptance probability of standard speculative sampling then simplifies to $\min\{1, q_j/p_j\}=q_j(\hat y_j)$, and on rejection the residual distribution becomes $r_j(v)=q_j(v)\mathbf{1}[v\neq\hat y_j]/(1-q_j(\hat y_j))$. Repeating this at every slot preserves the verifier distribution exactly — that is, stochastic target decoding does not require stochastic draft selection (App. D). No softmax normalization on transition scores and no candidate-distribution sampling are needed.

How It Works: Walking Through a Concrete Example

Let us reproduce the “New Diego” problem in miniature. Say the block size is $B=4$ and the number of candidates per slot is $K=3$. Suppose a single DFlash pass left the following (the same example tokens and scores as Fig. 2 in the paper).

• Slot 1: These(4.2), The(3.6), A(2.8)
• Slot 2: key(0.4), keys(2.7), locks(1.1)
• Slot 3: is(0.2), are(3.1), was(0.8)
• Slot 4: missing(2.8), here(1.6), broken(0.5)

If the drafter simply picks each slot’s top-1, the result is “These key is missing” — every word is plausible, yet the grammar is off. The identity of the problem is that the adjacent context — that “These” should be followed by “keys” and “keys” by “are” — is not reflected in the drafter at all.

DCOCO lays this block out into a $4\times 3$ grid. It builds each slot’s candidate representation $z_{ik}$ (embedding + slot hidden state) and pair state $s_p$, and scores all $3\times 3=9$ transitions between adjacent slots with the shared MLP $\Delta_\theta$. For instance, the 9 transitions of (Slot1→Slot2) are all computed in parallel as:

Since during training $\Delta_\theta$ was trained to assign high scores to adjacent pairs that the verifier actually selects, such as “These + keys” and “keys + are,” it adds a correction to the draft logit $\ell_2(\text{key})=0.4$ so that $S_1(\text{These}, \text{keys})$ overtakes $S_1(\text{These}, \text{key})$.

Then comes the prefix-first walk. Slot 1 is fixed to top-1 These (→ These). Slot 2 picks keys, which maximizes $S_1(\text{These}, \cdot)$ (→ These keys). Slot 3 now picks are, which maximizes $S_2(\text{keys}, \cdot)$ conditioned on the selected keys (→ These keys are). Finally, Slot 4 picks missing via $S_3(\text{are}, \cdot)$. The final proposal is “These keys are missing” — an ungrammatical string produced by the drafter’s top-1 combination has been turned into a grammatically correct sentence with no extra drafter forward pass.

The important point here is that only the selection is sequential, while the 9 scorings are all bundled into a single batched operation. To fix this example, DSpark and Domino must run a recurrent module that sweeps slot by slot, whereas DCOCO finishes with one parallel scoring plus an $O(BK)$ walk.

Performance Validation: Main Results

Setup

The verifiers are four backbones — Qwen3-4B, Qwen3-8B, Qwen3-Coder-30B-A3B, Llama-3.1-8B — and the benchmarks are 10 covering math (GSM8K, MATH-500, AIME-25), code (HumanEval, MBPP, LiveCodeBench), and general (MT-Bench, Arena-Hard-v2, Alpaca, SpeedBench). All drafters were trained with block size 16 on about 800K prompts (stem/math/code from Nemotron Post-Training v2 + CodeAlpaca-20k). Serving uses vLLM 0.27.0 on a single NVIDIA B200, BF16, TP=1 (App. C). The primary metric is end-to-end tokens per second (TPS), and the secondary metric is the average acceptance length $\tau$, which separates “quality improvement” from “serving overhead” (§4).

The Right Conclusion: Why You Should Read $\tau$, Not TPS

The most honest choice in this paper is making TPS the primary metric. The corrector can increase acceptance length while eroding wall-clock time with its own scoring cost, and that trade-off shows up only in TPS (§4). The paper therefore reports $\tau$ (mechanism) and TPS (outcome) side by side.

On Qwen3-8B (thinking off), DCOCO achieves an average acceptance length $\tau$ of 4.985, which is +8.1% over Domino (4.613), +17.2% over DSpark (4.254), and +27.4% over uncorrected DFlash (3.913). And this $\tau$ improvement translates directly into a TPS improvement (Fig. 3, Tab. 2).

Overall Results: First Place in All 12 Cells

In all 12 combinations of four backbones × (thinking on/off) × ($T{=}0$/$T{=}1$), DCOCO recorded the best average speedup (§4.1, Tab. 2). In each cell it leads the strongest baseline by 8–19%. For Qwen3-8B, math and code benchmarks go up to 3.7–6.1× over vanilla, while open-ended chat is 1.8–1.9× (greedy decoding). The up to 6.6× in the abstract comes from Qwen3-4B on AIME-25.

Averaging over the 10 suites for Qwen3-8B (thinking off, $T{=}0$):

That is, it is 14.3% faster than the strongest competing corrector Domino and 22.6% faster than uncorrected DFlash (Tab. 1).

The Decisive Difference Revealed Under Concurrency

LLM serving batches concurrent requests to raise GPU utilization. As concurrency grows, the verifier’s cycle time is amortized across streams, so the share of overhead the corrector adds per cycle keeps growing. This is exactly the graveyard of sequential correctors.

On Qwen3-8B (thinking off), DCOCO records the highest throughput across the entire range of concurrency from 4 to 32 streams, with a margin over the strongest baseline of +16.6%→+23.0% at $T{=}0$ and +16.1%→+12.4% at $T{=}1$ (Tab. 8). By contrast, Domino and DSpark keep an almost unchanged acceptance length, yet accumulate sequential-correction overhead and eventually become slower than uncorrected DFlash.

Fig. 3 makes this phenomenon clear. The vertical axis is the ratio relative to DFlash, the solid line is the acceptance length $\tau$, the dashed line is throughput, and the shaded region is the gap between them (= the corrector’s per-cycle overhead). Domino’s $\tau$ curve stays above parity (1.0), but its throughput curve slides below — overhead has swallowed the acceptance gain. In contrast, DCOCO’s two curves move nearly together with a thin shaded region and remain above parity. This phenomenon reproduces across all four backbones (App. B).

Path Selection and Ablation

Comparing path selection rules (§4.3) on the same learned score grid, Global-Sum (maximizing the total sum) yields lower $\tau$ and TPS than greedy on all 9 suites. Geometric discount ($\lambda=0.75$) is comparable to greedy in acceptance length (average 5.258 vs 5.254), but greedy achieves +2.8% on average higher throughput. In the end, greedy wins with its $O(BK)$ walk and no discount hyperparameter (Tab. 3).

Ablation (App. A) validates the corrector’s two input paths.

• Removing the pair state: acceptance length $-$7.3%/$-$5.6% ($T{=}0$/$T{=}1$), throughput $-$6.3%/$-$4.6%.
• Removing the candidate representation: acceptance length $-$4.8%/$-$3.7%, throughput $-$4.6%/$-$3.3%.
• Freezing the backbone (fixing the drafter for 6 epochs and training only the corrector): throughput $-$15.9%/$-$12.9%.

That is, both input paths are essential, and their contribution is largest on math and code. Co-adaptation is also key: even with a frozen backbone the corrector is a pure drop-in that yields +4.4% TPS over DFlash, but it is 15.9% below joint training. Conversely, extending the baselines to 12 epochs improves them by no more than 2%, confirming that DCOCO’s gain comes from the corrector itself, not from “training budget” (Tab. 4).

Our Take: Strengths, Limitations, and Why This Research Matters

Strengths

1. A clean answer to “restoring dependencies without sequencing.” Whereas DSpark and Domino take the compromise of “parallel draft + sequential correction,” DCOCO revives dependencies without sacrificing parallelism. The separation of fully parallel neural scoring with only an $O(BK)$ selection walk is this paper’s most elegant contribution (§1, §3).
2. Redefining the problem as “re-selection.” It is decisive that correction is reduced not to “generating new tokens” but to “re-selecting among the existing top-$K$ candidates.” The measurement of top-16 coverage 75.5% vs top-1 40.0% (Fig. 1) shows that this redefinition is not vacuous but grounded in real headroom.
3. Honesty about metrics. The design of using TPS as the primary metric and $\tau$ as the mechanism metric probes exactly the trap of correctors that “raise acceptance length but lose in wall-clock time.” The fact that the concurrency experiment demonstrates this trade-off makes the paper’s conclusion a system-level claim, not a single benchmark number.
4. Low cost. On Qwen3-8B the corrector is only about 2.23M parameters (drafter $d=4096$, 5 layers), and its per-cycle cost is merely 3–4% over uncorrected (App. A, §4.1). The concluding sentence “a correction made cheap enough is pure profit” is backed by the numbers.

Limitations

1. Single GPU, single framework. All experiments were performed on a single NVIDIA B200 + vLLM 0.27.0 + TP=1 (§4, App. C). Behavior under multi-GPU tensor parallelism or other runtimes (e.g., TRT-LLM) is unverified.
2. Single-seed training. Each method was trained once with seed 42, and variance across training seeds was not measured (App. A). The standard deviation for $T{=}1$ is only over “3 generation seeds,” not variability across independent training runs.
3. Room to argue over baseline fairness. Comparisons use “each method’s native configuration,” and DCOCO uses 6+6 epochs (drafter 6 + joint 6) whereas the baselines use 6 epochs. The authors defend fairness by showing that extended training gives the baselines less than 2% gain (Tab. 4), but this can still be a point of contention for readers demanding “exactly the same block size and the same budget.”
4. Drafter dependence. DCOCO is built on top of a specific parallel drafter, DFlash. Transfer to other parallel/diffusion drafters is natural but not empirically demonstrated, and sensitivity sweeps for hyperparameters such as block size, $K$, and $\rho$ are limited.
5. Limits on open-ended chat. Math and code speed up to 3.7–6.1×, but open-ended chat like MT-Bench and Alpaca stays at 1.8–1.9× (§4.1). This is a limitation of the drafter’s base acceptance rate itself, a structure DCOCO amplifies rather than overcomes.

Why It Matters

The real contribution of this research is not a specific number but a perspective that flips the problem. It replaces the implicit premise that “reversing the weakness of parallel drafting (intra-block independence) requires sequencing” with “the correct answer is already in the candidate pool, and scoring adjacent transitions in parallel restores it without sequencing.” And by showing that this replacement shines even under concurrent load — indeed all the more so — it draws the practical conclusion that “making correction cheap is the pure profit of speculative decoding.” At a time when speculative decoding is becoming standard in real serving stacks, this approach, which fills the last weakness of parallel drafting, is worth attention.

What’s Next?: The Road Ahead

The authors do not explicitly lay out a grand future roadmap, but next steps follow naturally from the paper’s limitations and structure.

1. Generalization to other drafters and backends. The most direct task is to verify whether DCOCO’s “parallel scoring + prefix-first walk” still holds an advantage on parallel drafters other than DFlash (e.g., other diffusion/CTC-based ones) and under multi-GPU and other runtimes.
2. Systematic sweeps of $K$, block size, and $\rho$. The number of top-$K$ candidates and the exponential decay $\rho$ are the knobs that directly control the trade-off between “candidate headroom” and “scoring cost.” Quantifying their cost-acceptance curves would yield tuning guidance for deployment.
3. Adaptive candidate counts. It is promising to exploit the fact that coverage drops sharply with slot depth (Fig. 1) in reverse, reducing scoring cost by using a small $K$ for shallow slots and a large $K$ for deep slots.
4. Acceptance-rate-based training weights. The current exponential weight with $\rho=4$ is a “fixed proxy” for prefix survival. Estimating actual acceptance probabilities and making the per-position training weights data-adaptive is the logical follow-up (explicitly mentioned in §3.2).
5. Designs targeting batch serving. Since the concurrency experiment already shows that correction overhead is the deciding variable, system optimizations that fuse scoring more finely at the CUDA-kernel level or exploit shared computation within a mini-batch will be the source of the next gain.

In short, DCOCO has offered a new baseline for speculative decoding: “the cheapest way to recover the lost context of parallel drafting.” What comes next will be extending this idea to more diverse hardware, drafters, and load conditions to draw the boundary of the proposition that “correction is always pure profit.”

Tables from the paper

Tables converted mechanically from the arXiv e-print LaTeX source. The numbers are the paper’s own and did not pass through a model.

Table 1. Throughput (output tokens/s, per-stream) and speedup over no-spec decoding, Qwen3-8B, thinking off, concurrency 1, max 2048 generated tokens. Each speculative cell reports TPS with its speedup over the no-spec column in parentheses. Bold: best method per row.

Table 2. All configurations at a glance: average speedup over vanilla decoding, unweighted mean over the ten evaluation suites. Bold: best per row and temperature. Qwen3-Coder-30B-A3B has no thinking mode. Per-dataset throughput and accepted length: App. .

Table 3. Path selection over identical trained scores (Qwen3-8B from-pretrained corrector, $T{=}0$, no thinking, 2048 max tokens). Bold: best rule per row and metric. Between greedy and geometric the $\tau$ differences carry both signs and are an order of magnitude smaller than either rule’s gap to Global-Sum; the throughput ordering is uniform.

Table 4. Extending baseline training to match our total budget, Q8B no-thinking, max generated tokens 2048, T=0 averaged across the 10 benchmarks. Baselines show no more than 2% improvement from a second six epochs; the joint phase accounts for our arm’s improvement.

Table 5. Ablations: throughput (output tokens/s, per-stream), Qwen3-8B, thinking off, concurrency 1, max 2048 generated tokens, all ten suites. Each ablation cell reports its value with the change relative to DCOCO in parentheses. $-$pair state removes the corrector’s dedicated previous-state channel; $-$cand state zeroes the candidate representations entering the transition scorer; Frozen backbone trains only the corrector on the fixed six-epoch DFlash backbone. Bold: best measured system per row.

Table 6. Ablations: accepted length $\tau$, Qwen3-8B, thinking off, concurrency 1, max 2048 generated tokens, all ten suites. Each ablation cell reports its value with the change relative to DCOCO in parentheses. $-$pair state removes the corrector’s dedicated previous-state channel; $-$cand state zeroes the candidate representations entering the transition scorer; Frozen backbone trains only the corrector on the fixed six-epoch DFlash backbone. Bold: best measured system per row.

Table 7. Mechanism: mean accepted length $\tau$, Qwen3-8B, thinking off, companion to Table ; no-spec commits one token per round by construction and is omitted. $\tau$ explains the TPS movement; the corrector’s own overhead is visible only in throughput.

Table 8. Throughput (output tokens/s, per-stream) and speedup over no-spec decoding, Qwen3-8B, thinking on, concurrency 1, max 4096 generated tokens. Same format as Table .

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