Self-Selected Attention Span for Accelerating Large Language Model Inference

大型语言模型（LLMs）可以解决具有挑战性的任务。然而，由于它们在生成新的token时必须考虑的token数量不断增加，它们在现代GPU上的推理计算非常低效。为了解决这个效率问题，我们利用LLMs的问题解决能力来优化它们自己的推理时间效率。我们展示了两个具体任务：（a）评估复杂的算术表达式和（b）总结新闻文章。对于这两个任务，我们创建了自定义数据集来微调LLM。微调的目标是双重的：首先，让LLM学习解决评估或总结任务，其次，训练它识别每个任务步骤所需的最小注意力跨度。结果，经过微调的模型能够在推理过程中将这些自我识别的最小注意力跨度转换为稀疏的注意力掩码。我们开发了一个自定义的CUDA内核来利用减少的上下文进行注意力计算。我们证明使用这个自定义的CUDA内核可以提高LLM推理的吞吐量28%。我们的工作提供了一个端到端的演示，展示了训练LLMs自行选择注意力跨度如何加速自回归推理以解决现实任务。

本论文旨在解决大型语言模型在推理计算中的效率低下问题，通过自我优化来提高效率。

通过针对特定任务创建自定义数据集，对大型语言模型进行微调，使其学习解决评估或摘要任务，并训练其识别每个步骤所需的最小注意力范围，从而在推理过程中动态生成稀疏的注意力掩码以提高效率。

论文提出的方法通过自我优化大型语言模型的注意力机制，使其能够在推理过程中动态生成稀疏的注意力掩码，从而提高效率。实验结果表明，使用自定义CUDA核心可以提高LLM推理的吞吐量28%。论文还使用了自定义数据集进行微调，并提供了开源代码。

近期的相关研究包括《Scaling Laws for Neural Language Models》、《Efficient Transformers: A Survey》等。