TAPS : Distributions de Propositions Sensibles à la Tâche pour l'Échantillonnage Spéculatif
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
Résumé
Le décodage spéculatif accélère la génération autorégressive en permettant à un modèle d'ébauche léger de proposer des tokens futurs qu'un modèle cible plus grand vérifie ensuite en parallèle. En pratique, cependant, les modèles d'ébauche sont généralement entraînés sur des corpus génériques étendus, ce qui laisse incertaine la mesure dans laquelle la qualité du décodage spéculatif dépend de la distribution d'entraînement de l'ébauche. Nous étudions cette question avec des ébauches légères HASS et EAGLE-2 entraînées sur MathInstruct, ShareGPT et des variantes à données mixtes, évaluées sur MT-Bench, GSM8K, MATH-500 et SVAMP. Mesurée par la longueur d'acceptation, l'entraînement spécifique à une tâche confère une spécialisation nette : les ébauches entraînées sur MathInstruct sont les plus performantes sur les benchmarks de raisonnement, tandis que celles entraînées sur ShareGPT le sont sur MT-Bench. L'entraînement sur données mixtes améliore la robustesse, mais des mélanges plus importants ne dominent pas à travers les températures de décodage. Nous étudions également comment combiner des ébauches spécialisées au moment de l'inférence. La moyenne naïve des checkpoints donne de mauvais résultats, tandis qu'un routage basé sur la confiance s'avère supérieur aux ébauches à domaine unique et que la vérification par arbre fusionné produit la longueur d'acceptation globale la plus élevée pour les deux architectures. Enfin, la confiance est un signal de routage plus utile que l'entropie : les tokens rejetés tendent à avoir une entropie plus élevée, mais la confiance produit des décisions de routage au niveau du benchmark beaucoup plus nettes. Ces résultats montrent que la qualité du décodage spéculatif dépend non seulement de l'architecture de l'ébauche, mais aussi de l'adéquation entre les données d'entraînement de l'ébauche et la charge de travail en aval, et que les ébauches spécialisées sont mieux combinées au moment de l'inférence que dans l'espace des poids.
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.