清华大学联合斯坦福大学提出混合注意力机制MoA，大模型解码速率提高6倍

夕小瑶科技说 原创
作者 | Richard

随着大语言模型的规模不断扩大，如何在保持模型性能的同时提高其效率，成为了当前研究的热点问题。最近，清华大学联合斯坦福大学提出一种名为"注意力混合"(Mixture of Attention, MoA)的新方法，用于自动压缩大语言模型。

MoA的核心思想是为模型中的每个注意力头设计独特的稀疏注意力模式。与传统方法不同，MoA考虑了不同注意力头的多样性，并根据输入长度动态调整注意力范围。通过精心设计的数据集和自动优化流程，MoA能够在大幅提高模型效率的同时，保持甚至超越原始密集模型的性能。

实验结果表明，MoA显著扩展了模型的有效上下文长度，提高了检索准确率，并在各项基准测试中表现出色。更重要的是，MoA实现了这些改进的同时，还大大降低了计算资源的需求。

这项研究不仅为大语言模型的压缩和优化提供了新的思路，也为提高模型在长文本处理方面的能力开辟了新的可能性。

论文标题：
MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression

论文链接：
https://arxiv.org/pdf/2406.14909

大语言模型的"减肥"之旅：如何让AI既聪明又高效？

近年来，大语言模型(LLMs)在各种应用中展现出惊人的能力，从回答问题到创作文章，再到辅助编程，似乎无所不能。然而，随着模型规模的不断扩大，一个棘手的问题逐渐浮出水面：如何在保持模型性能的同时提高其效率？特别是当处理长文本时，这个问题变得尤为突出。

想象一下，你正在使用一个强大的AI助手来帮你总结一本长篇小说。理想情况下，AI应该能够理解整本书的内容，并给出准确的摘要。但现实情况是，随着输入文本长度的增加，模型的计算需求呈指数级增长，导致处理速度变慢，甚至无法完成任务。这不仅影响用户体验，也限制了AI在更广泛领域的应用。

为了解决这个问题，研究人员提出了各种稀疏注意力方法。这些方法的核心思想是限制每个词只关注其周围固定范围内的其他词，从而减少计算量。乍一看，这似乎是个不错的解决方案。但正如文章中的实验所展示的，这种简单的限制可能会严重影响模型的长距离理解能力。

下图生动地展示了这一问题。在这个实验中，研究者比较了三种模型在信息检索任务上的表现：

• 原始模型能够从整个输入中有效检索信息，但效率较低；
• 采用统一稀疏注意力的StreamingLLM虽然提高了效率，但只能有效检索最近的一半输入内容，远距离检索能力大幅下降；
• 而本文提出的MoA方法则在保持高效率的同时，仍能从整个输入中有效检索信息。这个实验清楚地表明，简单地限制注意力范围并不是一个理想的解决方案。

更深入的分析揭示了问题的复杂性。大语言模型中的不同注意力头实际上有着多样化的行为模式。有些主要关注局部上下文，而另一些则倾向于关注全局信息。更有趣的是，随着输入长度的增加，这些行为模式还会发生变化。这意味着，一刀切的稀疏化策略很难满足所有注意力头的需求。

基于这些发现，研究者意识到需要一种更智能、更灵活的方法来压缩大语言模型。这种方法应该能够考虑不同注意力头的独特行为，并根据输入长度动态调整注意力模式。同时，它还需要在提高效率的同时，尽可能保持模型的长距离理解能力。这就是本文提出的MoA方法的核心思想。

通过这种方式，研究者希望在模型效率和性能之间取得更好的平衡，为大语言模型在实际应用中面临的挑战提供一个有效的解决方案。

MoA让大语言模型更苗条更聪明

在了解了大语言模型面临的效率挑战后，研究者们提出了一种名为"注意力混合"（Mixture of Attention，MoA）的创新方法。这个方法就像是给AI模型量身定制了一套"智能眼镜"，让它能够更加高效地关注重要信息。

MoA的核心思想是为模型中的每个注意力头设计独特的稀疏注意力模式。这就好比在一个大型会议上，不同的与会者会根据自己的兴趣和职责关注不同的话题和发言人。有些人可能只关注近期的讨论，而有些人则需要回顾整个会议的内容。MoA就是要让AI模型也能像这样灵活地分配注意力。

那么，MoA是如何实现这一目标的呢？首先，研究者们设计了一个灵活的搜索空间。这个空间包含了各种可能的注意力模式，有些模式可能更关注局部信息，有些则更侧重全局视角。更重要的是，这些模式还能根据输入的长度动态调整。这就好比我们阅读时会根据文章长短来调整阅读策略一样。

具体而言，MoA使用了一个简单而有效的线性函数来定义注意力范围：

这里， 是注意力头 的注意力范围， 是输入长度， 和 是可调整的参数。这个公式就像是给每个注意力头配备了一副可调焦的眼镜，能够根据"观察距离"(输入长度)自动调整"焦距"(注意力范围)。

接下来是如何从这么多可能的模式中选出最佳组合，研究者们采用了一个巧妙的自动优化流程。这个过程包括三个主要步骤：分析、优化和验证。

1. 在分析步骤中，MoA会仔细观察模型在处理不同长度的文本时的表现，它会评估每个注意力值对模型预测的影响。这就像是在观察每个与会者在会议中的贡献度。

   这里使用注意力影响矩阵的数学工具量化：

   这个复杂看似的公式实际上是在计算当我们"遮蔽"某个注意力值时，模型预测会发生多大变化。它考虑了直接影响和间接影响，就像是在评估会议中某个发言被忽略时对整体讨论的影响。

2. 优化步骤则是根据这些观察结果，为每个注意力头选择最合适的注意力模式。这个过程会考虑多个目标，比如准确性和效率，就像在平衡会议的质量和时间一样。

3. 验证步骤会测试选出的方案在新的、未见过的数据上的表现。这确保了MoA不只是在特定数据上表现良好，而是具有广泛的适用性。

一个特别值得注意的创新是MoA在选择训练数据时的考虑。研究者们发现，使用具有长距离依赖关系的数据集，并以模型自己的输出作为参考，能够更准确地评估压缩的影响。这就像是让模型在真实的长篇对话中练习，而不是简单的短句问答。

通过这种精心设计的方法，MoA成功地在保持模型性能的同时，大幅提高了效率。它不仅能处理更长的文本，还能在各种任务中保持高水平的表现。这就像是让AI既能快速阅读一本书，又能准确理解和总结其中的内容。

实验结果验证

研究者们进行了一系列全面的实验，不仅验证了MoA的有效性，还揭示了一些令人惊喜的发现。

首先，在整体性能方面，MoA展现出了令人印象深刻的表现。在保持50%密度的情况下，MoA在各种任务中都能与原始密集模型相媲美，甚至在某些情况下表现更优。特别是在检索任务中，MoA将相对性能下降控制在了8%以内，远优于其他方法。这意味着MoA不仅"减肥"成功，还保持了模型的"智力水平"。

长文本理解能力是大语言模型的一个关键指标。在这方面，MoA的表现尤为出色。下图左图展现出MoA成功地将模型的有效上下文长度扩展到了注意力范围的3.9倍。这就好比让一个原本只能阅读短篇文章的AI，现在能够轻松理解长篇小说了。更令人惊讶的是，如最右侧的图所示，MoA甚至能够处理长达60k的输入，并保持超过90%的检索准确率。这种长距离理解能力的提升，为AI在更复杂的任务中的应用打开了新的可能性。

效率提升是MoA的另一大亮点。MoA不仅减少了1.2倍到1.4倍的GPU内存使用，还将解码吞吐量提高了5.5倍到6.7倍。这意味着使用MoA后，AI模型不仅变得更"苗条"，还变得更"敏捷"了。

甚至与高度优化的vLLM框架相比，MoA仍然实现了1.4倍到1.5倍的吞吐量提升。这种效率的提升，无疑会大大降低AI应用的成本，使其更容易被广泛采用。

研究者们还进行了深入的消融研究，以了解MoA成功的关键因素。下表的结果清楚地表明，引入层级和头部的异质性，以及弹性规则，都对性能提升起到了重要作用。这就像是给AI配备了一套可以根据不同情况灵活调整的"智能眼镜"。

此外，研究者们还探讨了MoA如何捕捉和利用模型中的语义信息。研究者对MoA不同层和头部之间的分布进行可视化，这些发现不仅有助于我们更好地理解大语言模型的工作原理，还为未来的模型设计提供了valuable insights。

值得一提的是，MoA的自动优化过程也非常高效。整个过程通常只需要几个小时就能完成，这使得MoA成为一个实用的解决方案。

总的来说，这些实验结果充分证明了MoA是一个强大而灵活的方法。它不仅成功地减轻了大语言模型的计算负担，还在多个方面提升了模型的性能。

总结与展望

清华大学联合斯坦福大学提出的MoA方法为大语言模型的效率优化开辟了一条新路径。通过为每个注意力头设计独特的稀疏注意力模式，MoA成功地在保持模型性能的同时大幅提升了效率。

MoA的成功不仅解决了当前大语言模型面临的效率挑战，还为我们深入理解这些模型的工作原理提供了新的视角。它展示了如何通过智能的设计来平衡模型的性能和资源消耗，这对于AI技术的广泛应用至关重要。

未来，MoA的思路可能会被应用到更多类型的AI模型中，不仅限于语言处理领域。我们也期待看到这种方法如何与其他模型压缩技术结合，进一步推动AI技术的发展。