TAPS: Task-Aware Proposal Distributions for Speculative Sampling

Project Page · March 2026

Mohamad Zbib, Mohamad Bazzi, Ammar Mohanna, Bernard Ghanem*, Hasan Abed Al Kader Hammoud* (*equal advising)
KAUST · AUB — Contacts: mbz02@mail.aub.edu, hasanabedalkader.hammoud@kaust.edu.sa

---

Task-aware drafts that speed up speculative decoding without sacrificing quality.

Quick links

Contents

Overview • Abstract • Highlights • Figures • Results snapshot • Repository map • Setup • Citation • License

---

Overview

TAPS studies how the draft training distribution shapes speculative decoding quality. Using Meta-Llama-3-8B-Instruct as the verifier and lightweight LLaMA-style drafters (HASS and EAGLE-2, ~0.8B params, shared tokenizer), the work compares single-domain training, mixed-domain training, and inference-time composition (confidence routing, merged-tree verification).

TL;DR (At a glance)

• Goal: Quantify how draft training data and specialist composition affect speculative decoding.
• Verifier: Meta-Llama-3-8B-Instruct.
• Drafts: HASS and EAGLE-2 (single-layer, ~0.8B params, shared tokenizer).
• Workloads: MT-Bench (chat), GSM8K, MATH-500, SVAMP (reasoning).
• Metric: acceptance length (lossless speculative decoding).
• Compute: single node, 4×A100.

Abstract

Speculative decoding speeds up autoregressive generation by letting a lightweight drafter propose tokens that a larger verifier checks in parallel. We study how much draft quality depends on the training distribution using HASS and EAGLE-2 drafts trained on MathInstruct, ShareGPT, and mixed variants, evaluated on MT-Bench, GSM8K, MATH-500, and SVAMP. Task-matched drafts specialize; mixed data aids robustness but is not uniformly dominant across temperatures. Among composition strategies, weight averaging underperforms, confidence routing improves, and merged-tree verification attains the highest acceptance length. Confidence is a stronger routing signal than entropy. Results show speculative decoding quality hinges on both draft architecture and the alignment between draft training data and downstream workload.

Highlights

• Task-aware specialization: ShareGPT drafts lead MT-Bench (e.g., HASS 3.98 vs. 2.90 for MathInstruct), while MathInstruct drafts dominate GSM8K/MATH-500/SVAMP.
• Mixed data is nuanced: Mixed 70k+70k (HASS) peaks at temperature 0 with average acceptance length 5.18 but drops at temperature 1 (3.69), where Mixed 35k+35k is steadier.
• Composition wins: Weight-space averaging is weakest (≈2.4–2.6). Confidence routing reaches ~4.8 average acceptance length; merged-tree verification is strongest (HASS 5.11, EAGLE-2 5.02 at temperature 0).
• Routing signal: Confidence cleanly separates workloads (MathInstruct chosen for 90.8% of GSM8K; ShareGPT for 81.2% of MT-Bench). Entropy is diagnostic but weaker for routing.

Figures

Checkpoint averaging is unstable across interpolation weights and remains weaker than inference-time composition.

Results snapshot

Benchmarks: MT-Bench (chat), GSM8K, MATH-500, SVAMP. Metric: average acceptance length (higher is better) under lossless speculative decoding. Tables stay inline (no collapses).

Temperature 0

Temperature 1

Repository map

Setup

• Python 3.10+; GPU memory sufficient for Meta-Llama-3-8B-Instruct plus draft checkpoints (paper used 4×A100).
• Prefer per-project virtualenvs to avoid dependency clashes.

HASS install:

python -m venv .venv && source .venv/bin/activate
pip install -r Hass-Code/requirements.txt

EAGLE install:

python -m venv .venv && source .venv/bin/activate
pip install -r Eagle-Code/requirements.txt
pip install -e Eagle-Code

Citation

@article{zbib2026taps,
  title={TAPS: Task Aware Proposal Distributions for Speculative Sampling},
  author={Zbib, Mohamad and Bazzi, Mohamad and Mohanna, Ammar and Hammoud, Hasan Abed Al Kader and Ghanem, Bernard},
  journal={arXiv preprint arXiv:2603.27027},
  year={2026}
}