Mehrfachentwurf Spekulatives Sampling: Kanonische Architekturen und theoretische Grenzen

Multi-Draft Speculative Sampling: Canonical Architectures and Theoretical Limits

October 23, 2024

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

cs.AI

Zusammenfassung

Wir betrachten das mehrstufige spekulative Sampling, bei dem die Vorschlagssequenzen unabhängig voneinander aus verschiedenen Entwurfsmodellen ausgewählt werden. Bei jedem Schritt verwendet ein Token-Ebene Entwurfsauswahlverfahren eine Liste gültiger Token als Eingabe und erzeugt einen Ausgabetoken, dessen Verteilung der des Zielmodells entspricht. Frühere Arbeiten haben gezeigt, dass das optimale Verfahren (das die Wahrscheinlichkeit maximiert, einen der Eingabetoken zu akzeptieren) als Lösung eines linearen Programms formuliert werden kann. In dieser Arbeit zeigen wir, dass das optimale Verfahren in eine Zwei-Schritt-Lösung zerlegt werden kann: Im ersten Schritt wird ein Verfahren vom Typ Importance Sampling (IS) verwendet, um einen Zwischentoken auszuwählen; im zweiten Schritt wird das (einstufige) spekulative Sampling angewendet, um den Ausgabetoken zu generieren. Für den Fall von zwei identischen Entwurfsmodellen etablieren wir weiterhin 1) eine notwendige und hinreichende Bedingung für die Verteilungen des Ziel- und Entwurfsmodells, damit die Akzeptanzwahrscheinlichkeit gleich eins ist, und 2) geben einen expliziten Ausdruck für die optimale Akzeptanzwahrscheinlichkeit an. Unsere theoretische Analyse motiviert auch eine neue Klasse von Token-Ebene Auswahlverfahren basierend auf gewichteter Importance Sampling. Unsere experimentellen Ergebnisse zeigen konsistente Verbesserungen in der erreichbaren Blockeffizienz und Token-Raten gegenüber Baseline-Verfahren in einer Reihe von Szenarien.

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.