TAPS: Distribuciones de Propuestas Conscientes de la Tarea para Muestreo Especulativo
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
Resumen
La decodificación especulativa acelera la generación autoregresiva permitiendo que un modelo borrador ligero proponga tokens futuros que un modelo objetivo más grande verifica en paralelo. Sin embargo, en la práctica, los modelos borradores suelen entrenarse con corpus genéricos amplios, lo que deja incierto hasta qué punto la calidad de la decodificación especulativa depende de la distribución de entrenamiento del borrador. Estudiamos esta cuestión con los borradores ligeros HASS y EAGLE-2 entrenados en variantes de MathInstruct, ShareGPT y datos mixtos, evaluados en MT-Bench, GSM8K, MATH-500 y SVAMP. Medida por la longitud de aceptación, la formación específica de tareas produce una especialización clara: los borradores entrenados con MathInstruct son más fuertes en benchmarks de razonamiento, mientras que los entrenados con ShareGPT son más fuertes en MT-Bench. El entrenamiento con datos mixtos mejora la robustez, pero las mezclas más grandes no dominan en las diferentes temperaturas de decodificación. También estudiamos cómo combinar borradores especializados en el momento de la inferencia. El promediado simple de checkpoints tiene un rendimiento pobre, mientras que el enrutamiento basado en confianza mejora respecto a los borradores de dominio único y la verificación por árbol fusionado produce la mayor longitud de aceptación general para ambos modelos. Finalmente, la confianza es una señal de enrutamiento más útil que la entropía: los tokens rechazados tienden a tener mayor entropía, pero la confianza produce decisiones de enrutamiento a nivel de benchmark mucho más claras. Estos resultados demuestran que la calidad de la decodificación especulativa depende no solo de la arquitectura del borrador, sino también de la concordancia entre los datos de entrenamiento del borrador y la carga de trabajo objetivo, y que los borradores especializados se combinan mejor en el momento de la inferencia que en el espacio 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.