Amostragem Especulativa Multirrascunho: Arquiteturas Canônicas e Limites 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

Resumo

Consideramos a amostragem especulativa multi-rascunho, onde as sequências de propostas são amostradas independentemente de diferentes modelos de rascunho. Em cada etapa, um esquema de seleção de rascunho ao nível do token recebe uma lista de tokens válidos como entrada e produz um token de saída cuja distribuição corresponde à do modelo alvo. Trabalhos anteriores demonstraram que o esquema ótimo (que maximiza a probabilidade de aceitar um dos tokens de entrada) pode ser formulado como uma solução para um programa linear. Neste trabalho, mostramos que o esquema ótimo pode ser decomposto em uma solução de dois passos: no primeiro passo, um esquema do tipo amostragem de importância (IS) é utilizado para selecionar um token intermediário; no segundo passo, a amostragem especulativa (de um único rascunho) é aplicada para gerar o token de saída. Para o caso de dois modelos de rascunho idênticos, estabelecemos ainda 1) uma condição necessária e suficiente sobre as distribuições dos modelos alvo e de rascunho para que a probabilidade de aceitação seja igual a um e 2) fornecemos uma expressão explícita para a probabilidade de aceitação ótima. Nossa análise teórica também motiva uma nova classe de esquema de seleção ao nível do token baseado em amostragem de importância ponderada. Nossos resultados experimentais demonstram melhorias consistentes na eficiência de bloco alcançável e nas taxas de token em relação aos esquemas de referência em diversos cenários.

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.