破解LLM长上下文瓶颈：MoA混合注意力技术提升3.9倍上下文长度，性能飙升7倍！-轻识

Abstract

稀疏注意力能有效缓解大型语言模型（LLM）在长上下文中的显著内存和吞吐量需求。现有方法通常采用统一的稀疏注意力 Mask ，在不同注意力头和输入长度上应用相同的稀疏模式。然而，这种统一的方法未能捕捉到LLM中固有的多样化注意力模式，忽略了它们在准确性和延迟之间的不同权衡。为了应对这一挑战，作者提出了混合注意力（MoA），它自动为不同的头和层定制独特的稀疏注意力配置。MoA构建并导航了各种注意力模式及其相对于输入序列长度的缩放规则的搜索空间。它对模型进行剖析，评估潜在的配置，并确定最佳的稀疏注意力压缩计划。MoA适应不同的输入大小，发现一些注意力头扩展其关注范围以适应更长的序列，而其他头始终专注于固定长度的局部上下文。实验显示，MoA在相同的平均注意力跨度下，将有效上下文长度提高了，在Vicuna-7B、Vicuna-13B和Llama3-8B模型上，相较于统一注意力 Baseline ，将检索准确性提高了。此外，MoA缩小了稀疏模型和密集模型之间的能力差距，将最大相对性能下降从降低到两个长上下文理解基准上的以内。在单个GPU上，对于7B和13B密集模型，MoA实现了的GPU内存减少，并将解码吞吐量提升了，对性能的影响最小。

1 Introduction

大型语言模型（LLM）在众多应用中显示出卓越的多功能性[6; 57; 63]。LLM的核心是注意力机制[62]，它计算一定范围内的标记之间的交互，从而实现上下文理解。扩展输入长度对于增强LLM的能力至关重要[7; 60]，包括事实检索、总结、小样本学习、问题回答等等[4; 70]。然而，不断增长的注意力计算和键值缓存（KV-Cache）提出了重大的效率挑战[54; 69; 26; 33]。

先前的工作提出稀疏注意力方法来解决生成LLM中长上下文的效率挑战。这些方法通常在所有头和输入长度上采用统一的、固定跨度的滑动窗口 Mask ，将注意力限制在局部上下文[69; 26]中。这种方法允许LLM以固定注意力跨度接受长输入，保持有限的注意力计算和KV缓存开销。遵循先前的工作[7; 60]，作者将有效上下文长度量化为内容检索准确率超过90%阈值的最大输入长度。原则上，固定跨度的局部注意力可以通过多个模型层逐渐聚合全局信息，产生比每个注意力跨度更长的有效上下文长度[18]。然而，作者发现像StreamingLLM[69]这样的统一 Mask 几乎无法将有效上下文长度扩展到跨度之外，如图6所示。图1(b)进一步说明了这种限制：使用50%的注意力跨度 Mask ，StreamingLLM无法准确检索输入前半部分的内容，并且在更长的输入长度上表现更差。图2揭示了解释这一问题的可能原因：一些注意力头关注局部上下文，而另一些则涵盖了整个输入序列的全局范围。因此，统一方法无法实现长的有效上下文长度，因为它限制了关注全局上下文的头的注意力跨度，同时过度分配了局部上下文头的计算和内存预算。此外，随着输入长度的增加，一些注意力头需要比其他头更快地增加注意力跨度，以避免严重的性能下降，如表1所示。不幸的是，统一方法没有包括异构规则，以不同的方式缩放不同头的注意力跨度。此外，现有的模型压缩方法[44, 40, 68, 36, 31, 37]使用通用语言建模语料库来决定压缩计划，这不能准确分析压缩对长上下文任务的影响。

在这项工作中，作者提出了混合注意力（MoA），一种无需训练的稀疏注意力方法。如图3所示，MoA构建了异构弹性注意力跨度的搜索空间。对于自动LLM压缩，MoA首先利用基于梯度的分析来检查每个注意力位置对预测损失的影响。基于分析结果，MoA为每个模型层和注意力头定制异构稀疏注意力配置。在分析过程中，MoA使用具有长距离依赖关系的校准数据集，并使用原始密集模型的响应（而不是人工编写的响应）作为计算损失的参考。这确保了准确分析注意力影响，以便获得更好的压缩结果。作者的贡献总结如下。

异构弹性规则。作者提出了每个注意力头的 Mask 的异构弹性规则。作者制定了MoA压缩搜索空间，包括一系列不同的弹性规则，这些规则针对每个注意力头相对于输入长度的局部注意力跨度进行定制。与为每个头提供统一跨度和缩放函数的 Mask 相比，异构弹性规则将MoA的事实检索准确率从25%提高到了98%。
校准数据集构建 作者强调了数据工程在LLM压缩中的重要性。作者的研究发现，使用具有长距离依赖关系的数据集和参考原始LLM的响应对于准确分析压缩影响至关重要。
自动优化。作者提出了一种自动化的流程，以找到包含各种注意力头的异构弹性规则的优化压缩计划。该流程可以在几个小时内有效地找到最优计划，例如，压缩Vicuna-13B需要两小时。实验表明，MoA在7B和13B密集LLM上以50%的密度（即KV-Cache长度/输入长度的平均值）实现了5.5倍到6.7倍的吞吐量改进，而在检索准确率上的平均相对退化仅为1%。此外，MoA在仅25%的平均密度下实现了超过90%的检索准确率，远远超过了需要75%到100%的密度才能达到类似性能的稀疏注意力 Baseline 。在长上下文理解基准测试中，MoA的表现与密集模型相当，最大相对性能下降不到5%，约为使用统一稀疏注意力 Baseline 观察到的六分之一。

2 Preliminary and Related work

Attention mechanism

多头自注意力（MHA）机制[62]对LLM的功能至关重要。它从输入序列开始，通过线性投影转换为 Query （Q）、键（K）和值（V）矩阵。这些矩阵与前序序列中缓存的K和V（KV-Cache）结合，计算注意力矩阵（A）。这个计算通过因果 Mask （M）进行修改，以确保自回归属性，得到输出（O），如方程1所示：

LLM中的自回归推理包括两个阶段：预填充和解码。在预填充阶段，模型处理整个输入序列以生成初始响应标记。在随后的解码阶段，它使用新生成的标记和先前缓存的K和V矩阵来产生后续标记，直到生成结束。尽管这种方法有效，但这一迭代过程由于扩展的KV-Cache而增加了内存和计算需求。

Efficient Attention

提出了一些高效的方法来减轻与注意力机制相关的计算和内存成本。一方面的工作使用动态稀疏注意力 Mask ，在预填充阶段自适应地跳过注意力计算[46; 52; 53; 64; 43; 32]，或者在解码阶段根据输入序列丢弃KV-Cache[3; 75; 20; 54; 41]。然而，由于复杂的控制和计算流程，动态预填充通常需要特定硬件才能实现实质性的实时加速[52; 64; 43; 25; 44]。此外，在解码阶段动态剪枝KV-Cache可能需要广泛的重新训练[3]，或者需要额外的累积注意力分数计算[54; 75; 41; 20]。

另一方面的工作使用静态稀疏注意力，其中预定义的 Mask 在所有处理过的句子中一致应用。由于固定的计算流程，静态稀疏注意力通常更高效，对GPU更友好。对于像BERT[16]这样的语言理解模型，使用了各种 Mask [73; 5; 9; 76; 69; 26]。但对于生成式大型语言模型（LLMs），主流方法是固定跨度滑动窗口 Mask ，在几个初始 Token 上使用全局注意力[69; 26]。通过局部注意力模式，可以丢弃当前注意力范围之外的KV-Cache，为长序列场景节省大量内存。然而，不同注意力头和输入长度之间统一的静态 Mask 是模型和数据不可知的，这可能会损害LLMs的有效上下文长度，导致在长序列场景中性能不佳。作者的方法属于这一类别，从效率和免训练的优势中受益，同时解决之前方法遇到的性能限制。

除了稀疏注意力，还提出了其他机制来替代传统的注意力进行长序列建模。然而，这些新机制通常需要与普通 Transformer 不同的权重，对于LLMs来说，这带来了显著的重新训练开销。

以前的工作还提出了LLM加速框架，以及 Kernel Level 的优化[13; 12]。这些 Kernel 和系统优化与作者的工作是正交的，并且可以集成以进一步提高效率。

3 Mixture of Attention (MoA)

作者首先在3.1节中说明了在预训练的LLM中注意模式的不一致性。基于这一洞察，作者在3.2节定义了作者的混合注意力（MoA）方法的搜索空间。

Mixture of Attention Patterns and Elastic Rules

异质注意力模式。如图2所示，LLM中的不同注意力头显示出异质的注意力模式。例如，第一个注意力头主要关注具有窄跨度滑动窗口的局部上下文，而第三个注意力头几乎覆盖了整个输入，表明全局注意力。表1显示，在模型层中应用相同的滑动窗口稀疏注意力 Mask 可能会导致检索准确度有65%的方差。这符合多头自注意力设计原则，即捕捉不同信息[62]，以及同时期研究识别特定注意力头进行全局文本检索的发现[67]。

异质弹性规则。除了在特定长度上的异质性，不同的注意力头在输入长度变化时也表现出不同的弹性行为。图2说明了这种可变性：对于较短的输入（注意力矩阵的左上部分），第二个和第三个注意力头最初显示出全局注意力。然而，随着输入长度的增加，第二个注意力头保持中等跨度的局部关注，而第三个注意力头继续扩展为全局注意力。表1进一步证明了不同的弹性规则。例如，在4k输入长度下，第9到11层上的2k滑动窗口稀疏注意力 Mask 比第17到19层上的 Mask 产生更好的检索准确度。然而，对于8k输入长度，情况则相反。这些数据支持了图2中的视觉观察，强调了注意力模式对不同输入长度缩放的不同响应。利用这些洞察，MoA将异质弹性规则作为搜索空间。

Heterogeneous Elastic Rule Search Space

在设计MoA Mask 的搜索空间时，作者考虑了LLM注意力模式固有的异质性和弹性特性。如图3(a)所示，作者采用一种硬件友好的滑动窗口 Mask 作为作者基础的稀疏注意力 Mask [5]。借鉴之前的工作[69; 26]，最初的几个标记（对于MoA来说是64个标记）不被 Mask 。注意力跨度等于滑动窗口跨度加上初始未 Mask 标记的数量。为了简化，作者使用简单的线性函数定义了在输入长度下，头的注意力跨度：

其中和是控制基础跨度和随输入长度扩展速率的超参数。为了确保硬件友好性，除非有特别指定，否则作者限制MoA每个模型层的不同规则数量最多为两个。

每个注意力头的和超参数可以从多个离散选项中选择。对于具有多层的LLM，MoA搜索空间可能会变得相当大。例如，对于一个由个注意力头和层组成的7B模型，每个头的和有和个不同的选项，潜在的搜索空间扩展到种配置。因此，作者设计了一个自动化的 Pipeline ，以有效地为任何LLM定位最佳的和。## 4 MoA压缩的自动 Pipeline

本节概述了如图3(b)所示的MoA自动压缩 Pipeline 。从训练好的LLM和校准数据集开始，MoA首先分析了对于校准数据集中的各种输入序列，每个注意力值对模型预测损失的影响。Mask 后影响的总和表示了不同输入长度下与每个 Mask 相关的准确度损失，显示了在那种长度下，每个候选弹性规则可能造成的准确度损失。然后，MoA通过为每个头选择最佳弹性规则来优化压缩计划，该规则在遵循指定密度约束的同时，最小化各种长度下的准确度损失。以下各节详细讨论了此 Pipeline 中每个步骤的内容。

Attention Influence Profiling

在概要步骤中，MoA量化了单个注意力值对预训练LLM最终预测损失的影响。它告知后续步骤关于 Mask 每个注意力值的影响，揭示了每个注意力头候选弹性规则的准确性权衡。

每个注意力值的影响来源于注意力矩阵及其在校准数据集上计算的梯度。在应用稀疏注意力 Mask 时，作者使用基于注意力矩阵变化的首次泰勒展开来近似模型预测损失的变化：。这里，索引所有层的注意力头，是每个注意力矩阵中的行索引和列索引。

作者定义了注意力影响矩阵，作为如果 Mask 注意力值（即设为零）时，损失的估计变化。如方程3所示，这个度量考虑了 Mask 的直接和间接效果。为了简化表示，作者在这里省略了头索引。最初， Mask 直接将注意力值降为零，表示为。另外，注意力中的softmax函数将注意力矩阵的每一行的和规范化为1。因此，将列中的一个注意力值设为零会导致同一行中其他注意力值的增加。这两个效果被整合到以下公式中，其推导在附录D.1中提供：

在实践中，作者使用校准数据集上的反向传播来计算每个头在不同数据项上的平均注意力影响。分别对不同输入长度计算平均注意力影响。梯度是使用深度学习框架（如PyTorch [48]）中的链式导数计算的。详细的校准数据集设置将在第5节讨论。

有了每个头的平均注意力影响，MoA可以计算在特定输入长度应用候选弹性规则的准确性损失。该损失根据规则计算被 Mask 的平均注意力影响的总和。作者用表示对应于规则的头上的二进制 Mask ，被 Mask 的位置标记为0，其他位置标记为1。作者将准确性损失形式化为以下形式：

在概要阶段之后，MoA获得了弹性规则独特的准确度-密度权衡。它告知将更密集的 Mask 分配给更敏感的头，将较轻的 Mask 分配给较不敏感的头。在不同输入长度上进行概要分析，能够识别出对未见长度最有效的弹性规则。

Automatic Optimization

MoA自动选择每个注意力 Head 的最佳弹性规则，以在密度预算下减小各种序列长度的准确度损失。根据分析结果，MoA首先确定帕累托前沿压缩计划，在这些建议的长度中，任何在一个配置长度上准确度损失的改进都会使另一个变差。为了确保对那些未被配置的长度最佳的泛化，MoA然后选择在帕累托前沿解决方案中在未见长度下产生最小损失的方案作为最终计划。

具体来说，作者利用多目标优化搜索在配置长度范围内的一组帕累托最优压缩计划。每个长度的目标是最小化总准确度损失，同时符合任何用户定义的密度约束。该目标表述如下：

这里，上标表示不同长度的值；和分别表示配置长度和受密度约束的长度集合；表示候选规则集合；表示压缩导致的准确度损失；表示 Head 的规则的密度；表示平均密度约束；表示注意力 Head 的总数。

这种表述对应于经典的多目标混合整数规划问题，可以使用现有的线性求解器，如Gurobi [23]，在几分钟内有效解决。详细的问题表述和解题策略在附录D.2中讨论。

在帕累托最优压缩计划中，作者选择在未见验证长度下损失最小的计划作为最优解。这种方法使作者能够在不配置每个可能长度的情况下，提高计划对未见长度的有效泛化可能性。

得益于这个自动流程，作者高效地为每个注意力 Head 量身定制了弹性规则。通过这个流程，MoA在符合用户定义的密度约束的同时，最小化了由注意力稀疏化引起的准确度损失。

5 Dataset and Supervision

在本节中，作者强调了在LLM压缩过程中校准数据集设计及其监督目标被忽视的重要性。校准数据集对于各种压缩技术（包括权重剪枝[44; 34; 42]和量化[40; 68; 36; 31]）的敏感性分析至关重要。在这项工作中，MoA分析了校准数据集中注意力的影响，这对于后续的自动优化至关重要。

当前方法。通常使用通用语言建模数据集，如人类编写的文本语料库RedPajama [11]，作为校准数据集。这些数据集在整个语料库上进行下一 Token 预测的监督，主要捕捉与紧前 Token 一致的注意力模式。然而，它们缺乏长上下文依赖，无法解决像长距离检索等任务中至关重要的全局注意力。

此外，模型响应与人类编写的监督之间存在显著的错位。因此，在分析过程中使用人类响应来计算注意力和梯度会导致不准确。例如，对于同一个问题，人类可能会回答“蓝色”，而模型可能会生成“蓝色”。使用人类回答进行监督，基于预测“蓝色”的概率变化不准确量化注意力影响；这与保持原始模型预测“the”的关键注意力的目标相背离。这些不一致源于各种因素，包括位置、语气和同义词的不匹配。

MoA的方法。MoA通过整合长上下文依赖和模型对齐来增强校准数据集。它利用包含大量依赖长距离内容的摘要的长上下文MultiNews数据集[17]。摘要是原始密集模型生成的，并作为监督。与当前采用人类响应作为计算损失参考的方法相比，使用原始模型生成的响应作为监督可以促进准确的影响力分析，从而有利于压缩结果。

方法比较。作者通过改变数据集选择、监督类型和摘要参考，同时将数据项数量和长度标准化为50和8k个单词，来评估作者的设计效果。额外的设置和评估在附录A和B.3.1中。

作者通过比较使用不同数据集和监督方法的MoA压缩计划，展示了长距离依赖的重要性。表2中，RedPajama [11]代表通用语言建模数据集，而MultiNews [17]通过汇总关于单一事件的多个文档来强调长距离上下文。此外，每个MultiNews项还包括人类编写的摘要，提供了更强的长距离依赖和更好的性能。在MultiNews的摘要上计算损失会导致性能显著提高，检索准确率提高60%，困惑度降低0.98。

此外，使用原始密集模型生成的摘要作为监督，与自身的注意力模式与文本监督之间的对齐度更高。与可能不一致的人类摘要相比，它提高了性能，如表2的最后两行所示。

作者将MoA与最先进的静态和动态稀疏注意力 Baseline 进行比较，用于LLM：StreamingLLM [69]和H2O [75]。作者将LLM的“密度”定义为稀疏解码阶段KV-Cache长度与序列长度的平均值。值得注意的是，在MoA和StreamingLLM中，KV-Cache长度与它们在稀疏预填充阶段的注意力跨度相匹配。相比之下，H2O采用密集预填充，并根据累积的注意力分数生成动态 Mask ，这需要额外的计算。

模型和基准测试。作者从LMSys评估vicuna-7b-v1.5-16k、vicuna-13b-v1.5-16k [8]，以及从Gradient AI评估Llama-3-8B-Instruct-262k [1]。为了有效评估上下文长度，作者使用LongEval [35]测试每个长度 Level 的100个数据项的关键值检索准确性。为了全面评估能力，作者使用包含11个子任务的LV-Eval [70]和包含13个子任务的LongBench [4]。为了测试连贯性，作者在LongBench [15; 17; 38; 28; 45]的四个子集上测量困惑度，代表不同的能力。

MoA设置。作者在MultiNews [17]上对MoA进行概要分析，模型摘要长度分别为2k、4k和8k。使用12k验证数据集选择最优压缩计划。除非特别指定，对于每个模型的各个基准测试和长度，都使用相同的计划。这些模型没有经过微调。

Performance

总体性能。如表格3所示，MoA优于现有的稀疏注意力方法，并且在50%的密度下与原始密集模型的性能相当。作者计算了压缩模型相对于原始模型的性能下降。对于检索任务，MoA表现出最大8%的相对准确度下降（计算方式为，跨三种长度和LLM），远低于StreamingLLM和H2O的87%和44%。平均而言，MoA的相对准确度下降不到1%，而StreamingLLM和H2O分别高达51%和20%。在LV-Eval和LongBench的基准测试中，与MoA相比，StreamingLLM的分数最大下降了36%和27%，而MoA只下降了5%和3%。H2O的效率成本更高，最大下降分别为9%和4%。在困惑度测试中也显示出类似趋势，MoA保持相对困惑度增长不到1%，而StreamingLLM和H2O分别增长13%和4%。这一趋势在其他密度下也成立，如附录B.1.1和B.1.3所示。

长上下文检索。作者通过在LongEval [35]上进行检索测试来研究最终有效的上下文长度。如图4(a)所示，MoA将其有效上下文长度扩展到大约是注意力跨度的，在达到原始模型的12k限制之前。有效上下文长度计算的详细数据在附录B.1.2中。图4(b)显示，在固定的8k输入长度下，MoA在25%的密度下就能达到超过0.9的检索准确度，而StreamingLLM和H2O分别需要100%和75%的密度。图4(c)进一步研究了MoA在扩展和推广到极端输入长度方面的能力。通过在8k长度内进行性能分析，MoA在50%的密度下能够将能力扩展到60k输入，同时保持超过90%的检索准确度。长上下文理解。图5展示了在长上下文基准LV-Eval和LongBench上的性能比较。MoA在全面性能上与原始密集模型相当，以及效率成本更高的H2O。相比之下，StreamingLLM的性能不一致：尽管它偶尔会超越原始模型，但总体上会明显退化。

消融研究。作者在表5中评估了不同稀疏 Mask 搜索空间对性能的影响。从基本的均匀 Mask 开始，作者观察到通过顺序引入异质性显著增强：首先是层，然后是头，最后是弹性规则。

Efficiency

作者使用带有FlashAttention [13]的Huggingface [66]作为 Baseline ，并在其基础上实现了MoA。正如表5所示，在平均50%的密度下，MoA在各种模型尺寸和输入长度上，将总的GPU内存减少了到。此外，MoA将解码吞吐量提升了到，这主要归功于固定大小的KV-Cache（），减少了注意力计算（）以及增加了批处理大小（）。与使用密集预填充并需要额外成本生成动态 Mask 的H2O相比，MoA在相同的50%密度下实现了到的吞吐量提升。与高度优化的vLLM框架[33]相比，MoA仍然实现了到的吞吐量增加。更多效率结果展示在附录B.2中。

Rules Discovered by MoA

作者研究了MoA每个 Head 的弹性规则。如图8所示，初始层和中间层的 Mask 显示出高密度，这与之前关于LLM内在维度[61]和层敏感性[71]的研究结论相一致。相反，在最终层中，大多数 Head 需要低密度，而只有少数需要高密度。图9显示，平均 Mask 密度较低的层通常在 Head 之间显示出更广泛的密度范围，这证实了同一层内需要异质性的需求。更多细节和量化的见解在附录C中呈现。

7 Conclusion and Future Work

MoA为每个注意力头和输入长度自动选择异构弹性 Mask ，显著将LLM的有效上下文长度延长了。它将检索准确性提高了至，并在50%平均密度下将吞吐量增加到超过，同时在严格的基准测试中保持了与密集模型相当的性能。

局限性和未来工作。在极低密度预算下，MoA无法保持良好性能。设计一种动态MoA方法有潜力解决这一问题，作者将其留作未来工作。为了进一步提高MoA的效率，可以整合系统级和 Kernel 级的优化，例如 Kernel 融合[27]和KV-Cache管理[33]。使用具有有限注意力跨度的非线性弹性规则也值得探索。此外，MoA的分析方法可以适应于评估权重和其他激活的影响，从而促进其他压缩方法，如权重和激活量化。

参考

[1].MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression.