Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds

What it claims

For parallel and semi-autoregressive (semi-AR) drafters, the draft for a given position depends on where the current round started. Training them to match the target locally is therefore not the same as minimizing target-model forward passes. The paper models speculative decoding, conditioned on the output sequence, as a Markov reward process. From that it derives EDR, a loss that exactly equals expected decoding rounds. It also proves that a drafter maximizing expected accepted length per round can be strictly suboptimal for EDR. That result is the argument against block-local objectives such as Li et al. 2026c and the hand-weighted divergences of Chen et al. 2026a and Wu et al. 2026.

How it works

• State (n, t): the current round began after committed prefix x_1:n, and position t is being verified. Decoding starts at (0,1).
• Acceptance probability given the realized token: A = min(1, q_{n,t}(x_t | x_1:n; x_{n+1:t-1}) / p_t(x_t | x_<t)). Rejection probability is R = 1 − A. Past the block (t > n+B), A = 0 and R = 1.
• Transitions: an accept moves to (n, t+1). A reject starts a new round at (t, t+1) and costs 1.
• Occupancies: these are propagated forward. The expected number of rounds equals 1 plus the sum of occupancy × rejection probability (Proposition 1).
• Local rejection cost: the realized-token rejection is replaced by its expectation under the target with EOS excluded. That is the total-variation (TV) distance Σ_v [p_t(v) − q_{n,t}(v)]_+ divided by 1 − p_t(EOS). EOS incurs no cost, because producing it never starts another round.
• EDR: the occupancy-weighted sum of these costs, averaged over target rollouts. It has no auxiliary hyperparameters.
• Training: a temporal-difference form of the gradient allows unbiased stochastic optimization from target-model rollouts.
• Offline evaluation: the same recursion computes exact expected round counts on shared rollouts, so two drafters can be compared without running speculative decoding.

flowchart LR
  S["state (n,t): verify draft t"] -->|accept, prob A| N["(n,t+1) same round"]
  S -->|reject, prob R, cost 1| R["(t,t+1) new round"]
  N --> S
  R --> S
  S -->|t > n+B| R

Evidence

The experiments finetune DSpark and DFly with EDR for a single epoch, keeping architectures and inference procedures unchanged. They report higher mean accepted length (MAL) than both the original drafters and the competing training objectives, across nine math, code and chat benchmarks. MAL is defined as E[L+1]/E[T_q], so raising it is equivalent to cutting expected rounds for a fixed target.

How to use it

Code: https://github.com/y-x-zhao/AngelSpec-EDR. EDR is a finetuning loss for an existing drafter, not a new architecture. It needs target-model rollouts and access to the target's full next-token distributions p_t. The offline evaluator is a cheap way to A/B drafters on fixed rollouts before deploying.

Limits

• The fetched text stops in Section 3.2. The benchmark names, MAL numbers, baselines' scores, compute and target-model sizes are not visible here; check Section 5 and Appendix D before relying on the gains.
• The gains are measured as MAL, i.e. fewer rounds. Wall-clock speedup is not established in the visible text.
• The theory assumes output length has a universal upper bound.
• The method is validated only on DSpark and DFly.
• Appendix C.2 mentions "fixed-size anchor sampling", which suggests the practical estimator subsamples round starts. Its details are not in the fetched text.