TAPS: Task Aware Proposal Distributions for Speculative Sampling

This paper demonstrates that training draft models on task-specific data significantly improves speculative decoding performance for corresponding workloads and that combining these specialized drafters at inference time via confidence-based routing and merged-tree verification outperforms naive weight averaging.

The Gist

Imagine you are a brilliant but very slow architect (the Target Model) trying to build a skyscraper, one brick at a time. Every time you place a brick, you have to stop, think deeply, check the blueprints, and make sure it's perfect before moving to the next one. This process is accurate but incredibly slow.

To speed things up, you hire a fast, energetic intern (the Draft Model). The intern's job is to quickly guess the next 5 or 10 bricks you might want to place. You then quickly glance at the intern's guesses. If they look right, you accept them all at once and move on. If they look wrong, you discard them and do the work yourself.

This is Speculative Decoding. The goal is to get the intern to guess correctly as often as possible.

The Big Problem: The "Wrong Intern"

The paper TAPS asks a simple but crucial question: Does it matter what the intern studied in school?

In the past, researchers just hired interns who read a little bit of everything (news, chat logs, random facts). They assumed a "generalist" intern would be good at guessing anything.

The authors of this paper say: "No! That's like hiring a generalist to build a rocket ship."

They tested two types of interns:

1. The Math Whiz: Trained only on math problems and logic puzzles.
2. The Chatterbox: Trained only on casual conversations and social media chats.

The Results: Specialization Wins

When they put these interns to work:

• The Math Whiz was amazing at building the "rocket ship" (solving math problems) but terrible at writing a casual email.
• The Chatterbox was great at writing emails but failed miserably at the math problems.
• The Generalist (trained on a mix of both) was okay at everything, but not great at anything.

The Lesson: If you know you are going to be doing math, hire the Math Whiz. If you are writing a story, hire the Chatterbox. The quality of the "guessing" depends entirely on how well the intern's training matches the job you are giving them.

The "Merging" Mistake

The researchers then asked: "What if we have both a Math Whiz and a Chatterbox available? Can we just mix their brains together to make one super-intern?"

They tried Weight Averaging (literally mixing the intern's brain weights together).

• Result: Disaster. The new "Super-Intern" was confused. They tried to be a little bit of both, and ended up being mediocre at both. It was like mixing oil and water; you just get a messy sludge.

The Smart Solution: The "Smart Manager"

Instead of mixing their brains, the researchers tried a Smart Manager approach (Inference-Time Composition).

Imagine you have both interns sitting at the desk. Before you ask them to guess the next bricks, the Manager looks at the current task:

• "Oh, we are doing a math problem? Math Whiz, you take the lead!"
• "Oh, we are writing a poem? Chatterbox, you take the lead!"

They found two ways to do this:

1. Confidence Routing: The Manager asks, "Who feels most confident about this guess?" and picks that intern. This worked very well.
2. Merged-Tree Verification (The Best Way): The Manager lets both interns shout out their guesses at the same time, but keeps their lists separate. The Architect (Target Model) then checks all the guesses from both lists at once.
   • Analogy: It's like having two different teams of detectives solving a mystery. Instead of forcing them to agree on one theory, you let them both present their theories, and you pick the one that fits the evidence best. This gave the best results of all.

Why Entropy (Confusion) Didn't Work

The researchers also tried to use "Confusion" (Entropy) as a signal. They thought, "If the intern is confused, maybe we should switch to the other intern."

• Result: It didn't work well. The interns were often confused even when they were right.
• Better Signal: Confidence. If an intern says, "I am 99% sure this is the right brick," that is a much better signal to trust them than asking, "Are you confused?"

The "Depth" Surprise

Finally, they noticed something interesting about how deep the guesses go.

• Shallow Guesses (The first few bricks): It helps to have a broad mix of ideas.
• Deep Guesses (The 10th or 20th brick): You really need the specialist. If you are deep into a complex math proof, the Chatterbox will get lost. The Math Whiz is the only one who can keep the chain of logic going.

Summary in One Sentence

To make AI faster, don't just train a "jack-of-all-trades" assistant; instead, train specialized experts for specific jobs, and use a smart system to pick the right expert (or let them both work) at the exact moment you need them.

Technical Summary

1. Problem Statement

Speculative decoding accelerates Large Language Model (LLM) inference by using a lightweight "draft" model to propose future tokens, which a larger "target" model verifies in parallel. While prior research has focused on improving draft architectures (e.g., EAGLE, HASS) or verification mechanisms, a critical gap remains: the impact of the draft model's training distribution.

Most draft models are trained on broad, generic corpora (e.g., ShareGPT). The paper investigates whether this generic training is optimal or if task-specific training yields better performance. Furthermore, it addresses the practical challenge of combining multiple specialized drafters: Is it better to merge them via weight averaging during training, or to compose them dynamically at inference time?

2. Methodology

The authors conducted a controlled study using two state-of-the-art speculative decoding backbones:

• EAGLE-2: A feature-level drafter that predicts future hidden states and uses a dynamic draft tree.
• HASS (Harmonized Speculative Sampling): A method that reduces objective and context mismatch between training and inference.

Experimental Setup:

• Target Model: Meta-Llama-3-8B-Instruct (fixed verifier).
• Draft Models: Lightweight LLaMA-style decoders (~0.8B parameters).
• Training Distributions:
  • Single-Domain: MathInstruct (mathematical reasoning) vs. ShareGPT (conversational).
  • Mixed-Domain: Variants combining 35k+35k and 70k+70k examples from both domains.
• Evaluation Benchmarks: MT-Bench (conversational), GSM8K, MATH-500, and SVAMP (reasoning).
• Metrics: Primary metric is Acceptance Length (average number of consecutive tokens accepted by the target model).

Composition Strategies Investigated:

1. Checkpoint Averaging: Merging weights of specialized models ($\theta_{merge} = \lambda\theta_{math} + (1-\lambda)\theta_{chat}$).
2. Confidence-Based Routing: At inference, generating trees from both specialists and selecting the one with the highest mean confidence.
3. Merged-Tree Verification: Packing draft trees from both specialists under a shared root to allow the verifier to evaluate candidates from both domains simultaneously in a single pass.

3. Key Contributions & Findings

The paper answers five specific research questions (RQs) with consistent results across both EAGLE-2 and HASS backbones:

RQ1: Does task-specific training improve matched-domain acceptance?

• Finding: Yes. Specialization is clear.
  • MathInstruct-trained drafts significantly outperform generic drafts on reasoning benchmarks (GSM8K, MATH-500, SVAMP).
  • ShareGPT-trained drafts significantly outperform on conversational benchmarks (MT-Bench).
• Implication: Draft quality is not just an architectural property; it depends heavily on the alignment between the draft training data and the downstream workload.

RQ2: Can mixed-data training recover cross-domain robustness?

• Finding: Mixed-data training improves robustness but does not eliminate specialization.
  • Larger mixtures (70k+70k) performed best at temperature 0 but degraded significantly at temperature 1 compared to smaller mixtures.
  • No single mixed-data checkpoint uniformly dominated across all temperatures and tasks.

RQ3: How should multiple specialized drafters be combined?

• Finding: Inference-time composition is superior to weight-space averaging.
  • Checkpoint Averaging: Performed the worst, often failing to capture the strengths of either specialist.
  • Confidence Routing: Significantly improved over single-domain baselines by selecting the appropriate specialist per prompt.
  • Merged-Tree Verification: Achieved the highest acceptance length overall. By allowing the verifier to process candidates from both specialists simultaneously, it maximized proposal diversity without discarding potential correct tokens.

RQ4: Are confidence, entropy, and depth useful signals?

• Confidence vs. Entropy: Confidence is a far superior signal for routing. It successfully routed MathInstruct drafts to math tasks (~90-97% selection) and ShareGPT to chat tasks. Entropy was less discriminative (near 50/50 splits) but useful as a diagnostic (rejected tokens had higher entropy).
• Depth Analysis: Acceptance rates generally decline with speculative depth. However, task-matched specialists become increasingly dominant at deeper levels, suggesting that deep speculation relies on sustained agreement between the drafter and verifier.

RQ5: How does speculative depth affect exploration vs. exploitation?

• Finding: Shallow depths benefit from broad coverage (mixed data helps), while deeper depths rely on exploitation by the task-matched specialist. Merged-tree verification succeeds because it preserves diversity at shallow levels while allowing the verifier to filter for the best specialist at deeper levels.

4. Results Summary

• Best Performance: The Merged-Tree Verification strategy yielded the highest average acceptance lengths across all benchmarks and backbones (e.g., 5.11 for HASS on average at Temp 0, compared to ~4.12 for the best single-domain checkpoint).
• Worst Performance: Simple Checkpoint Averaging resulted in the lowest acceptance lengths (e.g., ~2.59 for HASS), demonstrating that naive weight interpolation destroys specialized capabilities.
• Routing Efficiency: Confidence-based routing allowed the system to dynamically switch between specialists, recovering most of the performance gap between a single generic model and the optimal specialist.

5. Significance

This paper fundamentally shifts the perspective on speculative decoding from a purely architectural problem to a data and systems problem.

1. Data-Centric Optimization: It proves that training draft models on domain-specific data is essential for high-performance speculative decoding in specialized domains.
2. Inference-Time Composition: It demonstrates that keeping specialized models separate and combining them at inference (via routing or merged trees) is vastly superior to merging them in weight space. This suggests a new paradigm for deploying LLM systems: maintain a "pool" of specialized drafters rather than a single monolithic one.
3. Practical Deployment: The findings offer a blueprint for building robust speculative decoding systems that can handle diverse workloads (e.g., a chatbot that also solves math problems) without sacrificing speed or accuracy.

Resources: The authors have released code, model weights, and datasets on Hugging Face and GitHub to facilitate further research in task-aware proposal distributions.