TAPS: Task Aware Proposal Distributions for Speculative Sampling

Mohamad Zbib; Mohamad Bazzi; Ammar Mohanna; Hasan Abed Al Kader Hammoud; Bernard Ghanem

cs.CLMar 27, 2026

TAPS: Task Aware Proposal Distributions for Speculative Sampling

Authors: Mohamad Zbib, Mohamad Bazzi, Ammar Mohanna, Hasan Abed Al Kader Hammoud, Bernard Ghanem

Organizations: King Abdullah University of Science and Technology (KAUST) · American University of Beirut (AUB)

Abstract

Speculative decoding accelerates autoregressive generation by letting a lightweight draft model propose future tokens that a larger target model then verifies in parallel. In practice, however, draft models are usually trained on broad generic corpora, which leaves it unclear how much speculative decoding quality depends on the draft training distribution. We study this question with lightweight HASS and EAGLE-2 drafters trained on MathInstruct, ShareGPT, and mixed-data variants, evaluated on MT-Bench, GSM8K, MATH-500, and SVAMP. Measured by acceptance length, task-specific training yields clear specialization: MathInstruct-trained drafts are strongest on reasoning benchmarks, while ShareGPT-trained drafts are strongest on MT-Bench. Mixed-data training improves robustness, but larger mixtures do not dominate across decoding temperatures. We also study how to combine specialized drafters at inference time. Naive checkpoint averaging performs poorly, whereas confidence-based routing improves over single-domain drafts and merged-tree verification yields the highest acceptance length overall for both backbones. Finally, confidence is a more useful routing signal than entropy: rejected tokens tend to have higher entropy, but confidence produces much clearer benchmark-level routing decisions. These results show that speculative decoding quality depends not only on draft architecture, but also on the match between draft training data and downstream workload, and that specialized drafters are better combined at inference time than in weight space.

Figures & tables

Figure 1: Schematic of the speculative decoding pipeline. Given contextual information from the target LLM, the draft model generates latent proposed tokens, which are converted by the LM head and sampling module into multiple candidate future tokens. These candidates are provisional and are later verified by the target model. Importantly, the trainable component in this framework is the draft model, whose role is to efficiently approximate the target model’s next-token behavior while preserving the target model’s final output distribution after verification.

Panel of Figure 2 — Figure 2: Two strategies for combining specialized draft models. Left: checkpoint weight merging in parameter space. Right: confidence-based routing at inference time.

Figure 3: Confidence Routing Between Specialized Trees. The MathInstruct and ShareGPT checkpoints generate separate draft trees from the same prefix, with node labels indicating draft confidence. Confidence routing selects the tree with the higher mean node confidence before verification.

		Temperature 0					Temperature 1
Model Variant	Method	MT-Bench	GSM8K	MATH-500	SVAMP	Average	MT-Bench	GSM8K	MATH-500	SVAMP	Average
RQ1. Task-specific training: single-domain checkpoints
MathInstruct	HASS	2.90	5.02	5.35	3.13	4.10	2.31	4.75	4.63	2.46	3.54
MathInstruct	EAGLE-2	2.54	5.04	5.28	4.81	4.42	2.43	4.71	4.61	4.53	4.07
ShareGPT	HASS	3.98	4.09	3.98	4.44	4.12	3.50	4.03	3.61	3.95	3.77
ShareGPT	EAGLE-2	3.57	3.72	3.81	3.71	3.70	3.38	3.72	3.43	3.65	3.54

Table 1: Main Results by Research Question. Average acceptance length on MT-Bench, GSM8K, MATH-500, and SVAMP for HASS and EAGLE-2 at temperatures

0

and

1

. Rows are grouped by the question they answer: RQ1 tests single-domain specialization, RQ2 mixed-data robustness, and RQ3 composition strategies. Higher is better.

Figure 6: Interpolation Sweep for Checkpoint Averaging. Acceptance length is plotted against the interpolation weight between the MathInstruct and ShareGPT draft checkpoints under a fixed verifier setup. Weight-space averaging is unstable and remains well below the strongest inference-time composition methods.

Panel of Figure 7 — Figure 7: Accepted vs. Rejected Token Entropy. Both panels compare draft entropy at temperature $0$ for EAGLE-2 and HASS on the same benchmark suite and checkpoint families. Entropy is a useful diagnostic of rejection, but Table 2 shows that it is weaker than confidence for routing.

	Confidence Routing			Entropy Routing
Benchmark	MathInstruct	ShareGPT	Total	MathInstruct	ShareGPT	Total
MT-Bench	15 (18.8%)	65 (81.2%)	80	42 (52.5%)	38 (47.5%)	80
GSM8K	1198 (90.8%)	121 (9.2%)	1319	720 (54.6%)	599 (45.4%)	1319
MATH-500	485 (97.0%)	15 (3.0%)	500	312 (62.4%)	188 (37.6%)	500
SVAMP	279 (93.0%)	21 (7.0%)	300	159 (53.0%)	141 (47.0%)	300

Table 2: Routing Decisions by Benchmark. Benchmark-level routing counts for EAGLE-2 under confidence-based and entropy-based selection. Confidence routing separates conversational and mathematical workloads much more clearly than entropy routing.

Figure 8: Acceptance by Speculative Depth. Acceptance rate is shown by draft depth for HASS and EAGLE-2 at temperatures $0$ and $1$ across MT-Bench, GSM8K, MATH-500, and SVAMP. Acceptance declines with depth for all variants, while domain specialization remains visible and often sharpens on reasoning-heavy tasks.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Checkpoint	Benchmark	Draft Accepted	Draft Rejected	$\Delta$ Draft	Verifier Accepted	Verifier Rejected	$\Delta$ Verifier
Averaged	GSM8K	8.0257	9.2397	$+1.2141$	0.1567	0.2698	$+0.1131$
Averaged	MATH-500	8.8405	9.8344	$+0.9938$	0.2013	0.3787	$+0.1774$
Averaged	MT-Bench	8.5128	9.4886	$+0.9758$	0.2427	0.5796	$+0.3368$
Averaged	SVAMP	7.8233	9.0891	$+1.2658$	0.1703	0.2833	$+0.1131$
MathInstruct	GSM8K	0.5284	1.0756	$+0.5473$	0.1500	0.4246	$+0.2746$
MathInstruct	MATH-500	0.4567	0.8555	$+0.3988$	0.1984	0.5067	$+0.3083$

Appendix

Table 3: EAGLE-2 Entropy at Temperature 0. Each row reports accepted-token and rejected-token entropy for one benchmark. Positive

\Delta

means higher entropy for rejected tokens.

Checkpoint	Benchmark	Draft Accepted	Draft Rejected	$\Delta$ Draft	Verifier Accepted	Verifier Rejected	$\Delta$ Verifier
MathInstruct	GSM8K	1.0731	1.4475	$+0.3744$	0.1613	0.5508	$+0.3895$
MathInstruct	MATH-500	2.0806	2.1836	$+0.1030$	0.1928	0.6625	$+0.4698$
MathInstruct	MT-Bench	4.4777	5.2238	$+0.7461$	0.2358	1.3139	$+1.0780$
MathInstruct	SVAMP	1.3927	1.6778	$+0.2851$	0.1765	0.5660	$+0.3896$
ShareGPT	GSM8K	1.0690	1.7551	$+0.6861$	0.1549	0.5937	$+0.4388$
ShareGPT	MATH-500	2.2460	2.3454	$+0.0995$	0.1814	0.7575	$+0.5761$

Appendix

Table 4: HASS Entropy at Temperature 0. Each row reports accepted-token and rejected-token entropy for one benchmark. Positive

\Delta

means higher entropy for rejected tokens.

Depth	Variant	MT-Bench	GSM8K	MATH-500	SVAMP	Avg
1	MathInstruct	72.0%	96.1%	97.9%	94.9%	90.2%
1	ShareGPT	89.9%	93.1%	90.6%	91.9%	91.4%
1	Mixed 35k+35k	88.7%	95.3%	96.4%	94.3%	93.7%
1	Averaged	71.3%	88.3%	89.3%	87.0%	84.0%
1	Mixed 70k+70k	80.0%	80.8%	86.6%	80.7%	82.0%
2	MathInstruct	52.6%	88.8%	92.5%	86.8%	80.2%

Appendix

Table 5: EAGLE-2 Acceptance by Depth at Temperature 0. Higher rows correspond to shallower draft positions in the speculative tree.

Depth	Variant	MT-Bench	GSM8K	MATH-500	SVAMP	Avg
1	MathInstruct	68.5%	94.1%	91.5%	93.1%	86.8%
1	ShareGPT	85.7%	91.7%	83.7%	90.7%	88.0%
1	Mixed 35k+35k	85.1%	93.8%	89.9%	93.3%	90.5%
1	Averaged	67.6%	86.1%	81.2%	84.0%	79.7%
1	Mixed 70k+70k	75.5%	76.9%	75.6%	78.1%	76.5%
2	MathInstruct	50.1%	86.3%	82.0%	84.7%	75.8%

Appendix

Table 6: EAGLE-2 Acceptance by Depth at Temperature 1. Higher rows correspond to shallower draft positions in the speculative tree.

Depth	Variant	MT-Bench	GSM8K	MATH-500	SVAMP	Avg
1	MathInstruct	61.6%	94.6%	95.9%	94.8%	86.7%
1	ShareGPT	84.5%	90.9%	88.0%	89.4%	88.2%
1	Mixed 35k+35k	84.9%	95.2%	96.1%	94.7%	92.7%
1	Averaged	68.1%	87.5%	90.9%	87.4%	83.5%
1	Mixed 70k+70k	82.3%	85.1%	90.6%	86.2%	86.0%
2	MathInstruct	50.9%	84.8%	88.9%	82.3%	76.7%

Appendix

Table 7: HASS Acceptance by Depth at Temperature 0. Higher rows correspond to shallower draft positions in the speculative tree.

Speculative Decoding

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

See all 309 papers

Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds
Oct 7, 2026Yunxiao Zhao, Changxiao CaiSpeculative Decoding Parallel Speculative Decoding
Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, MI, USA.
Draft-OPD: On-Policy Distillation for Speculative Draft Models
May 28, 2026Haodi Lei, Yafu Li, Haoran Zhang +8Speculative Decoding LLM Inference Acceleration
Shanghai Jiao Tong University · Shanghai AI Laboratory · The Chinese University of Hong Kong +3
D^2SD: Accelerating Speculative Decoding with Dual Diffusion Draft Models
Jun 3, 2026Liyuan Zhang, Jiarui Zhang, Jinwei Yao +6Speculative Decoding for Diffusion Language Models Speculative Decoding
Peking University · Tsinghua University · HKUST +2

Oct 7, 2026J/K move · Enter open · S save

Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds

Yunxiao Zhao, Changxiao Cai

Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, MI, USA.

Speculative Decoding Parallel Speculative Decoding LLM Inference Acceleration

Open paper AskPDF arXiv

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.

Abstract

Speculative Decoding

Explore similar work

Training Parallel Speculative Draft Models by Directly Minimizing Expected Decoding Rounds

Draft-OPD: On-Policy Distillation for Speculative Draft Models

D^2SD: Accelerating Speculative Decoding with Dual Diffusion Draft Models