随着大语言模型在长文本场景下的需求不断涌现，其核心的注意力机制（Attention Mechanism） 也获得了非常多的关注。注意力机制会计算一定跨度内输入文本（令牌，Token）之间的交互，从而实现对上下文的理解。随着应用的发展，高效处理更长输入的需求也随之增长[1][2]，这带来了计算代价的挑战：注意力高昂的计算成本和不断增长的键值缓存（KV-Cache）代价。稀疏注意力机制可以有效缓解内存和吞吐量的挑战。然而，现有稀疏注意力通常采用统一的稀疏注意力模式，即对不同的注意力头和输入长度应用相同的稀疏模式。这种统一的方法难以捕捉到大语言模型中多样的注意力模式，导致不同注意力头的不同的精度-代价权衡被忽略。 最近，来自清华大学、无问芯穹和上海交通大学的研究团队发表了《MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression》，提出通过混合不同稀疏度的注意力头，使用 25% 的注意力稠密度，就可以记忆几乎 100% 的上下文。
论文标题：
MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression
论文链接：
https://arxiv.org/abs/2406.14909
代码链接：
https://github.com/thu-nics/MoA
项目主页：
https://nics-effalg.com/MoA
一、总览
在大语言模型中，不同的注意力头表现出各异的注意力模式和扩展规则：有的关注全局信息，有的则聚焦局部；有的注意力范围随输入长度增加而扩展，有的则保持不变。然而，现有的统一稀疏注意力机制破坏了这些固有的特性。为应对这一挑战，研究团队提出了混合稀疏注意力（Mixture of Sparse Attention, MoA）方法，它能够为不同的头和层定制独特的稀疏注意力配置。MoA 构建了一个包含多种注意力模式及其扩展规则的搜索空间。通过分析模型，评估潜在配置，MoA可以为每个注意力头找到最优的稀疏注意力模式和扩展规则。
实验结果显示，无需任何训练，MoA 就可以在保持平均注意力跨度不变的情况下，将有效上下文长度提升约 3.9 倍。模型效果上，在 Vicuna-7B、Vicuna-13B 和 Llama3-8B 模型上，MoA将长文本信息检索准确率提高了1.5-7.1倍，优于统一注意力基线方法 StreamingLLM。
此外，MoA缩小了稀疏与稠密模型之间的能力差距，在 50% 平均注意力跨度下，长上下文理解基准测试集的最大相对性能下降从基线方法的9%-36%降低至5%以内。在运行效率上，MoA 的稀疏注意力使得生成过程中 KV-Cache 长度不扩大便于内存管理，减少了注意力计算量，降低了存储量从而可增大批大小。结合 CUDA GPU 算子优化，MoA 相比于 FlashAttention2 和 vLLM 将7B和13B稠密模型的生成吞吐量分别可提升 6.6-8.2 以及 1.7-1.9 倍。方法在 12K 以内的输入长度上搜索压缩方案，压缩后模型可以在长达 256K 的输入长度上高效且精准地检索信息。
二、研究背景
多头自注意力（MHA）机制是大型语言模型的核心功能之一[3]。该机制从输入序列出发，通过线性变换将其转换成查询（Q）、键（K）和值（V）三个矩阵。这些矩阵与之前序列的键值缓存（KV-Cache）相结合，共同计算出注意力矩阵（A）。为了保持自回归特性，这一计算过程会通过因果掩膜（M）进行调整，最终得到输出（O）。具体公式如下：
在大语言模型的自回归推理过程中，分为两个阶段：预填充和解码。在预填充阶段，模型会处理整个输入序列，以生成初始的响应令牌。随后进入解码阶段，模型利用新产生的令牌以及之前缓存的K和V矩阵，逐步生成后续令牌，直至完成整个序列的生成。虽然这种迭代方法效果显著，但随着KV-Cache的不断扩展，它也带来了内存和计算资源的需求增加。
三、现有方法
之前的研究通过引入稀疏注意力方法来应对大型语言模型处理长上下文时的效率挑战。对于生成式的大型语言模型，主流的稀疏模式是采用统一跨度滑窗：即不论注意力头还是输入长度如何，都使用固定、均匀跨度的滑动窗口掩膜，这样每个文本仅关注其邻近的上下文区域。 此外，还会对最初的几个文本施加全局注意力，以便它们能够关注到所有其他文本。这种局部注意模式通过丢弃当前注意跨度之外的KV-Cache，显著降低了长序列场景下的内存需求[4][5]。原则上，尽管单个模型层的注意力是局部的，但通过多层模型的逐步传递，每个词最终都能获取全局信息，从而在理论上可以实现比平均注意力跨度更长的有效上下文长度[6]。
图注：在LongEval数据集上[7]，使用不同注意力方法的Vicuna-7B模型在不同输入长度和检索位置的检索精度。大海捞针实验[8]将大量的键值对作为输入，并在不同位置测试给键检索值时的精度。(a) 原始稠密注意力模型；(b)统一跨度滑窗StreamingLLM，注意力跨度减半，超出跨度时检索效果降低；(c) MoA，平均注意力跨度减半，超出跨度时检索效果依然优秀。
下图揭示了这个现象的一个可能解释：虽然部分注意力头专注于局部上下文，但另一些注意力头则关注了整个输入序列。因此，采用统一跨度的方法限制了那些关注全局上下文的头的注意力跨度，并且给关注局部上下文的头分配了过多的计算和内存资源。此外，随着输入长度的增加，某些注意力头需要比其他头更快地扩展其注意力跨度，以防止性能显著下降。遗憾的是，统一跨度的方法没有考虑到这种异质性，未能针对不同注意力头的需求分别调整其注意力范围。另外，现有的模型压缩技术在制定压缩策略时通常基于通用语料库和人工撰写的输出结果，这并不能精确地反映稀疏化对处理长上下文任务的具体影响。
图注：来自Vicuna-7B模型不同注意力头的注意力矩阵示例。每个注意力矩阵是从LongEval数据集的256个输入上取平均得到的。
四、本文方法
图注：MoA概览：(a) 稀疏注意力的搜索空间涵盖了滑窗跨度的异构弹性规则；(b) 自动压缩过程始于精心设计的校准数据集。MoA通过分析这个数据集中每个注意力值对模型预测的影响，揭示了不同候选弹性规则在不同输入长度下的准确性损失。在优化阶段，MoA为每个注意力头挑选出最合适的弹性规则，使得其在满足平均跨度约束的同时，尽可能减少模型的性能损失。
4.1 异质弹性规则的搜索空间
4.2 自动注意力跨度规则搜索
对于有很多层的大语言模型，异质弹性规则的搜索空间可能会变得非常大。因此，本工作设计了自动注意力跨度规则搜索方法，来对大语言模型的每一个注意力头确定最优的和 。这包括分析和优化两步。
4.2.1 分析
在实际操作中，本工作推导了注意力影响的形式化表达，并利用深度学习框架的反向传播机制高效地计算所需的偏导数值。在校准数据集的不同输入长度上，MoA 对滑动窗口的影响分别取平均值，以体现同一种异质弹性规则在不同长度输入下的影响。
在分析阶段完成后，MoA能够明确每种异质弹性规则在精度和效率之间的平衡。基于此，MoA可以将较长的注意力跨度分配给那些对压缩更为敏感的注意力头，而将较短的注意力跨度分配给那些对此不太敏感的注意力头。
4.2.2 优化
具体而言，MoA构建了一个多目标优化问题：最小化不同输入长度下的预测损失，同时确保平均注意力密度不超过用户设定的限制。MoA 采用混合整数优化方法求解这一多目标优化问题。求解得到的所有压缩方案均达到了帕累托最优，意味着在没有增加其他长度预测损失的前提下，无法进一步降低任何特定长度的预测损失。为了确保模型能够有效泛化至未曾见过的输入长度，MoA 会从所有帕累托最优压缩方案中，选择在未见过的验证长度上损失最小的方案作为最终采用的压缩策略。
通过自动注意力跨度规则搜索，MoA 在遵守用户定义的密度约束的同时，找到合适的异质弹性规则来最小化由注意力稀疏化引起的准确性损失。
4.3 校准数据集的设计与选择
同时，MoA 也强调了数据工程在大语言模型压缩中的重要性。本工作发现，使用具有长距离依赖性的数据集并参考原始大语言模型的响应对于准确分析压缩的影响至关重要。
本工作指出了常用的通用语言建模数据集的主要问题。这类数据集，例如人类编写的文本语料库，通过在整个语料库上进行下一个词预测作为监督信号。但是这主要捕获的是临近上下文之间的注意力模式，而忽略了长期上下文依赖性，无法解决像长距离检索这样的全局注意力任务。同时，模型响应和人类编写的监督之间存在显著的不对齐。例如，对于同一个问题，人类可能会回答'Blue'，而模型可能会回答'The blue color'。使用人类的答案进行监督，注意力影响是基于预测'Blue'的概率转移量化的，这与最终目标背道而驰，即难以保持原始模型预测'The'的关键注意力。
图注：校准数据集的设计：数据集内容，监督，和参考来源。所有测试都在8k的输入长度和25%的平均注意力密度下完成。
五、实验结果
5.1 精度
MoA 超越基线稀疏注意力方法，并在50%的平均注意力密度下达到了与原始稠密模型相当的性能。我们计算了稀疏模型相对于原始模型的性能下降。
对于长上下文检索检索任务，MoA 最大相对准确性下降为8%，远小于StreamingLLM、InfLLM 和 H2O 的 87%、58% 和 44%。平均来说，MoA 的相对准确性下降在1%以下，而 StreamingLLM、InfLLM 和 H2O 的 51%、41% 和 20%。
图注：在LongEval上进行上下文检索准确性测试。使用Vicuna-7B模型：(a)改变注意力跨度，比较有效上下文长度，(b)将输入长度设定为8k，比较不同注意力密度下的检索准确性。使用Llama3-8B模型：(c)将密度设定为50%，比较输入长度增加时的检索准确性。
不同注意力方法在 50% 密度下进行 (a) LV-Eval和(b) LongBench 长上下文理解基准测试。测试使用 Vicuna-7B和13B模型，以及 Llama3-70B 模型。分数相对于原始稠密模型进行归一化。
消融实验
我们进一步评估了不同程度的稀疏注意力混合对于最终性能的影响。从基本的统一跨度滑窗开始，通过依次引入不同程度的异质性（层间，注意力头间，输入长度间），最终模型性能在不断提升。
图注：在 25% 注意力密度下对搜索空间进行消融实验
5.2 效率
我们使用Huggingface框架支持的FlashAttention[10]和vLLM框架作为基线，对比 MoA 的效率。
下表比较了 MoA 相对于各种注意力机制和LLM框架的运行效率，并对 MoA 的每个设计带来的效率提升进行了消融分析。在50% 注意力密度下，MoA 相比FlashAttention2 将解码吞吐量提升了6.6-8.2倍。相比H2O和InfLLM，解码吞吐量提升了 1.2-4.0倍。与包含高度系统级优化的 vLLM框架[11]相比，MoA 仍实现了1.7-1.9 倍的吞吐量提升。MoA 还将GPU总内存减少1.2到1.4倍。
图注：不同框架在7B和13B模型上的效率分析。MoA 每个设计带来的效率提升通过消融分析分为四个部分。所有稀疏注意力方法都使用50%的注意力密度。解码吞吐量在A100-80GB GPU 显存能容纳的最大批大小下进行评估。
作者介绍
引用
[1] Chen, Shouyuan, et al. "Extending Context Window of Large Language Models via Positional Interpolation." ArXiv, 2023, abs/2306.15595,https://api.semanticscholar.org/CorpusID:259262376.
[2] Tworkowski, Szymon, et al. "Focused Transformer: Contrastive Training for Context Scaling." ArXiv, 2023, abs/2307.03170,https://api.semanticscholar.org/CorpusID:259360592.
[3] Vaswani, Ashish, et al. "Attention is all you need." Advances in Neural Information Processing Systems, vol. 30, 2017.
[4] Xiao, Guangxuan, et al. "Efficient Streaming Language Models with Attention Sinks." The Twelfth International Conference on Learning Representations, 2024.
[5] Han, Chi, et al. "Lm-infinite: Simple on-the-fly length generalization for large language models." arXiv preprint arXiv:2308.16137, 2023.
[6] Zaheer, Manzil, et al. "Big bird: Transformers for longer sequences." Advances in Neural Information Processing Systems, vol. 33, 2020, pp. 17283-17297.
[7] Li, Dacheng, et al. "How Long Can Open-Source LLMs Truly Promise on Context Length?" lmsys.org, June 2023,https://lmsys.org/blog/2023-06-29-longchat.
[8] Fu, Yao, et al. "Data Engineering for Scaling Language Models to 128K Context." ArXiv, 2024, abs/2402.10171, https://api.semanticscholar.org/CorpusID:267682361.
[9] Zhang, Zhenyu (Allen), et al. "H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models." ArXiv, 2023, abs/2306.14048,https://api.semanticscholar.org/CorpusID:259263947.
[10] Dao, Tri, et al. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." Advances in Neural Information Processing Systems, 2022.
[11] Kwon, Woosuk, et al. "Efficient Memory Management for Large Language Model Serving with PagedAttention." Proceedings of the 29th Symposium on Operating Systems Principles, 2023,https://api.semanticscholar.org/CorpusID:261697361.
[12] Xiao, Chaojun et al. “InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory.” (2024).
扫码观看！
“AI技术流”原创投稿计划
TechBeat是由将门创投建立的AI学习社区（www.techbeat.net）。社区上线500+期talk视频，3000+篇技术干货文章，方向覆盖CV/NLP/ML/Robotis等；每月定期举办顶会及其他线上交流活动，不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台，希望为AI人才打造更专业的服务和体验，加速并陪伴其成长。
投稿内容
// 最新技术解读/系统性知识分享 //
// 前沿资讯解说/心得经历讲述 //
投稿须知
稿件需要为原创文章，并标明作者信息。
我们会选择部分在深度技术解析及科研心得方向，对用户启发更大的文章，做原创性内容奖励
投稿方式
发送邮件到
melodybai@thejiangmen.com
或添加工作人员微信（yellowsubbj）投稿，沟通投稿详情；还可以关注“将门创投”公众号，后台回复“投稿”二字，获得投稿说明。
