Échantillonnage spéculatif multi-brouillons : architectures canoniques et limites théoriques

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

Résumé

Nous considérons l'échantillonnage spéculatif multi-brouillon, où les séquences de propositions sont échantillonnées indépendamment à partir de différents modèles de brouillon. À chaque étape, un schéma de sélection de brouillon au niveau du jeton prend une liste de jetons valides en entrée et produit un jeton de sortie dont la distribution correspond à celle du modèle cible. Des travaux antérieurs ont démontré que le schéma optimal (qui maximise la probabilité d'accepter l'un des jetons en entrée) peut être formulé comme une solution à un programme linéaire. Dans ce travail, nous montrons que le schéma optimal peut être décomposé en une solution en deux étapes : dans la première étape, un schéma de type échantillonnage d'importance (IS) est utilisé pour sélectionner un jeton intermédiaire ; dans la deuxième étape, un échantillonnage spéculatif (à un seul brouillon) est appliqué pour générer le jeton de sortie. Dans le cas de deux modèles de brouillon identiques, nous établissons en outre 1) une condition nécessaire et suffisante sur les distributions des modèles cible et de brouillon pour que la probabilité d'acceptation soit égale à un et 2) fournissons une expression explicite pour la probabilité d'acceptation optimale. Notre analyse théorique motive également une nouvelle classe de schéma de sélection au niveau du jeton basée sur l'échantillonnage d'importance pondéré. Nos résultats expérimentaux démontrent des améliorations cohérentes dans l'efficacité des blocs réalisables et les taux de jetons par rapport aux schémas de base dans plusieurs scénarios.

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.