Inference acceleration for large language models using "stairs" assisted greedy generation

📄 arXiv: 2407.19947v1 📥 PDF

作者: Domas Grigaliūnas, Mantas Lukoševičius

分类: cs.CL, cs.LG

发布日期: 2024-07-29

备注: Accepted at the 29th International Conference on Information Society and University Studies (IVUS 2024)

💡 一句话要点

提出“阶梯”辅助贪婪生成方法，加速大型语言模型推理。

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

关键词: 大型语言模型 推理加速 辅助生成 贪婪解码 文本生成

📋 核心要点

  1. 大型语言模型计算资源需求高昂，小型模型精度可能不足，如何在精度和效率间取得平衡是核心问题。
  2. 论文提出“阶梯”辅助贪婪生成方法，利用小模型的快速生成能力和大模型的批量预测能力，加速推理过程。
  3. 实验结果表明，该方法在文本生成任务中，与独立的大型语言模型相比，推理时间减少了9.58%到17.24%，且不损失精度。

📝 摘要（中文）

大型语言模型（LLM）拥有数十亿参数，以其卓越的预测能力而闻名，但也需要大量的计算资源。随着它们越来越受欢迎，即使是所需资源的少量减少也可能对环境产生影响。另一方面，较小的模型需要的资源较少，但可能会牺牲准确性。本文提出了一种“阶梯”辅助贪婪生成方法。这是一种改进的辅助生成方法，它利用较小模型的快速生成、大型模型的批量预测和“阶梯”验证，从而实现预测生成的加速。结果表明，在文本生成任务中，与独立的LLM预测相比，推理时间减少了9.58%到17.24%，且没有损失准确性。

🔬 方法详解

问题定义：大型语言模型（LLM）虽然具有强大的预测能力，但推理过程需要消耗大量的计算资源。现有方法要么依赖于计算资源充足的硬件，要么牺牲模型精度来换取更快的推理速度。如何在保证模型精度的前提下，降低LLM的推理成本，是本文要解决的核心问题。

核心思路：本文的核心思路是利用一个小型模型辅助大型模型的推理过程。小型模型生成速度快，可以快速生成候选token，然后利用大型模型进行批量验证，从而减少大型模型的推理次数，加速整体推理过程。同时，引入“阶梯”验证机制，确保生成质量，避免精度损失。

技术框架：该方法主要包含三个阶段：1) 小型模型快速生成：使用一个小型语言模型快速生成多个候选token；2) 大型模型批量预测：将小型模型生成的多个候选token输入到大型模型中进行批量预测，得到每个token的概率分布；3) “阶梯”验证：设计一种“阶梯”验证机制，根据大型模型的预测结果，选择最优的token作为最终输出。

关键创新：该方法最重要的创新点在于“阶梯”验证机制。该机制通过设置不同的阈值，逐步筛选候选token，从而在保证生成质量的同时，尽可能地减少大型模型的推理次数。与传统的辅助生成方法相比，该方法能够更有效地利用小型模型的生成能力，并避免精度损失。

关键设计：论文中没有明确说明具体的参数设置、损失函数或网络结构等技术细节。但是，可以推断，“阶梯”验证机制中的阈值设置是关键的设计参数，需要根据具体的任务和模型进行调整。此外，小型模型的选择也会影响最终的性能，需要选择与大型模型具有一定相关性，但规模较小的模型。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，与独立的LLM预测相比，该方法在文本生成任务中，推理时间减少了9.58%到17.24%，且没有损失准确性。这一结果表明，该方法能够有效地加速大型语言模型的推理过程，并在保证生成质量的前提下，显著降低计算成本。

🎯 应用场景

该研究成果可广泛应用于各种需要快速文本生成的场景，例如智能客服、机器翻译、文本摘要等。通过降低大型语言模型的推理成本，可以使其更容易部署在资源受限的设备上，从而推动人工智能技术的普及和应用。此外，该方法还可以用于加速其他类型的序列生成任务。

📄 摘要（原文）

Large Language Models (LLMs) with billions of parameters are known for their impressive predicting capabilities but require lots of resources to run. With their massive rise in popularity, even a small reduction in required resources could have an impact on environment. On the other hand, smaller models require fewer resources but may sacrifice accuracy. In this work, we are proposing an implementation of ``stairs'' assisted greedy generation. It is a modified assisted generation methodology that makes use of a smaller model's fast generation, large model's batch prediction, and "stairs" validation in order to achieve a speed up in prediction generation. Results show between 9.58 and 17.24 percent inference time reduction compared to a stand-alone large LLM prediction in a text generation task without a loss in accuracy.