Campionamento Speculativo Multidraft: Architetture Canoniche e Limiti Teorici

Multi-Draft Speculative Sampling: Canonical Architectures and Theoretical Limits

October 23, 2024

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

cs.AI

Abstract

Consideriamo il campionamento speculativo multi-bozza, in cui le sequenze di proposte sono campionate indipendentemente da diversi modelli di bozza. Ad ogni passo, uno schema di selezione di bozza a livello di token prende in input un elenco di token validi e produce un token di output la cui distribuzione corrisponde a quella del modello target. Lavori precedenti hanno dimostrato che lo schema ottimale (che massimizza la probabilità di accettare uno dei token in input) può essere formulato come soluzione di un programma lineare. In questo lavoro mostriamo che lo schema ottimale può essere decomposto in una soluzione a due passaggi: nel primo passaggio viene utilizzato uno schema di tipo campionamento di importanza (IS) per selezionare un token intermedio; nel secondo passaggio viene applicato il campionamento speculativo (a singola bozza) per generare il token di output. Nel caso di due modelli di bozza identici, stabiliamo inoltre 1) una condizione necessaria e sufficiente sulle distribuzioni dei modelli target e di bozza affinché la probabilità di accettazione sia uguale a uno e 2) forniamo un'espressione esplicita per la probabilità di accettazione ottimale. La nostra analisi teorica motiva anche una nuova classe di schema di selezione a livello di token basata sul campionamento di importanza pesata. I nostri risultati sperimentali dimostrano miglioramenti consistenti nell'efficienza del blocco ottenibile e nei tassi di token rispetto agli schemi di base in diversi scenari.

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.