Multi-Concept Speculatieve Steekproef: Canonieke Architecturen en Theoretische Grenzen

Multi-Draft Speculative Sampling: Canonical Architectures and Theoretical Limits

October 23, 2024

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

cs.AI

Samenvatting

We overwegen multi-draft speculatieve steekproeven, waarbij de voorstelreeksen onafhankelijk worden genomen uit verschillende conceptmodellen. Bij elke stap neemt een conceptselectieschema op tokenniveau een lijst met geldige tokens als invoer en produceert een uitvoertoken waarvan de verdeling overeenkomt met die van het doelmodel. Eerdere werken hebben aangetoond dat het optimale schema (dat de kans maximaliseert om een van de invoertokens te accepteren) kan worden opgevat als een oplossing voor een lineair programma. In dit werk tonen we aan dat het optimale schema kan worden opgesplitst in een tweestapsoplossing: in de eerste stap wordt een schema van het type importance sampling (IS) gebruikt om een tussenliggend token te selecteren; in de tweede stap wordt (enkelvoudige concept) speculatieve steekproeven toegepast om de uitvoertoken te genereren. Voor het geval van twee identieke conceptmodellen stellen we verder 1) een noodzakelijke en voldoende voorwaarde vast voor de verdelingen van het doel- en conceptmodel waarbij de acceptatiekans gelijk is aan één en 2) bieden we een expliciete uitdrukking voor de optimale acceptatiekans. Onze theoretische analyse motiveert ook een nieuwe klasse van tokenniveau-selectieschema's op basis van gewogen importance sampling. Onze experimentele resultaten tonen consistente verbeteringen in de haalbare blokkefficiëntie en tokentarieven ten opzichte van basisschema's in verschillende scenario's.

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.