FR-Spec: Ускорение языковых моделей с большим словарным запасом с помощью частотного ранжированного спекулятивного сэмплирования

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

February 20, 2025

Авторы: 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

Аннотация

Спекулятивная выборка стала важной техникой для ускорения авторегрессивного процесса генерации больших языковых моделей (LLM) за счет использования механизма "черновик-затем-проверка" для создания нескольких токенов за один прямой проход. Хотя современные методы спекулятивной выборки используют только один слой и голову языкового моделирования (LM) в качестве черновой модели для достижения впечатляющего сжатия слоев, их эффективность значительно снижается для LLM с большим словарным запасом, таких как Llama-3-8B с 128 тысячами токенов. Для решения этой проблемы мы представляем FR-Spec — частотно-ранжированную структуру спекулятивной выборки, которая оптимизирует выбор кандидатов для черновика за счет сжатия пространства словаря. Ограничивая поиск черновика подмножеством токенов, упорядоченных по частоте, наш метод снижает вычислительную нагрузку на голову LM на 75%, сохраняя эквивалентность итогового распределения выходных данных. Эксперименты на нескольких наборах данных демонстрируют среднее ускорение в 1,12 раза по сравнению с современным методом спекулятивной выборки 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.