VOCABTRIM : Élagage du vocabulaire pour un décodage spéculatif efficace dans les LLM

Résumé

Dans cet article, nous présentons une technique simple et sans entraînement pour améliorer les performances des méthodes de décodage spéculatif (SpD) basées sur des modèles rédacteurs, en intégrant une tête de modélisation du langage (LM head) lors du processus de rédaction. Le décodage spéculatif basé sur des rédacteurs utilise un ou plusieurs modèles de langage plus petits, appelés rédacteurs ou modèles rédacteurs, pour échantillonner une séquence ou un arbre de jetons, suivi d'une vérification par un modèle de langage de base (LLM), le modèle cible, qui accepte un sous-ensemble comme génération valide. Comme il est généralement admis que le décodage spéculatif nécessite une correspondance un à un entre les vocabulaires du modèle cible et du modèle rédacteur, il est naturel de partager le vocabulaire entre eux, voire de partager la LM head comme dans EAGLE ou Medusa. Nous identifions d'abord que ce schéma d'échantillonnage de jetons rédacteurs contient intrinsèquement une surcharge d'inférence inutile lors de la rédaction, en particulier pour certains LLM cibles avec des vocabulaires très étendus. Ensuite, nous proposons une technique simple, VocabTrim, pour atténuer cette surcharge et améliorer la vitesse de génération dans un environnement limité par la mémoire. VocabTrim reconstruit la LM head du rédacteur pour contenir uniquement un ensemble limité de jetons, sélectionnés parmi les plus fréquemment échantillonnés dans le vocabulaire du modèle cible. Bien que la limitation du vocabulaire lors de la rédaction dégrade légèrement le taux d'acceptation, elle réduit significativement la latence de rédaction dans un processus limité par la mémoire, ce qui est souvent le cas sur les appareils périphériques, entraînant une accélération plus importante dans des conditions de mémoire limitée (MBSU). Nous montrons que notre méthode peut augmenter l'accélération en mémoire limitée pour les modèles Llama-3 sur Spec-Bench, spécifiquement de 16 % pour Llama-3.2-3B-Instruct.

English

In this paper, we introduce a simple training-free technique to improve the performance of drafter-based speculative decoding (SpD) methods that incorporates language modeling head (LM head) during drafting process. A drafter-based speculative decoding leverages one or more smaller language models, a.k.a. drafters or draft models, to sample a draft sequence or tree consisting of multiple tokens, followed by verification by a base LLM, a target model, accepting a subset as its valid generation. As it is usually considered that the speculative decoding requires one-to-one mapping between vocabularies of the target model and the draft model, it has been natural to share the vocabulary between them, or even share the LM head as in EAGLE or Medusa. We first identify that this draft token sampling scheme inherently contains an unnecessary inference overhead in drafting, especially for some target LLMs with very large vocabularies. Then, we propose a simple technique, VocabTrim, to mitigate the drafting overhead to improve the generation speed in memory-bound environment. VocabTrim reconstructs the drafter LM head to contain only a limited set of tokens, selected by the most frequently sampled from the vocabulary of the target model. While limiting the vocabulary in drafting slightly degrades the acceptance rate, it significantly reduces the drafting latency in memory-bound process which is often the case on edge devices, resulting in higher memory-bound speed up (MBSU). We show that our method can boost the memory-bound speed-up for Llama-3 models on Spec-Bench, specifically by 16% for Llama-3.2-3B-Instruct.

VOCABTRIM : Élagage du vocabulaire pour un décodage spéculatif efficace dans les LLM

VOCABTRIM: Vocabulary Pruning for Efficient Speculative Decoding in LLMs

Résumé

Support