ImportanceWeighted Multi-Draft Speculative Sampling

We consider multi-draft speculative sampling, where the proposal sequences are sampled independently from the same underlying draft model. 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. Applying our decomposition result to the case of two drafts we 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. We study the performance of such schemes via experiments involving Llama 2-7B chat model for a natural language task and demonstrate improvements over prior approaches.