VOCABTRIM: Poda de Vocabulário para Decodificação Especulativa Eficiente em LLMs

Resumo

Neste artigo, apresentamos uma técnica simples e livre de treinamento para melhorar o desempenho de métodos de decodificação especulativa (SpD) baseados em rascunhos que incorporam o cabeçalho de modelagem de linguagem (LM head) durante o processo de rascunho. A decodificação especulativa baseada em rascunhos utiliza um ou mais modelos de linguagem menores, conhecidos como rascunhos ou modelos de rascunho, para amostrar uma sequência ou árvore de rascunho composta por múltiplos tokens, seguida pela verificação por um LLM base, o modelo alvo, que aceita um subconjunto como sua geração válida. Como geralmente se considera que a decodificação especulativa requer um mapeamento um-para-um entre os vocabulários do modelo alvo e do modelo de rascunho, tem sido natural compartilhar o vocabulário entre eles, ou até mesmo compartilhar o LM head, como em EAGLE ou Medusa. Primeiro, identificamos que esse esquema de amostragem de tokens de rascunho contém inerentemente uma sobrecarga desnecessária de inferência durante o rascunho, especialmente para alguns LLMs alvos com vocabulários muito grandes. Em seguida, propomos uma técnica simples, o VocabTrim, para mitigar a sobrecarga de rascunho e melhorar a velocidade de geração em ambientes com restrição de memória. O VocabTrim reconstrói o LM head do rascunho para conter apenas um conjunto limitado de tokens, selecionados pelos mais frequentemente amostrados do vocabulário do modelo alvo. Embora a limitação do vocabulário no rascunho degrade ligeiramente a taxa de aceitação, ela reduz significativamente a latência de rascunho em processos com restrição de memória, o que é frequentemente o caso em dispositivos de borda, resultando em um maior aumento de velocidade com restrição de memória (MBSU). Mostramos que nosso método pode impulsionar o aumento de velocidade com restrição de memória para modelos Llama-3 no Spec-Bench, especificamente em 16% para o 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: Poda de Vocabulário para Decodificação Especulativa Eficiente em LLMs

VOCABTRIM: Vocabulary Pruning for Efficient Speculative Decoding in LLMs

Resumo

Support