TAPS: Распределения предложений с учетом задачи для спекулятивной выборки

TAPS: Task Aware Proposal Distributions for Speculative Sampling

March 27, 2026

Авторы: Mohamad Zbib, Mohamad Bazzi, Ammar Mohanna, Hasan Abed Al Kader Hammoud, Bernard Ghanem

cs.AI

Аннотация

Спекулятивное декодирование ускоряет авторегрессионную генерацию, позволяя легковесной проектной модели предлагать будущие токены, которые затем параллельно проверяются более крупной целевой моделью. Однако на практике проектные модели обычно обучаются на широких общедоступных корпусах, из-за чего остается неясным, насколько качество спекулятивного декодирования зависит от распределения данных для обучения проектной модели. Мы исследуем этот вопрос на примере легковесных проектных моделей HASS и EAGLE-2, обученных на данных MathInstruct, ShareGPT и их смешанных вариантах, с оценкой на MT-Bench, GSM8K, MATH-500 и SVAMP. При измерении по длине принятия (acceptance length) обучение для конкретной задачи демонстрирует явную специализацию: проектные модели, обученные на MathInstruct, показывают наилучшие результаты на бенчмарках логических рассуждений, тогда как модели, обученные на ShareGPT, лидируют на MT-Bench. Обучение на смешанных данных повышает устойчивость, но более крупные смеси не доминируют при различных температурах декодирования. Мы также исследуем способы комбинирования специализированных проектных моделей во время вывода. Наивное усреднение чекпоинтов показывает низкую эффективность, тогда как маршрутизация на основе уверенности (confidence-based routing) превосходит модели, обученные на одном домене, а верификация по объединенному дереву (merged-tree verification) обеспечивает наибольшую длину принятия в целом для обеих архитектур. Наконец, уверенность является более полезным сигналом для маршрутизации, чем энтропия: отклоненные токены обычно имеют более высокую энтропию, но уверенность приводит к гораздо более четким решениям по маршрутизации на уровне бенчмарков. Эти результаты показывают, что качество спекулятивного декодирования зависит не только от архитектуры проектной модели, но и от соответствия между данными ее обучения и последующей рабочей нагрузкой, а также что специализированные проектные модели эффективнее комбинировать во время вывода, чем в пространстве весов.

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.