FR-Spec: Acelerando Modelos de Linguagem de Grande Vocabulário por meio de Amostragem Especulativa Baseada em Frequência

FR-Spec: Accelerating Large-Vocabulary Language Models via Frequency-Ranked Speculative Sampling

February 20, 2025

Autores: Weilin Zhao, Tengyu Pan, Xu Han, Yudi Zhang, Ao Sun, Yuxiang Huang, Kaihuo Zhang, Weilun Zhao, Yuxuan Li, Jianyong Wang, Zhiyuan Liu, Maosong Sun

cs.AI

Resumo

A amostragem especulativa emergiu como uma técnica importante para acelerar o processo de geração auto-regressiva de grandes modelos de linguagem (LLMs) ao utilizar um mecanismo de rascunho-verificação para produzir múltiplos tokens por passagem direta. Embora os métodos de amostragem especulativa mais avançados utilizem apenas uma única camada e uma cabeça de modelagem de linguagem (LM) como modelo de rascunho para alcançar uma impressionante compressão de camadas, seus ganhos de eficiência são substancialmente reduzidos para LLMs com vocabulários extensos, como o Llama-3-8B, que possui um vocabulário de 128k tokens. Para resolver isso, apresentamos o FR-Spec, um framework de amostragem especulativa com classificação de frequência que otimiza a seleção de candidatos de rascunho por meio da compressão do espaço de vocabulário. Ao restringir a busca de rascunho a um subconjunto de tokens priorizados por frequência, nosso método reduz a sobrecarga de computação da cabeça LM em 75%, garantindo a equivalência da distribuição final de saída. Experimentos em múltiplos conjuntos de dados demonstram uma aceleração média de 1,12 vezes em relação ao método de amostragem especulativa mais avançado, o EAGLE-2.

English

Speculative sampling has emerged as an important technique for accelerating the auto-regressive generation process of large language models (LLMs) by utilizing a draft-then-verify mechanism to produce multiple tokens per forward pass. While state-of-the-art speculative sampling methods use only a single layer and a language modeling (LM) head as the draft model to achieve impressive layer compression, their efficiency gains are substantially reduced for large-vocabulary LLMs, such as Llama-3-8B with a vocabulary of 128k tokens. To address this, we present FR-Spec, a frequency-ranked speculative sampling framework that optimizes draft candidate selection through vocabulary space compression. By constraining the draft search to a frequency-prioritized token subset, our method reduces LM Head computation overhead by 75% while ensuring the equivalence of the final output distribution. Experiments across multiple datasets demonstrate an average of 1.12times speedup over the state-of-the-art speculative sampling method EAGLE-2.