PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds

Yunxiao Zhao, Changxiao Cai

arXiv:2610.10411v1cs.LGcs.AIcs.CLstat.ML

TL;DR

Parallel and semi-AR drafters create cross-round dependencies that existing local objectives do not directly capture when optimizing decoding efficiency. The paper models speculative decoding as a Markov reward process and derives EDR, an exact objective for expected decoding rounds. Across nine benchmarks, EDR finetuning consistently improves DSpark and DFly and outperforms E2E in overall mean accepted length.

  • Problem

    For parallel and semi-AR drafters, proposal distributions depend on round starts, making expected decoding rounds the relevant objective amid cross-round dependencies.

    Source: §1 · p. 2

  • Method

    The paper formulates speculative decoding as a Markov reward process and derives EDR, an exact expected-round objective with a temporal-difference gradient for training.

    Source: §1 · p. 3

  • Results

    Across nine benchmarks, EDR consistently improves DSpark and DFly and outperforms E2E training in overall mean accepted length.

    Source: §1 · p. 3, §6 · pp. 13–14

  • Takeaways & Limitations

    Explicitly accounting for cross-round dependencies is a principled and practical direction for training more efficient parallel speculative draft models.

    Source: §6 · p. 14

    The framework assumes a fixed draft block size, and computing occupancy weights and value functions is more expensive than block-local methods.

    Source: §6 · p. 14

Abstract

from arXiv · show

Speculative decoding accelerates large language model inference by using a low-cost draft model to propose tokens that the full-size target model verifies in parallel. Parallel and semi-autoregressive (semi- AR) drafters improve drafting efficiency by proposing an entire block in a single forward pass, but training them raises a new difficulty: the draft distribution for a given position depends on where the decoding round starts, and where rounds start depends on how many tokens earlier rounds accepted. Existing training objectives typically rely on block-local surrogates that ignore this cross-round coupling, and therefore do not directly optimize the global decoding efficiency. In this work, we develop a theoretical framework for training and evaluating these drafters by representing speculative decoding as a Markov reward process. This formulation yields the Expected Decoding Rounds (EDR) objective, which weights local rejection costs by state occupancies and exactly equals the expected number of decoding rounds. Unlike prior surrogate objectives, EDR introduces no auxiliary hyperparameters. We then derive an exact temporal-difference gradient that supports unbiased stochastic optimization from target-model rollouts. The same framework also yields an exact offline evaluator for round counts, enabling paired drafter comparisons on shared target rollouts without running speculative decoding. Finetuning two state-of-the- art drafters, DSpark and DFly, with EDR consistently improves mean accepted length and outperforms existing training objectives across nine benchmarks spanning math reasoning, code generation, and chat.

1 Introduction

Speculative decoding reduces the sequential cost of autoregressive generation, but parallel and semi-AR drafters make local proposal quality insufficient for directly optimizing global decoding efficiency. The paper models cross-round dependencies to derive an exact training objective and evaluator, and reports consistent gains from EDR finetuning across nine benchmarks.

1 Introduction

  • Autoregressive generation requires sequential token production, while speculative decoding proposes multiple tokens for parallel target verification and parallel drafters produce blocks in one pass.Semi-AR variants add limited within-block dependence; decoding rounds correspond to expensive target-model verifications.Source: §1 · p. 2
  • Parallel and semi-AR proposals depend on the round’s starting position, so earlier acceptance outcomes affect which draft distributions are encountered in later rounds.This cross-round coupling breaks the direct alignment between local distribution matching and global decoding efficiency.Source: §1 · p. 2
  • Existing parallel and semi-AR drafter objectives use local surrogates, such as within-round accepted-token counts or hand-weighted draft-target discrepancies, rather than directly minimizing expected rounds.The paper poses direct minimization of expected decoding rounds as its central question.Source: §1 · pp. 2–3
  • EDR minimizes expected decoding rounds exactly, with a temporal-difference gradient for training and an offline evaluator for paired drafter comparisons on shared target rollouts.The framework also proves that maximizing expected accepted length in each round can be strictly suboptimal for EDR.Source: §1 · p. 3
  • Across nine benchmarks, one EDR finetuning epoch consistently improves mean accepted length for DSpark and DFly over the original drafters and existing training objectives.The drafters’ architectures and inference procedures remain unchanged.Source: §1 · p. 3

2 Preliminaries

This section defines the target-rollout and speculative-decoding setup, distinguishes drafter conditioning patterns, and introduces decoding rounds and mean accepted length. Rejection-sampling verification preserves the target distribution while determining how drafts are accepted or corrected.

2 Preliminaries

  • The notation defines the token vocabulary to include EOS and introduces total variation distance for comparing target and draft distributions.It also specifies sequence-prefix and segment notation.Source: §2 · p. 3
  • The target model samples tokens autoregressively until EOS, and the setup assumes a universal constant upper bound on target-sequence length.The complete sequence includes EOS.Source: §2 · p. 3
  • A draft block is proposed from the committed prefix, with draft distributions potentially conditioned on earlier within-block proposals.The notation suppresses dependence on the sampled sequence for brevity.Source: §2 · p. 4
  • Rejection-sampling verification accepts or corrects draft tokens from left to right; rejected blocks are discarded, and an all-accepted block can yield a bonus token.The procedure preserves the target model’s output distribution for any drafter.Source: §2 · p. 4
  • An AR drafter conditions on the full prefix, a parallel drafter conditions on the round-start prefix, and a semi-AR drafter can partly use earlier within-block drafts.Parallel and semi-AR draft distributions can differ at the same output position depending on the round start.Source: §2 · p. 4
  • A decoding round uses one target-model forward pass, and mean accepted length is the ratio of expected output length to expected decoding rounds.For a fixed target model and evaluation corpus, expected output length is drafter-independent, so maximizing mean accepted length is equivalent to minimizing expected rounds.Source: §2 · p. 4

3 Main Results

The main results frame speculative decoding as a Markov reward process and use state occupancies and local rejection costs to characterize expected decoding rounds. The framework also provides a temporal-difference training gradient, offline evaluation statistics, and a proof that block-local accepted-length objectives can be suboptimal.

3.1 Speculative Decoding as a Markov Reward Process

  • Conditioned on an output sequence, speculative decoding is a Markov reward process whose states track the round start and verification position, with rejection incurring unit cost.The expected round count follows from the occupancy-weighted sum of state costs.Source: §3.1 · p. 5

3.2 The EDR Objective

  • The EDR objective weights local rejection costs by state occupancies and exactly equals the expected number of decoding rounds.Its EOS-aware costs and actual drafter-dependent occupancies distinguish it from heuristic weighting in prior surrogate objectives.Source: §3.2 · p. 6

3.3 TD-Form Gradient of EDR

  • The EDR gradient has a one-step temporal-difference form, and averaging it over target rollouts gives an unbiased stochastic gradient for minimizing expected decoding rounds.Its terms capture immediate rejection-cost changes and downstream acceptance effects weighted by the value gap.Source: §3.3 · pp. 7–8

3.4 Exact Offline Evaluation

  • Exact offline evaluators compute decoding statistics for any drafter from target-model rollouts, without running speculative decoding.The framework covers expected round counts, MAL, rounds committing at least k tokens, and position-wise acceptance rates.Source: §3.4 · p. 8, §3.4.2 · pp. 8–9

3.4.2 Rounds Committing at Least k Tokens and Acceptance Rate

  • The expected number of rounds committing at least k tokens yields an estimator of position-wise acceptance rates.A round commits at least k tokens when its first k−1 drafts are accepted.Source: §3.4 · p. 8, §3.4.2 · pp. 8–9

3.4.1 Expected Decoding Rounds and MAL

  • On shared target rollouts, the expected-round estimator is unbiased and the MAL estimator is consistent, enabling paired drafter comparisons free of rollout variation.The rollouts do not depend on the drafter, so each drafter can be evaluated on the same dataset.Source: §3.4.1 · p. 8

3.5 Suboptimality of Block-Local Surrogates

  • Maximizing accepted length independently within each block can be strictly suboptimal for the global MAL because current acceptances affect where later rounds start.The theorem establishes this suboptimality for first-order drafters, while the text attributes the gap to cross-round interactions.Source: §3.5 · p. 9

4 Related Work

Related work spans parallel and semi-autoregressive drafter architectures and training losses, while the comparison of objectives distinguishes local surrogates from EDR’s cross-round-aware weighting and gradient.

4.1 Literature Review

  • Parallel and semi-AR drafters reduce sequential drafting, with DSpark and DFly combining parallel backbones and sequential heads for within-block consistency.Diffusion language models provide a basis for many parallel and semi-AR drafters, though conditionally independent sampling can trade off generation quality and sampling efficiency.Source: §4.1 · pp. 9–10
  • Drafter training ranges from CE/KL distribution matching to acceptance-aware and multi-token losses, but these objectives remain surrogates for mean accepted length.Examples include DFlash’s fixed position weights, D-PACE’s adaptive weights, E2E’s within-block acceptance structure, and verification-feedback methods.Source: §4.1 · p. 10

4.2 Comparison of Training Objectives

  • Token-level CE, KL, TV, and LK objectives align or compare draft and target distributions independently at each position.LK interpolates between KL and TV according to acceptance probability.Source: §4.2 · p. 11
  • Block-level surrogates weight positions within a block but still treat anchors equally, thereby ignoring cross-round coupling.EDR’s conditional-survival position weight is exact in expectation, unlike the E2E product surrogate.Source: §4.2 · p. 12
  • EDR’s occupancy-weighted local costs equal expected decoding rounds, and its drafter-dependent anchor weights capture how acceptance changes later round starts.The TD form has the same drafter-parameter gradient and is used for more efficient training.Source: §4.2 · p. 12

5 Experiments

The experiments compare EDR with E2E finetuning for DSpark and DFly using shared target trajectories. EDR achieves higher or equal mean accepted length across all nine benchmarks.

5.1 Experimental Setup

  • The evaluation pairs Qwen3-4B with DSpark and Qwen3-8B with DFly, using block size B = 7 for both drafters.Training finetunes only drafter parameters and compares EDR with E2E using the same architecture, data, optimizer configuration, and update steps.Source: §5 · p. 12, §5.1 · pp. 12–13

5.2 Experimental Results

  • EDR yields higher or equal MAL than E2E on every benchmark for DSpark and DFly, with larger gains on AIME25, LiveCodeBench, and chat benchmarks.The comparison uses one-epoch finetuning results and does not modify either drafter’s architecture.Source: Table 2 · p. 13, §5.2 · p. 13

5.1 Experimental Setup

  • Nine math, code, and chat benchmarks use at least 1,500 target trajectories each, with methods evaluated on the same trajectories for paired comparisons.MAL is computed with the paper’s offline estimator, avoiding additional variance from independently sampled evaluation outputs.Source: §5.1 · p. 13

6 Conclusion

The conclusion presents EDR as an exact global objective for parallel and semi-AR speculative drafters, with empirical gains over E2E across nine benchmarks. It identifies fixed block size and occupancy/value computation cost as directions for further work.

6 Conclusion

  • EDR combines a tractable gradient for minimizing expected verification rounds with an exact offline evaluator for paired drafter comparisons.The conclusion reports consistent improvement for DSpark and DFly across nine benchmarks, without changing their architectures or inference procedures, and higher overall MAL than E2E.Source: §6 · pp. 13–14
  • The framework assumes fixed draft block size; occupancy and value-function computations cost more than block-local objectives.The conclusion also identifies adaptive block sizing and a fuller characterization of target-distribution difficulty as future work.Source: §6 · p. 14

A Speculative Decoding Algorithm

B Proofs

The proof conditions speculative decoding on a fixed target-supported output and represents decoding progress with states recording the round start and current position. It derives transitions, establishes the Markov property, and relates state occupancy to the round-count objective.

B.1 Markov Process Formulation in Section 3.1

  • A state records the round-start index and current position, with the round-start index determining the committed prefix and which proposals have already been accepted.Only positions through the block’s bonus-token position are reachable from a given round start.Source: §B.1 · p. 17
  • Rejection sampling preserves the target conditional distribution for each committed token, whether the token is accepted or drawn from the residual distribution.At a bonus-token position, the token is drawn directly from the target distribution.Source: §B.1 · p. 18
  • Conditioned on the fixed output, transition probabilities depend on the past only through the current state, establishing the Markov property.Acceptance continues the current round; rejection or a non-EOS bonus token starts the next round after the processed position.Source: §B.1 · p. 19
  • State occupancy depends only on the committed prefix, and the proof derives forward recursions by summing paths that continue or begin a new round.The final EOS is produced in the current round and does not start another round.Source: §B.1 · pp. 19–20

B.2 Proofs of Proposition 1 and Proposition 2

The proofs establish that decoding-round counts can be characterized by visits to states in the speculative-decoding process. Proposition 1 follows from Proposition 2 by setting k = 1.

B.2 Proofs of Proposition 1 and Proposition 2

  • Proposition 1 is the k = 1 case, since a round starts after position t exactly when state (t, t + 1) is visited.Substituting the occupancy recursion for these states yields the stated result.Source: §B.2 · p. 20
  • A round starting after position n commits at least k tokens exactly when state (n, n + k) is visited.The proof uses the transition rule: reaching that state requires accepting the first k − 1 drafts, while the next processed position commits one more token.Source: §B.2 · p. 20
  • Conditioned on output sequence x, the probability that state (n, n + k) is visited is given by its occupancy.The proof then applies linearity of expectation to the indicators of these state visits.Source: §B.2 · p. 20

B.3 Proof of Theorem 1

The proof handles random output length by conditioning on the generated sequence and its prefixes, then taking expectations and summing over positions.

B.3 Proof of Theorem 1

  • To handle random sequence length L, the proof rewrites Proposition 1 and takes expectation over X ∼ p.It conditions on prefixes to express the relevant occupancies and costs.Source: §B.3 · p. 21
  • The occupancy ωq_n,t(X_<t) is determined by X_<t, and t ≤ L holds exactly when X_t ≠ <EOS>.Conditioning on X_<t lets the proof use the definition of c^q_n,t.Source: §B.3 · p. 21
  • Taking expectations and summing over n and t completes the proof of equation (9).Source: §B.3 · p. 21

B.4 Proof of Theorem 2

The proof differentiates the Bellman equation and propagates its local temporal-difference gradients from the initial state. Each state's contribution is weighted by its probability of being visited.

B.4 Proof of Theorem 2

  • The authors prove equation (13), which relates the initial-state value to per-sequence EDR loss, and then use (13) to obtain (14).Source: §B.4 · p. 21
  • Differentiating the Bellman equation at one state yields a local temporal-difference gradient and recursive terms.These recursive terms are propagated from the initial state using the transition probabilities.Source: §B.4 · p. 21
  • Unrolling the recursion weights each state's local gradient by that state's occupancy.The proof shows that the total coefficient on the local term at (n, t) is exactly ωq_n,t.Source: §B.4 · p. 21

B.5 Proof of Theorem 3

The appendix constructs a first-order-drafter example to contrast block-local optimization with continuation-aware choices, then clarifies terminal-token costs and unbiased fixed-size anchor sampling.

B.5 Proof of Theorem 3

  • For any block-local optimal drafter bq, each round commits an entire subsequence and the next starts at Wi+1; after m subsequences, one additional round generates <EOS>, so E[T bq] = m + 1.Source: §B.5 · p. 22
  • The continuation-aware drafter may sacrifice accepted length to obtain more favorable future round starts.For groups beginning with a type-1 subsequence, a group of ℓ subsequences uses ℓ − 1 rounds, saving one round; rounds starting inside a subsequence advance by B + 1 tokens unless generation ends earlier.Source: §B.5 · pp. 22–23

C.1 Draft EOS Probability Does Not Reduce EDR

  • Matching target EOS mass is not directly rewarded by EDR because producing EOS through acceptance, correction, or the bonus token terminates within the same verification pass.Accordingly, the EDR sum includes only ordinary-token positions.Source: §C.1 · p. 24

C.2 Fixed-Size Anchor Sampling

  • Fixed-size anchor sampling preserves unbiased stochastic gradients while restricting differentiation to sampled anchor contributions.The occupancies and value functions still use full recursions; the experiments set the maximum number of anchors per rollout to 512.Source: §C.2 · pp. 24–25

D Additional Details for Experiments

The appendix points to Table 3 for the training hyperparameters used across all experimental settings and methods.

D Additional Details for Experiments

  • All training settings and methods use the hyperparameters summarized in Table 3.The table is titled as covering hyperparameters used in all experiments.Source: D Additional Details for Experiments · p. 25, Table 3 · p. 25
  • Table 3 presents the training hyperparameters in a hyperparameter-and-value format.The supplied table text identifies the two column headings but gives no values.Source: Table 3 · p. 25
  • The table’s stated scope is the training hyperparameters used in all experiments.Source: Table 3 · p. 25
Loading 2610.10411v1…