TAPS: Distribuzioni di Proposte Consapevoli del Compito per il Campionamento Speculativo
TAPS: Task Aware Proposal Distributions for Speculative Sampling
March 27, 2026
Autori: Mohamad Zbib, Mohamad Bazzi, Ammar Mohanna, Hasan Abed Al Kader Hammoud, Bernard Ghanem
cs.AI
Abstract
Il decoding speculativo accelera la generazione autoregressiva facendo sì che un modello draft leggero proponga token futuri che un modello target più grande verifica poi in parallelo. Tuttavia, nella pratica, i modelli draft sono solitamente addestrati su corpora generici e ampi, il che lascia incerto quanto la qualità del decoding speculativo dipenda dalla distribuzione dei dati di addestramento del draft. Studiamo questa questione con modelli draft leggeri HASS ed EAGLE-2 addestrati su varianti di MathInstruct, ShareGPT e dati misti, valutati su MT-Bench, GSM8K, MATH-500 e SVAMP. Misurata dalla lunghezza di accettazione, l'addestramento specifico per task produce una chiara specializzazione: i draft addestrati su MathInstruct sono più forti sui benchmark di ragionamento, mentre quelli addestrati su ShareGPT sono più forti su MT-Bench. L'addestramento su dati misti migliora la robustezza, ma miscele di dati più ampie non dominano attraverso le temperature di decoding. Studiamo anche come combinare modelli draft specializzati al momento dell'inferenza. La semplice media dei checkpoint ha prestazioni scarse, mentre un routing basato sulla confidenza migliora rispetto ai draft a dominio singolo e la verifica ad albero fuso (merged-tree) produce la lunghezza di accettazione più alta in assoluto per entrambe le architetture. Infine, la confidenza è un segnale di routing più utile dell'entropia: i token rifiutati tendono ad avere un'entropia più alta, ma la confidenza produce decisioni di routing a livello di benchmark molto più chiare. Questi risultati dimostrano che la qualità del decoding speculativo dipende non solo dall'architettura del draft, ma anche dalla corrispondenza tra i dati di addestramento del draft e il carico di lavoro downstream, e che i draft specializzati sono meglio combinati al momento dell'inferenza che nello spazio dei pesi.
English
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.