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

📄 arXiv: 2502.14856v2 📥 PDF

作者: 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.CL, cs.AI, cs.LG

发布日期: 2025-02-20 (更新: 2025-03-11)

🔗 代码/项目: GITHUB

💡 一句话要点

提出FR-Spec，通过频率排序推测采样加速大词汇量语言模型

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 推测采样 大语言模型 词汇空间压缩 频率排序 模型加速

📋 核心要点

  1. 现有推测采样方法在大词汇量LLM中效率降低，主要瓶颈在于LM Head的计算开销。
  2. FR-Spec通过频率排序压缩词汇空间，优先选择高频token作为draft候选，降低LM Head的计算负担。
  3. 实验表明，FR-Spec在保证输出分布等价性的前提下，显著加速了LLM的生成速度，优于现有方法。

📝 摘要（中文）

推测采样是一种重要的加速大型语言模型(LLM)自回归生成过程的技术，它利用draft-then-verify机制，在每次前向传播中生成多个token。目前最先进的推测采样方法仅使用单层和一个语言模型(LM)头作为draft模型，以实现显著的层压缩，但对于大型词汇表的LLM（如具有128k词汇量的Llama-3-8B），其效率提升会大大降低。为了解决这个问题，我们提出了FR-Spec，一个频率排序的推测采样框架，通过词汇空间压缩来优化draft候选选择。通过将draft搜索限制在频率优先的token子集，我们的方法减少了75%的LM Head计算开销，同时确保了最终输出分布的等价性。在多个数据集上的实验表明，相比最先进的推测采样方法EAGLE-2，平均加速了1.12倍。

🔬 方法详解

问题定义：论文旨在解决大词汇量语言模型中，现有推测采样方法因LM Head计算开销过大而导致效率降低的问题。现有方法在draft阶段需要对整个词汇表进行搜索，计算量巨大，成为性能瓶颈。

核心思路：FR-Spec的核心思路是利用词汇表中token的频率信息，优先选择高频token作为draft阶段的候选token。通过限制搜索空间，显著减少LM Head的计算量，从而加速整个推测采样过程。这样设计的依据是，高频token在生成过程中出现的概率更高，优先选择它们可以提高draft的准确率，减少后续的验证步骤。

技术框架：FR-Spec的整体框架仍然遵循推测采样的draft-then-verify流程。首先，使用一个小的draft模型（通常是原模型的一个浅层版本）生成一个token序列（draft）。然后，使用原始的LLM对draft序列进行验证。如果验证通过，则接受draft序列中的多个token；如果验证失败，则回退到原始LLM进行单步生成。FR-Spec的关键在于优化了draft阶段的token选择过程，通过频率排序来限制搜索空间。

关键创新：FR-Spec最重要的创新点在于提出了频率排序的词汇空间压缩方法。与现有方法直接在整个词汇表上进行搜索不同，FR-Spec首先根据token的频率进行排序，然后只在频率最高的token子集上进行搜索。这种方法在保证输出分布等价性的前提下，显著降低了LM Head的计算开销。

关键设计：FR-Spec的关键设计在于确定合适的频率阈值或token子集大小。论文中可能采用了动态调整策略，根据不同的数据集或模型来选择最佳的阈值。此外，为了保证输出分布的等价性，FR-Spec可能需要对draft模型的输出进行校正，以补偿由于词汇空间压缩带来的偏差。具体的校正方法可能涉及到对概率分布进行重新归一化或使用其他技术。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，FR-Spec在多个数据集上实现了平均1.12倍的加速，优于最先进的推测采样方法EAGLE-2。该方法在保证生成质量的同时，显著降低了LM Head的计算开销，尤其是在大词汇量模型（如Llama-3-8B）上表现出色。这些结果验证了FR-Spec在加速LLM推理方面的有效性。

🎯 应用场景

FR-Spec可应用于各种需要加速大型语言模型推理的场景，例如：在线对话系统、文本生成、机器翻译等。通过提高LLM的生成速度，可以降低延迟，提升用户体验，并降低计算成本。该研究对于推动LLM在资源受限环境下的部署具有重要意义。

📄 摘要（原文）

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.12$\times$ speedup over the state-of-the-art speculative sampling method EAGLE-2. Code available at https://github.com/thunlp/FR-Spec.