Многократная спекулятивная выборка: канонические архитектуры и теоретические пределы

Multi-Draft Speculative Sampling: Canonical Architectures and Theoretical Limits

October 23, 2024

Авторы: Ashish Khisti, M. Reza Ebrahimi, Hassan Dbouk, Arash Behboodi, Roland Memisevic, Christos Louizos

cs.AI

Аннотация

Мы рассматриваем многодрафтовую спекулятивную выборку, где последовательности предложений выбираются независимо из различных моделей черновиков. На каждом шаге схема выбора черновика на уровне токенов принимает список допустимых токенов на вход и генерирует токен на выходе, чье распределение соответствует целевой модели. Предыдущие работы показали, что оптимальную схему (максимизирующую вероятность принятия одного из входных токенов) можно представить в виде решения линейной программы. В данной работе мы показываем, что оптимальную схему можно разложить на двухэтапное решение: на первом этапе используется схема типа важности выборки (IS) для выбора одного промежуточного токена; на втором этапе применяется спекулятивная выборка (однодрафтовая) для генерации выходного токена. Для случая двух идентичных моделей черновиков мы дополнительно 1) устанавливаем необходимое и достаточное условие для равенства вероятности принятия единице у распределений целевой и черновых моделей и 2) предоставляем явное выражение для оптимальной вероятности принятия. Наше теоретическое исследование также мотивирует новый класс схем выбора на уровне токенов на основе взвешенной важности выборки. Наши экспериментальные результаты демонстрируют последовательные улучшения в достижимой эффективности блока и скорости токенов по сравнению с базовыми схемами в ряде сценариев.

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.