Source-linked AI summary
TAPS: Task Aware Proposal Distributions for Speculative Sampling
Mohamad Zbib, Mohamad Bazzi, Ammar Mohanna, Hasan Abed Al Kader Hammoud, Bernard Ghanem
TL;DR
Speculative decoding depends on how well a draft model’s training distribution matches the target workload, but evidence for this dependence and for combining specialized drafters remains limited. This paper evaluates task-specific and mixed-data HASS and EAGLE-2 drafters, then compares weight averaging, confidence routing, and merged-tree verification across four benchmarks. Task-specific training produces clear domain specialization, while inference-time composition outperforms weight-space averaging and merged-tree verification achieves the highest acceptance length overall.
Problem
Evidence is limited on how draft training distributions affect speculative decoding and how multiple specialized drafters should be combined.
Method
The paper evaluates HASS and EAGLE-2 drafters trained on MathInstruct, ShareGPT, and mixed data, comparing checkpoint averaging, confidence routing, and merged-tree verification under a fixed verifier.
Results
Task-specific training produces clear domain specialization, while inference-time composition is stronger than weight-space averaging and merged-tree verification achieves the highest acceptance length overall.
Takeaways & Limitations
Speculative decoding quality depends on both draft architecture and training-distribution alignment, and separate specialized drafters are better composed at inference time than collapsed through naive averaging.
Takeaways & Limitations
The evidence covers one target model, two source domains, two backbones, and four benchmarks, and does not establish end-to-end deployment trade-offs.
Abstract
from arXiv · showhide
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.
Introduction
This paper asks whether draft training distributions affect speculative decoding and how specialized drafters should be combined. Across both backbones, task-specific training specializes performance, while inference-time composition outperforms naive checkpoint averaging.
- Autoregressive decoding is a major inference bottleneck because each token depends on the full previously generated prefix.
- Speculative decoding uses a lightweight draft model to propose future tokens that a larger target model verifies in parallel without changing the target output distribution.
- The paper investigates whether task-specific or mixed-data training improves proposal quality and whether specialized drafters should be averaged, routed, or jointly verified.
- MathInstruct-trained drafts are strongest on GSM8K, MATH-500, and SVAMP, whereas ShareGPT-trained drafts are strongest on MT-Bench.
- Mixed-data training improves robustness, but larger mixtures are not uniformly better across temperatures.
- Confidence routing improves over single-domain baselines, while merged-tree verification achieves the highest acceptance length overall.
Preliminaries
The paper fixes the verifier and lossless acceptance rule while varying the EAGLE-2 or HASS drafter and its training or composition strategy. These backbones generate proposals that are verified against the target model while preserving its output distribution.
- The study reviews lossless verification and the EAGLE-2 and HASS drafting backbones, keeping verification fixed throughout experiments.
- Speculative decoding has a draft model propose K future tokens and a target model verify them in parallel, reducing expensive target-model calls.
- Rejection-sampling correction preserves the target model’s output distribution while allowing multiple proposals to be checked in one verifier call.
- EAGLE-2: EAGLE-2 predicts future features, maps them to token probabilities through the target LM head, and expands high-confidence frontier nodes in a dynamic tree.
- HASS: HASS reduces objective and context mismatch through top-K distillation and training on imperfect draft features.
Experimental Setup and Composition Strategies
The experiments compare fixed-verifier drafting across domains, mixture sizes, backbones, and composition strategies. Specialized checkpoints can be averaged, routed by confidence, or merged into a jointly verified tree, with merged verification increasing proposal coverage but requiring separate systems analysis.
- Experimental Setup: Experiments evaluate MT-Bench, GSM8K, MATH-500, and SVAMP at temperatures 0 and 1 using acceptance length as the primary metric.
- Experimental Setup: The study varies speculative backbone, training domain, mixed-data supervision, and test-time composition strategy.
- Inference-Time Composition: Confidence routing generates separate specialist trees and passes only the tree with the higher mean draft confidence to the verifier.
- Checkpoint Averaging: Checkpoint averaging combines MathInstruct and ShareGPT parameters into one merged draft controlled by λ, with λ = 0.5 used in the main table.
- Inference-Time Composition: Merged-tree verification packs both specialist subtrees under a shared root and preserves each subtree’s ancestry through attention masks and depth-based position ids.
- Inference-Time Composition: Merged trees increase proposal diversity, but the paper does not claim end-to-end latency improvement without separate systems analysis.
Experiments
Across both HASS and EAGLE-2, acceptance length depends on task-matched training, decoding temperature, and how specialized drafters are combined. Inference-time composition and confidence routing outperform weight-space averaging, while depth exposes a coverage-to-specialization shift.
- Acceptance length measures the average number of draft tokens accepted per verifier call, with higher values indicating better drafter–verifier alignment.
- RQ1: Single-domain specialization: Task-matched training produces clear specialization: ShareGPT leads on MT-Bench, while MathInstruct leads on GSM8K, MATH-500, and SVAMP.Under HASS at temperature 0, the corresponding MT-Bench values are 3.98 versus 2.90, while GSM8K and MATH-500 are 5.02 versus 4.09 and 5.35 versus 3.98.
- RQ2: Mixed-data robustness: Mixed-data training improves robustness, but larger mixtures do not uniformly improve generalization across decoding temperatures.For HASS, Mixed 70k+70k reaches 5.18 at temperature 0 but 3.69 at temperature 1, below Mixed 35k+35k at 4.29.
- RQ3: Specialist composition: Merged-tree verification reaches the highest average acceptance length, while confidence routing improves over single-domain drafts and checkpoint averaging remains weakest.At temperature 0, merged-tree verification reaches 5.11 for HASS and 5.02 for EAGLE-2; confidence routing reaches 4.80 and 4.63, versus 2.34–2.62 for averaged checkpoints.
- RQ4: Routing and diagnostics: Confidence separates benchmark-level routing decisions more clearly than entropy, although rejected tokens consistently have higher entropy than accepted tokens.Confidence routing selects MathInstruct for 90.8% of GSM8K, 97.0% of MATH-500, and 93.0% of SVAMP, and ShareGPT for 81.2% of MT-Bench.
- RQ4: Routing and diagnostics: Acceptance declines with speculative depth, while task specialization often becomes more pronounced on reasoning-heavy tasks.Shallow depths often favor mixed-data coverage, whereas deeper accepted paths increasingly favor the task-matched specialist.
Discussion and Limitations
The paper argues that speculative decoding quality depends on matching draft training distributions to workloads and that specialized drafters are better composed at inference time than merged in weight space. However, the evidence is limited in scope and does not establish end-to-end deployment trade-offs.
- Discussion: Speculative decoding depends on the relationship between draft training distribution and target workload, not only on the drafting backbone.A mismatched draft is predictably weaker on specific task families.
- Discussion: Inference-time composition is substantially stronger than collapsing specialized drafters into one averaged checkpoint.Naive weight-space interpolation fails to preserve the relevant behavior.
- Trade-offs: 0.32× and 0.35× under EAGLE-2 and 0.40× and 0.47× under HASS are the confidence-routing speedup reductions at temperatures 0 and 1.Merged-tree verification incurs larger drops of 0.59× and 0.62× under EAGLE-2 and 0.72× and 0.78× under HASS.
- Limitations: The evidence covers one target model, two source domains, two speculative backbones, and four benchmarks, with acceptance length as the primary metric.The routing policy is simple and confidence-based rather than learned or cost-aware.
Conclusion
The paper asks whether task-trained drafters improve speculative decoding and how specialized drafters should be combined. It finds clear domain specialization, stronger inference-time composition than weight averaging, and the highest acceptance length from merged-tree verification.
- Conclusion: Task-specific training produces clear domain specialization in speculative decoding.This answers the paper’s first research question affirmatively.
- Conclusion: Inference-time composition is substantially stronger than weight-space averaging for combining specialized drafters.Confidence routing improves over single-domain baselines, while merged-tree verification achieves the highest acceptance length overall across HASS and EAGLE-2.
- Conclusion: Proposal quality should be studied as a function of both draft architecture and draft training distribution.The conclusion extends the paper’s focus beyond architecture alone.
A. Appendix
The appendix collects supporting materials for the paper’s main claims and is intended to make the evidence easier to audit.
- A. Appendix: The appendix contains entropy tables, the tree-merging utility, and depth-wise acceptance tables.These materials complement Figure 8 and support the main-text claims.
A.1. EAGLE-2 Entropy (Temperature 0)
Table 3 reports EAGLE-2 entropy at temperature 0 for accepted and rejected tokens across benchmarks.
- A.1. EAGLE-2 Entropy (Temperature 0): Each row reports accepted-token and rejected-token entropy for one benchmark.A positive ∆ indicates higher entropy for rejected tokens.
A.2. HASS Entropy (Temperature 0)
Table 4 compares accepted-token and rejected-token entropy for each benchmark at temperature 0. Positive Δ indicates higher entropy for rejected tokens.
- Each row reports accepted-token and rejected-token entropy for one benchmark.
- The table is evaluated at temperature 0.
- Positive Δ means rejected tokens have higher entropy than accepted tokens.
A.3. Tree Merge Utility
The tree-merge utility combines two draft trees into one packed representation. It preserves each subtree’s structure while adjusting indices, masks, and positions for joint verification.
- The merge routine receives two drafts, retrieval-index arrays, tree masks, and position arrays.
- The merged tree size is computed as N = n1 + n2 + 1, sharing one root across both input trees.
- The merged attention mask keeps each subtree’s ancestry while allowing both subtrees to attend to the shared root.
- Retrieval indices for the second subtree are offset by n1 before the two index arrays are concatenated.
- The routine returns merged draft tokens, retrieval indices, attention mask, and position identifiers.
A.4. Correctness of routing and merged-tree verification
The paper establishes that valid speculative trees, draft-side routing, and merged-tree verification preserve the target model’s output distribution. These mechanisms can therefore alter proposal quality and acceptance behavior without changing the verifier’s sampling law.
- Assumptions and losslessness: A valid packed tree must reproduce the target-side conditional distributions along every node’s path.
- Mixtures and routing: Random mixtures over valid trees remain lossless under the standard speculative-decoding guarantee.
- Mixtures and routing: Any routing rule based on pre-verification draft-side quantities preserves the target-model output distribution.
- Merged-tree verification: Merged-tree verification shares the root, concatenates non-root nodes, preserves subtree ancestry, masks cross-subtree attention, and retains source depth positions.
- Merged-tree verification: Merged verification gives each subtree the same target-side conditional distributions as standalone verification and therefore preserves the verifier’s sampling law.