Muestreo Especulativo Multidraft: Arquitecturas Canónicas y Límites Teóricos

Multi-Draft Speculative Sampling: Canonical Architectures and Theoretical Limits

October 23, 2024

Autores: Ashish Khisti, M. Reza Ebrahimi, Hassan Dbouk, Arash Behboodi, Roland Memisevic, Christos Louizos

cs.AI

Resumen

Consideramos el muestreo especulativo multi-borrador, donde las secuencias de propuestas se muestrean de forma independiente a partir de diferentes modelos de borrador. En cada paso, un esquema de selección de borrador a nivel de token toma una lista de tokens válidos como entrada y produce un token de salida cuya distribución coincide con la del modelo objetivo. Trabajos previos han demostrado que el esquema óptimo (que maximiza la probabilidad de aceptar uno de los tokens de entrada) se puede formular como una solución a un programa lineal. En este trabajo mostramos que el esquema óptimo se puede descomponer en una solución de dos pasos: en el primer paso se utiliza un esquema tipo muestreo de importancia (IS) para seleccionar un token intermedio; en el segundo paso se aplica un muestreo especulativo (de un solo borrador) para generar el token de salida. Para el caso de dos modelos de borrador idénticos, establecemos además 1) una condición necesaria y suficiente sobre las distribuciones de los modelos objetivo y de borrador para que la probabilidad de aceptación sea igual a uno y 2) proporcionamos una expresión explícita para la probabilidad de aceptación óptima. Nuestro análisis teórico también motiva una nueva clase de esquema de selección a nivel de token basado en muestreo de importancia ponderado. Nuestros resultados experimentales demuestran mejoras consistentes en la eficiencia de bloque alcanzable y las tasas de tokens sobre los esquemas base en varios escenarios.

English

We consider multi-draft speculative sampling, where the proposal sequences are sampled independently from different draft models. At each step, a token-level draft selection scheme takes a list of valid tokens as input and produces an output token whose distribution matches that of the target model. Previous works have demonstrated that the optimal scheme (which maximizes the probability of accepting one of the input tokens) can be cast as a solution to a linear program. In this work we show that the optimal scheme can be decomposed into a two-step solution: in the first step an importance sampling (IS) type scheme is used to select one intermediate token; in the second step (single-draft) speculative sampling is applied to generate the output token. For the case of two identical draft models we further 1) establish a necessary and sufficient condition on the distributions of the target and draft models for the acceptance probability to equal one and 2) provide an explicit expression for the optimal acceptance probability. Our theoretical analysis also motives a new class of token-level selection scheme based on weighted importance sampling. Our experimental results demonstrate consistent improvements in the achievable block efficiency and token rates over baseline schemes in a number of scenarios.