Attention Is All You Need But You Don't Need All Of It For Inference of Large Language Models
作者: Georgy Tyukin, Gbetondji J-S Dovonon, Jean Kaddour, Pasquale Minervini
分类: cs.LG, cs.CL
发布日期: 2024-07-22
💡 一句话要点
通过选择性丢弃层，加速大语言模型Llama-v2的推理过程
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大语言模型 推理加速 注意力机制 Llama-v2 层丢弃 低延迟 模型优化
📋 核心要点
- 大语言模型推理对算力需求高，注意力机制的复杂度是瓶颈，降低推理延迟是关键挑战。
- 该论文探索推理时选择性丢弃Llama-v2模型的MLP层和注意力层，以加速推理并尽量减少性能损失。
- 实验表明，丢弃较深的注意力层能在性能损失较小的情况下显著加速推理，为实际部署提供了优化方向。
📝 摘要（中文）
近年来，大语言模型（LLM）的推理需求急剧增长，但由于注意力层的复杂度与输入长度呈平方关系，低延迟的模型服务仍然面临挑战。本文研究了在Llama-v2模型推理时，丢弃MLP层和注意力层对性能的影响。研究发现，仅丢弃较深的注意力层只会略微降低性能，但在加速方面效果最佳，同时也可以丢弃整个层。例如，在一个13B的Llama2模型中，移除33%的注意力层会导致OpenLLM基准测试上的平均性能下降1.8%。此外，我们观察到，跳过除后几层之外的其他层会降低性能（跳过的层越多，性能越低），但跳过注意力层的情况除外。
🔬 方法详解
问题定义：大语言模型（LLM）的推理成本高昂，特别是注意力机制的计算复杂度随输入长度呈平方增长，导致推理延迟较高。现有方法难以在保证性能的同时显著降低推理延迟。因此，如何高效地加速LLM推理，使其能够以低延迟服务于用户，是一个亟待解决的问题。
核心思路：该论文的核心思路是通过在推理阶段选择性地丢弃Llama-v2模型的部分层（包括MLP层和注意力层），以减少计算量，从而加速推理过程。关键在于找到一种策略，使得在丢弃尽可能多的层的情况下，对模型性能的影响最小。论文发现，丢弃较深的注意力层对性能的影响相对较小，但能带来显著的加速效果。
技术框架：该研究的技术框架主要包括以下几个步骤：1) 选择Llama-v2模型作为研究对象。2) 设计不同的层丢弃策略，包括丢弃MLP层、注意力层以及整个层。3) 在OpenLLM基准测试上评估不同策略下的模型性能，包括准确率、困惑度等指标。4) 测量不同策略下的推理延迟，并分析性能与延迟之间的权衡关系。
关键创新：该论文的关键创新在于发现了丢弃较深的注意力层能够在性能损失较小的情况下显著加速Llama-v2模型的推理。与以往的研究不同，该论文没有采用复杂的模型压缩或知识蒸馏技术，而是通过一种简单而有效的方法实现了推理加速。此外，该论文还系统地研究了不同层丢弃策略对模型性能的影响，为实际部署提供了有价值的参考。
关键设计：论文的关键设计在于对不同层进行选择性丢弃。具体而言，研究者尝试了以下几种策略：1) 随机丢弃层；2) 丢弃浅层；3) 丢弃深层；4) 仅丢弃注意力层；5) 仅丢弃MLP层。通过实验对比，发现丢弃较深的注意力层在性能和延迟之间取得了最佳的平衡。具体的丢弃比例也是一个关键参数，需要根据实际需求进行调整。论文中，作者尝试了丢弃33%的注意力层，并取得了较好的效果。
🖼️ 关键图片
📊 实验亮点
实验结果表明，在一个13B的Llama2模型中，移除33%的注意力层会导致OpenLLM基准测试上的平均性能下降仅1.8%，但可以显著加速推理过程。这一结果表明，通过选择性丢弃层，可以在性能和延迟之间取得良好的平衡，为实际部署提供了可行的优化方案。
🎯 应用场景
该研究成果可应用于各种需要低延迟大语言模型服务的场景，例如在线对话系统、智能客服、实时翻译等。通过选择性丢弃层，可以在保证服务质量的前提下，降低推理成本，提高用户体验。此外，该方法也可以推广到其他基于Transformer架构的模型上，具有广泛的应用前景。
📄 摘要（原文）
The inference demand for LLMs has skyrocketed in recent months, and serving models with low latencies remains challenging due to the quadratic input length complexity of the attention layers. In this work, we investigate the effect of dropping MLP and attention layers at inference time on the performance of Llama-v2 models. We find that dropping dreeper attention layers only marginally decreases performance but leads to the best speedups alongside dropping entire layers. For example, removing 33\% of attention layers in a 13B Llama2 model results in a 1.8\% drop in average performance over the OpenLLM benchmark. We also observe that skipping layers except the latter layers reduces performances for more layers skipped, except for skipping the attention layers.