TAPS: Distribuições de Propostas Conscientes da Tarefa para Amostragem Especulativa
TAPS: Task Aware Proposal Distributions for Speculative Sampling
March 27, 2026
Autores: Mohamad Zbib, Mohamad Bazzi, Ammar Mohanna, Hasan Abed Al Kader Hammoud, Bernard Ghanem
cs.AI
Resumo
A decodagem especulativa acelera a geração autoregressiva permitindo que um modelo leve (draft) proponha tokens futuros que um modelo alvo maior então verifica em paralelo. Na prática, porém, os modelos draft são geralmente treinados em corpora genéricos e amplos, o que deixa incerto o quanto a qualidade da decodagem especulativa depende da distribuição de treinamento do draft. Investigamos esta questão com os modelos draft leves HASS e EAGLE-2 treinados no MathInstruct, ShareGPT e variantes de dados mistos, avaliados no MT-Bench, GSM8K, MATH-500 e SVAMP. Medido pelo comprimento de aceitação, o treinamento específico para a tarefa produz uma especialização clara: os drafts treinados no MathInstruct são mais fortes em benchmarks de raciocínio, enquanto os drafts treinados no ShareGPT são mais fortes no MT-Bench. O treinamento com dados mistos melhora a robustez, mas misturas maiores não dominam em todas as temperaturas de decodagem. Também estudamos como combinar drafts especializados durante a inferência. A média ingênua de checkpoints tem um desempenho fraco, enquanto o roteamento baseado em confiança supera os drafts de domínio único e a verificação por árvore mesclada (merged-tree) produz o maior comprimento de aceitação geral para ambas as arquiteturas. Finalmente, a confiança é um sinal de roteamento mais útil do que a entropia: tokens rejeitados tendem a ter entropia mais alta, mas a confiança produz decisões de roteamento a nível de benchmark muito mais claras. Estes resultados mostram que a qualidade da decodagem especulativa depende não apenas da arquitetura do draft, mas também da correspondência entre os dados de treinamento do draft e a carga de trabalho alvo, e que drafts especializados são melhor combinados durante a inferência do que no espaço de pesos.
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.