Self-Selected Attention Span for Accelerating Large Language Model Inference

📄 arXiv: 2404.09336v1 📥 PDF

作者: Tian Jin, Wanzin Yazar, Zifei Xu, Sayeh Sharify, Xin Wang

分类: cs.CL, cs.AI

发布日期: 2024-04-14

💡 一句话要点

提出自选注意力范围以加速大语言模型推理

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

关键词: 大语言模型 推理效率 自选注意力 CUDA内核 自然语言处理 稀疏注意力 微调

📋 核心要点

  1. 现有的大语言模型在推理时计算效率低，尤其是在处理大量标记时，导致资源浪费和延迟。
  2. 本文提出了一种自选注意力范围的方法，使LLM能够在推理过程中动态识别和应用最小注意力范围，从而提升效率。
  3. 实验结果表明，使用自定义CUDA内核后，LLM推理吞吐量提高了28%，显示出显著的性能提升。

📝 摘要（中文）

大语言模型（LLMs）能够解决复杂任务，但在现代GPU上推理计算效率低下，主要由于生成新标记时需要关注的标记数量不断增加。为了解决这一低效问题，本文利用LLMs的自我优化能力，提出了一种方法来优化推理时间效率。通过两个具体任务（复杂算术表达式评估和新闻文章摘要），我们创建了定制数据集以微调LLM，使其学习解决任务的同时识别每个步骤所需的最小注意力范围。微调后的模型能够在推理过程中动态生成稀疏注意力掩码。我们开发了自定义CUDA内核，利用减少的上下文进行注意力计算，结果显示使用该内核可以提高LLM推理的吞吐量28%。

🔬 方法详解

问题定义：本文旨在解决大语言模型在推理过程中由于标记数量增加而导致的计算效率低下问题。现有方法在处理复杂任务时，往往需要关注过多的标记，造成资源浪费和延迟。

核心思路：论文的核心思路是通过微调LLM，使其能够自我识别在特定任务中所需的最小注意力范围，并在推理时动态应用这些范围，从而提高推理效率。

技术框架：整体架构包括两个主要模块：首先是微调阶段，利用定制数据集训练LLM以完成特定任务；其次是推理阶段，模型根据任务动态生成稀疏注意力掩码。

关键创新：最重要的技术创新点在于LLM能够自我选择注意力范围，并将其转化为稀疏注意力掩码，这一过程在推理时实时进行，显著提升了推理效率。与现有方法相比，减少了不必要的计算。

关键设计：在微调过程中，采用了特定的损失函数以鼓励模型识别最小注意力范围。此外，开发了自定义CUDA内核，以优化稀疏注意力的计算，确保在推理时能够高效利用减少的上下文。

📊 实验亮点

实验结果显示，使用自定义CUDA内核后，LLM推理的吞吐量提高了28%。这一显著提升表明，通过自选注意力范围的策略，能够有效减少计算资源的消耗，并加快任务处理速度。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的各种任务，如文本摘要、问答系统和对话生成等。通过提高大语言模型的推理效率，能够在实际应用中减少延迟，提升用户体验，具有重要的实际价值和广泛的应用前景。

📄 摘要（原文）

Large language models (LLMs) can solve challenging tasks. However, their inference computation on modern GPUs is highly inefficient due to the increasing number of tokens they must attend to as they generate new ones. To address this inefficiency, we capitalize on LLMs' problem-solving capabilities to optimize their own inference-time efficiency. We demonstrate with two specific tasks: (a) evaluating complex arithmetic expressions and (b) summarizing news articles. For both tasks, we create custom datasets to fine-tune an LLM. The goal of fine-tuning is twofold: first, to make the LLM learn to solve the evaluation or summarization task, and second, to train it to identify the minimal attention spans required for each step of the task. As a result, the fine-tuned model is able to convert these self-identified minimal attention spans into sparse attention masks on-the-fly during inference. We develop a custom CUDA kernel to take advantage of the reduced context to attend to. We demonstrate that using this custom CUDA kernel improves the throughput of LLM inference by 28%. Our work presents an end-to-end demonstration showing that training LLMs to self-select their attention spans speeds up autoregressive inference in solving real-world tasks.