Dynamic Vocabulary Pruning in Early-Exit LLMs

📄 arXiv: 2410.18952v2 📥 PDF

作者: Jort Vincenti, Karim Abdel Sadek, Joan Velja, Matteo Nulli, Metod Jazbec

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-10-24 (更新: 2024-10-30)

期刊: NeurIPS 2024 ENLSP Workshop

💡 一句话要点

提出动态词汇表剪枝方法，加速早退LLM的推理效率并保持性能。

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

关键词: 大型语言模型 早退机制 词汇表剪枝 动态剪枝 推理加速

📋 核心要点

  1. 大型语言模型推理效率低，早退方法通过中间层预测token来加速，但大词汇表导致置信度估计计算成本高。
  2. 论文提出动态词汇表剪枝方法，在初始层剪枝词汇表，后续层使用剪枝后的词汇表，降低计算量。
  3. 实验结果表明，该方法在提高早退LLM置信度估计效率的同时，保持了具有竞争力的模型性能。

📝 摘要（中文）

大型语言模型（LLM）的规模增大通常能带来更好的性能，但同时也导致推理速度变慢和成本增加。早退（Early-exiting）是一种有前景的方法，它允许在中间层进行下一个token的预测，从而提高LLM推理的效率。然而，现代LLM中庞大的词汇量使得退出决策所需的置信度估计计算成本高昂，降低了效率提升。为了解决这个问题，我们提出了一种在测试时为每个token动态剪枝词汇表的方法。具体来说，在初始层之一剪枝词汇表，然后在剩余的前向传播中使用较小的词汇表。实验表明，这种事后动态词汇表剪枝提高了早退LLM中置信度估计的效率，同时保持了具有竞争力的性能。

🔬 方法详解

问题定义：论文旨在解决早退LLM中，由于词汇表过大导致的置信度估计计算量过高的问题。现有早退方法虽然能通过在中间层提前预测来加速推理，但庞大的词汇表使得在每一层进行置信度计算的开销很大，抵消了早退带来的效率提升。因此，如何在保证模型性能的前提下，降低词汇表大小，是本研究要解决的核心问题。

核心思路：论文的核心思路是在推理过程中动态地剪枝词汇表。具体来说，对于每一个token，在模型的前几层就确定一个较小的候选词汇子集，然后在后续的层中只考虑这个子集。这样可以显著减少每一层置信度估计的计算量，从而加速推理过程。这种动态剪枝是post-hoc的，即不需要重新训练模型。

技术框架：整体框架可以分为两个主要阶段：1) 词汇表剪枝阶段：在LLM的初始层（例如第一层或第二层），使用某种策略（例如基于token embedding的相似度）对词汇表进行剪枝，得到一个较小的候选词汇子集。2) 推理阶段：在后续的层中，只考虑剪枝后的词汇子集进行下一个token的预测。早退机制仍然存在，即在每一层都可以根据置信度决定是否提前退出。

关键创新：该方法最重要的创新点在于动态性和post-hoc性。动态性是指词汇表的剪枝是针对每个token进行的，而不是全局固定的。Post-hoc性是指该方法不需要重新训练模型，可以直接应用于现有的预训练LLM。与静态词汇表剪枝方法相比，动态剪枝可以更好地适应不同的输入token，从而在保证性能的同时，实现更高的效率提升。

关键设计：关键设计包括：1) 剪枝策略的选择：可以使用基于token embedding相似度、频率或其他指标的剪枝策略。2) 剪枝层数的选择：需要在效率和性能之间进行权衡，过早的剪枝可能会导致性能下降，过晚的剪枝可能效率提升不明显。3) 候选词汇子集的大小：子集越大，性能越好，但效率提升越小。4) 早退机制的集成：如何将动态词汇表剪枝与现有的早退机制有效结合。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

论文提出的动态词汇表剪枝方法能够在保持模型性能的同时，显著提高早退LLM的推理效率。具体的性能数据（例如加速比、精度损失）需要在论文中查找。该方法与没有词汇表剪枝的基线模型相比，在推理速度上应该有明显的提升，同时精度损失控制在可接受的范围内。实验结果将验证该方法在实际应用中的有效性。

🎯 应用场景

该研究成果可广泛应用于对推理速度有较高要求的场景，例如移动设备上的LLM部署、实时对话系统、以及需要快速响应的在线服务。通过降低推理计算量，可以减少延迟、降低能耗，并提高用户体验。未来，该方法可以与其他模型压缩技术（如量化、知识蒸馏）相结合，进一步提升LLM的部署效率。

📄 摘要（原文）

Increasing the size of large language models (LLMs) has been shown to lead to better performance. However, this comes at the cost of slower and more expensive inference. Early-exiting is a promising approach for improving the efficiency of LLM inference by enabling next token prediction at intermediate layers. Yet, the large vocabulary size in modern LLMs makes the confidence estimation required for exit decisions computationally expensive, diminishing the efficiency gains. To address this, we propose dynamically pruning the vocabulary at test time for each token. Specifically, the vocabulary is pruned at one of the initial layers, and the smaller vocabulary is then used throughout the rest of the forward pass. Our experiments demonstrate that such post-hoc dynamic vocabulary pruning improves the efficiency of confidence estimation in early-exit LLMs while maintaining competitive performance.