FR-Spec: Versnelling van Taalmodellen met Grote Vocabulaire via Frequentie-Gerangschikt Speculatief Samplen

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

February 20, 2025

Auteurs: 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

Samenvatting

Speculatieve sampling is naar voren gekomen als een belangrijke techniek voor het versnellen van het auto-regressieve generatieproces van grote taalmodellen (LLMs) door gebruik te maken van een draft-then-verify mechanisme om meerdere tokens per forward pass te produceren. Hoewel state-of-the-art speculatieve samplingmethoden slechts één enkele laag en een taalmodelleringskop (LM head) gebruiken als het draft-model om indrukwekkende laagcompressie te bereiken, worden hun efficiëntiewinsten aanzienlijk verminderd voor LLMs met een grote vocabulaire, zoals Llama-3-8B met een vocabulaire van 128k tokens. Om dit aan te pakken, presenteren we FR-Spec, een frequentie-gerangschikt speculatief sampling framework dat de selectie van draft-kandidaten optimaliseert door compressie van de vocabulaire-ruimte. Door de draft-zoekopdracht te beperken tot een frequentie-geprioriteerde token-subset, vermindert onze methode de rekenoverhead van de LM Head met 75% terwijl de gelijkwaardigheid van de uiteindelijke outputverdeling wordt gegarandeerd. Experimenten over meerdere datasets tonen een gemiddelde snelheidswinst van 1,12 keer ten opzichte van de state-of-the-art speculatieve samplingmethode 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.