L'inférence LLM est un problème de plusieurs milliards de dollars. Notre nouvel article présente un algorithme SOTA en échantillonnage spéculatif multi-draft, Global Resolution, qui fait des avancées significatives dans ce problème. Déballage ci-dessous 🧵👇

Une approche pour une inférence efficace s'appelle l'échantillonnage spéculatif. Cela utilise un modèle « brouillon » peu coûteux pour produire des « suppositions » sur ce que le modèle cible plus grand aurait produit.

En exploitant les efficacités de parallélisme des GPU modernes, cela peut réduire le nombre de passes avant du modèle cible de plus de 5x.

L'échantillonnage spéculatif peut être généralisé pour prendre en compte plusieurs suppositions provenant de plusieurs modèles de brouillon. Mais il n'est pas clair quel est le meilleur algorithme pour combiner ces multiples suppositions.

Dans le cas à une étape, des travaux antérieurs ont montré que la solution optimale peut être trouvée en résolvant un programme linéaire de transport optimal, le OTLP.

Cependant, l'OTLP est extrêmement difficile à résoudre presque exactement car il croît de manière exponentielle en taille de vocabulaire. Alors, comment pouvons-nous le résoudre ?

La clé est d'exploiter une structure supplémentaire dans la construction de l'arbre de brouillon.

Les travaux précédents [Hu et al.] ont montré que lorsque l'arbre de brouillon est formé par un échantillonnage i.i.d., en dualisant l'OTLP, la valeur objective optimale peut être calculée en temps presque linéaire grâce à la minimisation sous-modulaire.

Cependant, jusqu'à notre travail, aucune méthode n'était capable de résoudre la solution qui atteignait cette valeur objective optimale. Sans ce morceau manquant, tout ce que les travaux précédents nous donnent est l'efficacité des blocs, le gain de vitesse théorique maximum. Cela ne nous dit pas comment atteindre ce gain de vitesse.

Notre travail est le premier à réduire significativement la dimensionnalité de l'OTLP, en utilisant trois idées.

Nous renversons la dualisation de l'OTLP dans les travaux précédents [Hu et al.] avec la complémentarité des marges, pour formuler l'OTLP comme un problème de faisabilité de flux.

De nombreuses contraintes d'inégalité de flux sont redondantes. En utilisant un algorithme glouton issu de la théorie des polymatroides, nous pouvons les fusionner.

Ce problème de flux réduit a une solution qui peut être paramétrée comme le softmax d'un vecteur de faible dimension, et ce vecteur peut être calculé via une minimisation convexe. Cela réduit l'OTLP dans V^{n+1} variables en un problème de minimisation convexe dans V variables.

V peut cependant rester assez grand, c'est pourquoi dans notre article, nous appliquons d'autres approximations avec un taux d'erreur de modèle cible borné pour réduire encore le temps de calcul.

Dans de nombreux cas où V est limité aux top-k et aux n modèles de brouillon, comme indiqué ci-dessus, la Résolution Globale est le _seul_ solveur capable de résoudre l'OTLP dans un délai raisonnable.

De plus, en utilisant Global Resolution, nous pouvons améliorer les taux d'acceptation sur Llama et Gemma jusqu'à 6 % : En résumé, Global Resolution est SOTA pour une vérification multi-draft optimale dans le décodage spéculatif.

Il reste encore beaucoup de travail à faire ici, en assouplissant le cadre iid, ou en s'étendant à plusieurs étapes.