TAPS: Taakbewuste Voorstelverdelingen voor Speculatief Bemonsteren

TAPS: Task Aware Proposal Distributions for Speculative Sampling

March 27, 2026

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

cs.AI

Samenvatting

Speculatief decoderen versnelt autoregressieve generatie door een lichtgewicht draft-model toekomstige tokens te laten voorspellen, die een groter target-model vervolgens parallel verifieert. In de praktijk worden draft-modellen echter meestal getraind op brede, generieke corpora, waardoor het onduidelijk blijft in hoeverre de kwaliteit van speculatief decoderen afhangt van de draft-trainingsdistributie. Wij onderzoeken deze vraag met lichtgewicht HASS- en EAGLE-2 drafters, getraind op MathInstruct, ShareGPT en gemengde-data varianten, geëvalueerd op MT-Bench, GSM8K, MATH-500 en SVAMP. Gemeten aan de hand van acceptatielengte leidt taakspecifieke training tot duidelijke specialisatie: op MathInstruct getrainde drafts presteren het best op redeneerbenchmarks, terwijl op ShareGPT getrainde drafts het sterkst zijn op MT-Bench. Training met gemengde data verbetert de robuustheid, maar grotere mengsels domineren niet over verschillende decodeertemperaturen. Wij bestuderen ook hoe gespecialiseerde drafters tijdens inferentie gecombineerd kunnen worden. Naïeve checkpoint-averaging presteert slecht, terwijl confidence-gebaseerde routing verbetert ten opzichte van enkel-domein drafts en merged-tree-verificatie de hoogste algemene acceptatielengte oplevert voor beide backbones. Ten slotte is confidence een nuttiger routingsignaal dan entropie: verworpen tokens hebben tendentieel een hogere entropie, maar confidence produceert veel duidelijkere benchmark-level routingbeslissingen. Deze resultaten tonen aan dat de kwaliteit van speculatief decoderen niet alleen afhangt van de draft-architectuur, maar ook van de match tussen de draft-trainingsdata en de downstream-workload, en dat gespecialiseerde drafters beter gecombineerd kunnen worden tijdens inferentie dan in gewichtsruimte.

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.