KIMI 月之暗面提出（MoBA）大模型长文本处理新解法：块注意力混合 精华

（MoBA）大语言模型长文本处理新解法：块注意力混合

在自然语言处理领域，高效处理长文本一直是个老大难问题。随着大语言模型在阅读、理解和生成文本方面的能力不断提升，其处理输入的核心——注意力机制，却逐渐成了瓶颈。在典型的Transformer架构里，这种机制要把每个词元（token）和其他所有词元进行比较，这就导致计算成本会随着序列长度的增加呈二次方增长。如今，我们把语言模型应用到长篇文档、多章节书籍、法律文书或是大型代码库这类需要处理海量文本信息的任务中时，这个问题就更加突出了。要是模型得处理几万甚至几十万个词元，简单粗暴地计算全注意力的成本实在高得离谱。

以往解决这个问题的方法，常常是设置固定结构或者采用近似算法，但这些在某些场景下可能会影响效果。比如说，滑动窗口机制只能让词元关注局部区域，这样就容易忽略掉重要的全局关系。还有些方法，像直接用全新结构替换softmax注意力机制，从根本上改变了基础架构，可这就需要从头开始大规模重新训练模型，很难利用现有的预训练模型。所以，研究人员一直都在寻找一种新方法，既能保留原始Transformer设计的优势，也就是适应性强、能捕捉广泛依赖关系，又不会在处理超长序列时，产生传统全注意力机制带来的巨额计算开销。

来自Moonshot AI、清华大学和浙江大学的研究人员，带来了块注意力混合（MoBA）这个创新方法，它把专家混合（MoE）的原理运用到了注意力机制当中。MoBA会把输入划分成一个个便于处理的“块”，再通过一个可训练的门控系统，来确定每个查询词元应该关注哪些块，这样就解决了模型挨个比较词元时效率低下的问题。和那些强制使用局部注意力或者窗口注意力的方法不同，MoBA能让模型自己学习该重点关注哪里。这个设计遵循“less structure”原则，意思就是架构不会提前规定哪些词元必须相互作用，而是把这些决策交给训练好的门控网络。

MoBA有个很关键的特点，就是能和现有的基于Transformer的模型完美配合。它没有抛弃标准的自注意力接口，而是像个“插件”一样，直接就能替换使用。MoBA的参数数量和原模型一样，不会让架构变得臃肿，而且还保留了因果掩码，保证自回归生成的准确性。在实际应用中，MoBA可以在稀疏注意力和全注意力之间灵活切换。处理超长输入时，用稀疏注意力能提高速度；在训练的某些层或者阶段，如果有需要，还能切换回标准的全注意力模式。

可训练的区块稀疏注意力

完整的上下文被划分成多个区块，每个查询词元都能学会关注最相关的键值（KV）区块，这样就能高效处理长序列数据。

无参数门控机制

MoBA引入了一种很新颖的无参数top - k门控机制，专门为每个查询词元挑选最相关的区块，确保模型只聚焦在信息量最大的区块上。

在完全注意力和稀疏注意力之间无缝过渡

MoBA设计得非常灵活，可以完美替代全注意力机制，在全注意力模式和稀疏注意力模式之间自由切换。

技术细节和优势

MoBA的核心操作，就是把上下文划分成多个区块，每个区块都包含一连串连续的词元。门控机制会计算查询词元和每个区块之间的“亲和度”分数，一般是通过把查询和区块里键的聚合表示进行对比来实现。然后，它会选出得分最高的那些区块。这样一来，最终的注意力分配就只和最相关区块里的词元有关。而且，包含查询词元本身的那个区块一定会被纳入，保证局部上下文信息随时都能获取到。同时，MoBA还采用了因果掩码，让词元不会关注未来的位置，维持从左到右的自回归特性。

经过这样的处理，MoBA的注意力矩阵比原始Transformer的要稀疏得多。但它依然很灵活，要是有需要，查询也能获取到远距离的信息。举个例子，要是在文本末尾提出的问题，得参考文本开头的细节才能回答，门控机制就能给开头相关的区块打高分。从技术层面来讲，这种基于区块的方法把词元比较的数量减少到了亚二次方级别，随着上下文长度增加到几十万甚至上百万个词元，效率提升就更加明显了。

MoBA还有个很大的优势，就是和现代加速器以及专用内核的兼容性很好。研究人员把MoBA和FlashAttention结合起来，FlashAttention是一个高性能库，能实现快速、低内存消耗的精确注意力计算。他们根据所选的区块，仔细对查询 - 键 - 值操作进行分组，进一步优化了计算过程。研究人员表示，处理一百万个词元时，MoBA相比传统的全注意力机制，速度能提升大约6倍，这在实际应用场景里优势非常大。

结果和洞察

根据技术报告，MoBA在各种任务中的表现和全注意力机制不相上下，但处理长序列时，计算成本却低得多。在语言建模数据测试里，当序列长度达到8192或者32768个词元时，MoBA的困惑度和全注意力Transformer的很接近。更重要的是，当研究人员把上下文长度逐渐增加到128000甚至更长时，MoBA对长上下文的理解能力依然很强。研究人员还做了“trailing token”评估，主要测试模型预测长提示末尾词元的能力，这个测试通常能暴露那些过度依赖近似算法的方法的弱点。而MoBA在处理这些末尾位置的词元时，预测质量并没有明显下降。

研究人员还研究了MoBA对区块大小和门控策略的敏感度。有些实验发现，细化粒度（也就是用更小的区块，但选更多的区块）能让模型的效果更接近全注意力机制。就算MoBA忽略了大部分上下文，自适应门控也能找到对查询真正重要的区块。此外，还有一种“混合”模式，采取了平衡策略：一部分层继续用MoBA来提高速度，少数层则切换回全注意力模式。在进行有监督微调时，这种混合模式特别有用，因为训练目标可能会屏蔽掉输入里的某些位置。在少数上层保留全注意力，模型就能覆盖更广泛的上下文，对那些需要全局视角的任务很有帮助。

总的来说，这些研究结果表明，MoBA特别适合处理需要大量上下文的任务，像是长篇文档阅读理解、大规模代码补全，还有需要参考完整对话历史的多轮对话系统。MoBA能有效提升效率，而且对性能的影响很小，是让大语言模型在大规模应用中更高效的理想选择。

结论

块注意力混合（MoBA）为大语言模型处理长文本提供了一种更高效的方法，而且不用大幅修改Transformer架构，也不会降低性能。通过在注意力模块里融入专家混合的理念，MoBA用一种可学习的稀疏方式，让模型聚焦在超长输入的关键部分。它的设计非常灵活，尤其是能在稀疏注意力和全注意力之间无缝切换，这对正在进行的和未来的训练流程都很有吸引力。研究人员可以灵活调整修剪注意力模式的程度，要是任务需要全面覆盖信息，也可以有选择地使用全注意力模式。

虽然目前对MoBA的研究主要集中在文本领域，但它的底层机制在其他数据模态上可能也有很大潜力。只要序列长度长到会引发计算或者内存方面的问题，用MoBA这种把查询分配给不同区块“专家”的思路，就能在不影响处理全局依赖关系的前提下，有效缓解瓶颈。随着语言应用中的序列长度不断增加，像MoBA这样的方法，在提升神经语言模型的可扩展性和性价比方面，可能会发挥至关重要的作用。

链接

github链接：https://github.com/MoonshotAI/MoBA论文：  https://arxiv.org/html/2406.14909v1

本文转载自 ​​柏企科技圈​​，作者：柏企