In speculative decoding, a lightweight draft model proposes tokens that the target model verifies in parallel, speeding up generation without changing the target distribution. D-PACE (Wu et al., 2026) trains such drafters with a loss that weights each position by its contribution to the expected accepted length. Both its weights and its cross-entropy (CE) loss use only the draft's probability of the recorded target token, treating the target as a point mass. Under rejection sampling, however, a sampled draft token is accepted with probability equal to the overlap of the draft and target distributions, one minus their total variation (TV). We derive D-PAL (Dynamic Position-wise Acceptance Loss) from the exact expected accepted length, using this overlap as the acceptance signal. This yields overlap-based position weights with log acceptance (LA) as the per-position loss, and the objective reduces to D-PACE for a point-mass target. Since the weights are detached, any per-position loss can replace LA in practice. We show that forward KL divergence and Rényi divergence of order 1/2 match or exceed LA, and that Rényi-1/2 provably bounds the true acceptance rate from above and below. Replacing CE with Rényi-1/2 improves acceptance length under both D-PACE's weights and ours. On six benchmarks with a Qwen3-4B target, D-PAL raises mean acceptance length over D-PACE from 4.50 to 4.57 at T=0 and from 4.20 to 4.28 at T=1, and consistently achieves higher throughput at both temperatures.
Figures & tables
Figure 1: D-PAL uses the full target distribution for position weighting and, with KL or Rényi, for local fitting. Top: D-PACE weights by confidence and fits one token with CE; faded bars are target mass CE ignores. Bottom: D-PAL weights by overlap and, with a distributional loss such as KL or Rényi, fits qt to all of pt . Each weighted term wtℓt is the product of the acceptance through t (left), the local loss at t (center), and the continuation value after t (right).
Figure 2: Near q=p , LA’s slope stays near ±1 and flips sign, while KL and Rényi- 1/2 flatten out. Binary example with target p=(0.5,0.5) and draft q=(x,1−x) ; slopes are with respect to x , and all three losses are zero at x=0.5 .
for each nonempty valid block b :
Select proposals with mb,t=1 in order; index them by t=1,…,Db .
αt←∑vmin(pt,v,qt,v) ; st←ρ+(1−ρ)αt for all t
R0←1 ; wDb+1←0
for t=1,…,Db : Rt←stRt−1▹ Smoothed prefix product
for t=Db,…,1 : wt←wt+1+Rt▹ Position weights
Detach w1:Db .
Table 3
Math
Code
Chat
GSM8K
MATH-500
HumanEval
MBPP
MT-Bench
Alpaca
Mean
Method
SR
τ
SR
τ
SR
τ
SR
τ
SR
τ
SR
τ
SR
τ
Baselines
Temperature T=0
DFlash
3.20
4.50
3.94
5.46
3.33
4.67
3.06
4.28
1.97
3.15
1.77
2.71
2.88
4.13
D-PACE
3.50
4.96
4.35
6.06
3.60
5.09
3.34
4.65
2.05
3.36
1.83
2.87
3.11
4.50
D-PAL weights
Table 1: D-PAL[Rényi] , D-PAL[LA] , and D-PAL[KL] beat D-PACE; D-PAL[TV] and D-PAL[CE+TV] fall short. Qwen3-4B target, three-layer DFlash drafts, B=16 , 100K training examples, ρ=0.3 for D-PAL and 0.5 for D-PACE. SR is speedup over autoregressive decoding, τ is acceptance length, and bold marks the best value.
K=1
K=2
K=4
K=8
K=16
Method
TPS
τ
TPS
τ
TPS
τ
TPS
τ
TPS
τ
Temperature T=0
DFlash
310.24
5.46
423.17
5.48
607.93
5.46
990.60
5.49
1628.15
5.49
D-PACE
320.13
6.06
466.27
6.06
669.64
6.04
1120.19
6.06
1714.12
6.06
D-PAL[Rényi]
339.73
6.10
477.46
6.13
686.75
6.10
1094.14
6.13
1797.82
6.11
Temperature T=1
Table 2: D-PAL[Rényi] has the highest τ at every batch size and the highest throughput in 9 of 10 settings. 128 MATH-500 prompts on Qwen3-4B with the 100K checkpoints; K is the inference batch size (number of requests) and TPS is output tokens per second.
Table 6Table 7Table 8
Acceptance length τ
Method
GSM8K
MATH-500
HumanEval
MBPP
MT-Bench
Alpaca
Mean
Temperature T=0
DFlash
3.64
3.53
3.92
3.55
3.12
3.04
3.47
D-PACE
3.94
3.82
4.14
3.96
3.28
3.17
3.72
D-PAL[LA]
3.96
3.91
4.21
3.85
3.47
3.22
3.77
D-PAL[KL]
3.98
3.90
4.20
3.96
3.46
3.23
3.79
Table 9: With thinking enabled, D-PAL[Rényi] leads mean τ at T=0 and D-PAL[KL] at T=1 . DFlash is the reference checkpoint (protocol in Appendix A.1 ).
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Training
Optimizer / epochs
AdamW / 6
Learning rate / schedule
6×10−4 ; cosine decay, 4% warm-up
Global batch / seed
8 (Tables 1 and 2 ), 4 (other runs); seed 42
Sequence / anchor limits
3072 tokens / 512 anchors
Default block size
Qwen: B=16 ; Llama: B=10
Loss temperature
1 (target and draft)
Appendix
Table 10: Training and evaluation settings. Each block includes one anchor token excluded from the loss. Both draft depths use the same three target feature layers.
Objective
Local loss ℓt
Cross-entropy
ℓCE,t=−logqt(yt⋆)
D-PAL[LA]
−log[1−TV(pt,qt)]
D-PAL[KL]
KL(pt∥qt)
D-PAL[Rényi]
D1/2(pt∥qt)
D-PAL[TV]
TV(pt,qt)
D-PAL[CE+TV]
0.1ℓCE,t+0.9TV(pt,qt)
Appendix
Table 11: Local losses compared in this paper. D-PAL[KL+TV] computes λt=exp(−3αt) for each block and position and detaches it during optimization.
Figure 3: A CE+TV descent step can raise TV. Negative weighted gradients of TV ( a=0.9 ) and CE ( b=0.1 ) at one parameter point of the three-token example, on one shared scale and rotated 25∘ for display. The two gradients have a negative inner product, and the mixture step −gΣ raises TV to first order ( gT⊤gΣ<0 ).
Speculative decoding accelerates LLM inference by having a small drafter propose tokens that a larger target model verifies in parallel. Recent diffusion-based parallel drafters such as DFlash predict the full B-token block in one forward pass, enabling deeper drafters and longer accepted blocks. However, existing multi-token drafter objectives often use fixed position-dependent weighting schedules, such as head-dependent weights or block-position decays, which do not adapt as the positions limiting acceptance change during training. To address this, we derive per-position training weights from a differentiable surrogate of expected accepted draft length, matching the weight of each position to its log-probability gradient contribution. The resulting loss, D-PACE (Dynamic Position-Aware Cross-Entropy), shifts training signal toward positions that currently limit acceptance as the drafter improves. Across six benchmarks, two Qwen3-4B draft depths, two decoding temperatures, and two additional target models, D-PACE consistently improves both wall-clock speedup and average emitted length, with 2.3% measured training-time overhead and no changes to the drafter architecture or inference procedure.
Speculative decoding accelerates large language model (LLM) inference by using a lightweight draft model to generate multiple candidate tokens that are verified by the target model in a single forward pass. Its speedup is largely determined by the acceptance length, yet existing draft-model training methods mainly optimize cross-entropy or Kullback-Leibler (KL) divergence as proxies. These objectives encourage distribution matching but do not directly optimize acceptance length, and the acceptance mechanism also differs between greedy and sampling-based decoding. In this work, we propose acceptance-length-aware training losses that directly optimize the expected number of accepted tokens within a speculative window. For greedy verification, we derive an expected accepted length (EAL) loss that explicitly maximizes expected acceptance length. For sampling-based decoding, we introduce a window total variation (WTV) loss that optimizes the overlap between temperature-scaled draft and target distributions while accounting for sequential acceptance dependencies. Both objectives can be further combined with a group-relative reinforcement learning stage (GRPO) using simulated acceptance length as the reward. Experiments across different target and draft models, tasks, and decoding settings show that our losses consistently improve acceptance length over KL-based training. WTV provides particularly strong gains under sampling-based decoding, while EAL better matches greedy verification. These results show that directly optimizing the acceptance objective, with losses tailored to the decoding mode, is more effective than conventional distribution-matching objectives.
Tianhua Xia, Mugilan Ganesan, Yifei Feng +3
New York University · University of Waterloo · Independent Researcher
Speculative decoding accelerates language model inference by using a fast drafter to propose candidate tokens that are then verified by a larger target model. Existing theory largely studies the stochastic, distribution-preserving setting, where the goal is to exactly sample from the target distribution. In contrast, many practical systems use greedy decoding, relaxed acceptance rules, or tree-based candidate sets, where success is governed by local ranking and threshold events rather than exact distributional equality. We develop a theory for these regimes. We identify that many common acceptance criteria have rejection regions that can be characterized as lower level sets of the target distribution. For these, we characterize the exact KL divergence required for rejection yielding exact certificates and sharp margin-based bounds for strict greedy decoding, additive and multiplicative relaxed acceptance, top-(m) relaxed criteria, and entropy-thresholded acceptance. We then extend the framework to greedy tree decoding, deriving exact and margin-only certificates for when the target greedy token remains covered by the drafter's top-(m) candidates. Finally, we evaluate the resulting certificates on Qwen3 models, showing that relaxed and tree-based criteria substantially enlarge the region of certified acceptance, especially on decoding steps with low target model distribution margin. These results complement existing distribution-preserving analyses of speculative decoding by characterizing the deterministic local acceptance events common in practical inference systems.