FR-Spec: Aceleración de Modelos de Lenguaje de Gran Vocabulario mediante Muestreo Especulativo Basado en Frecuencia

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

Resumen

El muestreo especulativo ha surgido como una técnica importante para acelerar el proceso de generación auto-regresiva de los grandes modelos de lenguaje (LLMs, por sus siglas en inglés) mediante la utilización de un mecanismo de borrador-verificación para producir múltiples tokens por paso hacia adelante. Si bien los métodos de muestreo especulativo más avanzados utilizan solo una capa y una cabecera de modelado de lenguaje (LM Head) como modelo borrador para lograr una impresionante compresión de capas, sus ganancias de eficiencia se reducen considerablemente en LLMs con vocabularios extensos, como Llama-3-8B, que cuenta con un vocabulario de 128k tokens. Para abordar este problema, presentamos FR-Spec, un marco de muestreo especulativo basado en frecuencia que optimiza la selección de candidatos borradores mediante la compresión del espacio de vocabulario. Al restringir la búsqueda de borradores a un subconjunto de tokens priorizados por frecuencia, nuestro método reduce la sobrecarga computacional de la LM Head en un 75% mientras garantiza la equivalencia de la distribución final de salida. Los experimentos en múltiples conjuntos de datos demuestran una aceleración promedio de 1.12 veces sobre el método de muestreo especulativo más avanzado, 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.