Beyond the Point Mass: Acceptance-Derived Weights and Distributional Losses for Speculative Decoding

Tianyu Wu; Zhuohan Wang; Jonathan Geuter; Yu Yao; Weifan Jiang; Chengxi Zhang; Ju Li; Minlan Yu

cs.LGOct 20, 2026

Beyond the Point Mass: Acceptance-Derived Weights and Distributional Losses for Speculative Decoding

Authors: Tianyu Wu, Zhuohan Wang, Jonathan Geuter, Yu Yao, Weifan Jiang, Chengxi Zhang, Ju Li, Minlan Yu

Organizations: Harvard University · Massachusetts Institute of Technology

Abstract

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

Panel of Figure 1 — 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 $q_{t}$ to all of $p_{t}$ . Each weighted term $w_{t}\ell_{t}$ is the product of the acceptance through $t$ (left), the local loss at $t$ (center), and the continuation value after $t$ (right).

Panel of Figure 2 — Figure 2: Near $q=p$ , LA’s slope stays near $\pm 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 $m_{b,t}=1$ in order; index them by $t=1,\ldots,D_{b}$ .
	$\widehat{\alpha}_{t}\leftarrow\sum_{v}\min(p_{t,v},q_{t,v})$ ; $s_{t}\leftarrow\rho+(1-\rho)\widehat{\alpha}_{t}$ for all $t$
	$\widetilde{R}_{0}\leftarrow 1$ ; $w_{D_{b}+1}\leftarrow 0$
	for $t=1,\ldots,D_{b}$ : $\widetilde{R}_{t}\leftarrow s_{t}\widetilde{R}_{t-1}$ $\triangleright$ Smoothed prefix product
	for $t=D_{b},\ldots,1$ : $w_{t}\leftarrow w_{t+1}+\widetilde{R}_{t}$ $\triangleright$ Position weights
	Detach $w_{1:D_{b}}$ .

Table 3

D-PAL weights
	Math				Code				Chat
	GSM8K		MATH-500		HumanEval		MBPP		MT-Bench		Alpaca		Mean
Method	SR	$\tau$	SR	$\tau$	SR	$\tau$	SR	$\tau$	SR	$\tau$	SR	$\tau$	SR	$\tau$
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

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,

\rho=0.3

for D-PAL and

0.5

for D-PACE. SR is speedup over autoregressive decoding,

\tau

is acceptance length, and bold marks the best value.

Temperature $T=0$
	$K=1$		$K=2$		$K=4$		$K=8$		$K=16$
Method	TPS	$\tau$	TPS	$\tau$	TPS	$\tau$	TPS	$\tau$	TPS	$\tau$
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

\tau

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 6

Table 7

Table 8

Temperature $T=0$
	Acceptance length $\tau$
Method	GSM8K	MATH-500	HumanEval	MBPP	MT-Bench	Alpaca	Mean
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

\tau

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\times 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 $\ell_{t}$
Cross-entropy	$\ell_{\mathrm{CE},t}=-\log q_{t}(y_{t}^{\star})$
D-PAL[LA]	$-\log[1-\operatorname{TV}(p_{t},q_{t})]$
D-PAL[KL]	$\operatorname{KL}(p_{t}\\|q_{t})$
D-PAL[Rényi]	$D_{1/2}(p_{t}\\|q_{t})$
D-PAL[TV]	$\operatorname{TV}(p_{t},q_{t})$
D-PAL[CE+TV]	$0.1\ell_{\mathrm{CE},t}+0.9\operatorname{TV}(p_{t},q_{t})$

Appendix

Table 11: Local losses compared in this paper. D-PAL[KL+TV] computes

\lambda_{t}=\exp(-3\widehat{\alpha}_{t})

for each block and position and detaches it during optimization.

Panel of Figure 3 — 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^{\circ}$ for display. The two gradients have a negative inner product, and the mixture step $-g_{\Sigma}$ raises TV to first order ( $g_{T}^{\top}g_{\Sigma}<0$ ).

Speculative Decoding

Speculative decoding accelerates generation by drafting multiple tokens and verifying them in parallel with a larger target model.

See all 308 papers

D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting
May 12, 2026Tianyu Wu, Yu Yao, Zhenting Qi +7Speculative Decoding LLM Inference Acceleration
1Harvard · 2MIT
Acceptance-Aware Draft Model Training for Speculative Decoding
Sep 21, 2026Tianhua Xia, Mugilan Ganesan, Yifei Feng +3Speculative Decoding LLM Inference Acceleration
New York University · University of Waterloo · Independent Researcher
When Is a Draft Accepted? A Theory of Acceptance in Speculative Decoding
Jun 29, 2026Aaryam SharmaSpeculative Decoding LLM Inference Acceleration
Independent Researcher Waterloo, Ontario, Canada

May 12, 2026J/K move · Enter open · S save

D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting

Tianyu Wu, Yu Yao, Zhenting Qi, Han Zheng+6

1Harvard · 2MIT

Speculative Decoding LLM Inference Acceleration Parallel Decoding Diffusion Language Model Decoding

Open paper AskPDF arXiv

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.

Abstract

Speculative Decoding

Explore similar work

D-PACE: Dynamic Position-Aware Cross-Entropy for Parallel Speculative Drafting

Acceptance-Aware Draft Model Training for Speculative Decoding

When Is a Draft Accepted? A Theory of Acceptance in Speculative Decoding