MoBA：kimi团队发布超强注意力方法，让大模型的上下文理解长度突破10M！

本文链接：https://blog.youkuaiyun.com/sherlockMa/article/details/145731255

1.简介

有效扩展上下文长度对于将大型语言模型（LLM）推向人工通用智能（AGI）至关重要。然而，传统注意力机制中固有的二次计算复杂度显著增加了计算开销，这成为了实现这一目标的一大障碍。

在这项工作中，作者了混合块注意力（MoBA），一种创新的方法，适用于混合专家（MoE）的注意力机制的原则。这种新颖的架构在长上下文任务上表现出了极佳性能，同时提供了一个关键优势：能够在完全注意力和稀疏注意力之间无缝转换，提高效率而不会影响性能。MoBA已经被部署来支持Kimi的长上下文请求，并在LLM的有效注意力计算方面取得了重大进展。

2.论文详解

简介

对人工通用智能（AGI）的追求推动了大型语言模型（LLM）的发展，使其达到了前所未有的规模，并有望处理模拟人类认知的复杂任务。实现AGI的一个关键能力是处理、理解和生成长序列的能力，这对于从历史数据分析到复杂推理和决策过程的广泛应用至关重要。

然而，由于标准注意力机制相关的计算复杂性呈二次增长，因此在LLM中扩展序列长度并非易事。这一挑战激发了一波旨在提高效率而不牺牲性能的研究。一个突出的方向是利用注意力分数的内在稀疏性。这种稀疏性在数学上来自softmax操作。

Softmax函数的稀疏性是指在计算过程中，对于给定的输入向量，只有少数几个元素（即概率值）会显著大于零，而其他大多数元素则接近于零。这种现象在处理长序列和高维数据时尤为明显，因为Softmax函数会将输入向量转换为一个概率分布，其中大部分概率质量集中在少数几个元素上。

为什么Softmax具有稀疏性？

指数函数的性质：Softmax函数是通过指数函数和归一化操作来计算的。对于输入向量 x，Softmax函数定义为： $Softmax(x_i)=\frac{e^{x_i}}{\sum_{j}e^{x_j}}$

其中 xi 是输入向量的第 i 个元素。指数函数 exi 会放大输入值的差异，使得较大的输入值对应的概率值显著大于较小的输入值。

归一化操作：Softmax函数的分母是所有指数值的和，这确保了输出向量是一个概率分布，即所有元素的和为1。这种归一化操作进一步加剧了概率值的集中，因为较小的指数值在归一化后会变得更小。

长尾分布：在许多实际应用中，如自然语言处理和图像识别，数据的分布往往是长尾的。这意味着只有少数几个类别或元素出现的频率很高，而大多数类别或元素出现的频率很低。Softmax函数在这种情况下会将大部分概率质量分配给那些高频的类别或元素，而低频的类别或元素则分配到接近于零的概率值。

现有的方法通常利用预定义的结构约束，如sink-based的或滑动窗口注意，以利用这种稀疏性。虽然这些方法可能是有效的，但它们往往是依赖于特定任务，可能会阻碍模型的整体通用性。或者使用一系列动态稀疏注意机制，例如Quest，Minference和RetrievalAttention，这些方法在推理时选择标记的子集。虽然这样的方法可以减少长序列的计算，但它们并没有实质上减轻长上下文模型的密集训练成本，这使得将LLM有效地扩展到数百万个令牌的上下文具有挑战性。

另一种有前途的替代方法最近以线性注意力模型的形式出现，如Mamba，RWKV和RetNet。这些方法用线性近似代替了基于规范softmax的注意力，从而减少了长序列处理的计算开销。然而，由于线性注意力和传统注意力之间的实质性差异，调整现有的Transformer模型通常会导致高转换成本或需要从头开始训练全新的模型。更重要的是，它们在复杂推理任务中的有效性的证据仍然有限。

因此，一个关键的研究问题出现了：我们如何设计一个强大的和适应性强的注意力架构，保留原来的Transformer框架，同时坚持“少结构”的原则，允许模型来确定在哪里参加，而不依赖于预定义的结果？理想情况下，这种架构将在完全注意力模式和稀疏注意力模式之间无缝转换，从而最大限度地提高与现有预训练模型的兼容性，并在不影响性能的情况下实现高效推理和加速训练。

因此，作者引入了一种新型架构，称为块注意力混合（MoBA），它通过动态选择历史片段（块）进行注意力计算来扩展Transformer模型的功能。MoBA受到混合专家（MoE）和稀疏注意力技术的启发。前一种技术主要应用于Transformer架构中的前馈网络（FFN）层，而后者已被广泛用于缩放Transformer以处理长上下文。

方法

初步准备：Transformer中的标准注意力

首先回顾一下Transformer中的注意力标准。为了简单起见，我们重新考虑单个Query令牌 $q\in \mathbb{R}^{1\times d}$ 涉及N个Key令牌和Value令牌的情况，分别表示 $K,V\in \mathbb{R}^{N\times d}$ 。标准注意力计算如下： $Attn(q,K,V)=Softmax(qK^T)V$ ，其中d表示单个注意力头部的尺寸。

MoBA架构

与每个Query令牌关注整个上下文的标准关注不同，MoBA使每个Query令牌仅关注Key和Value的子集： $MoBA(q,K,V)=Softmax(qK[I]^T)V[I]$ ，其中 $I\subseteq [N]$ 是所选Key和Value的集合。

MoBA的关键创新是块划分和选择策略。作者将长度为N的完整上下文划分为n个块，其中每个块表示后续令牌的子集。这里假设上下文长度N可被块的数量n整除。可以进一步将 $B=\frac{N}{n}$ 表示为块大小，而 $I_i=[(i-1)\times B+1,i\times B]$ 是第i个块的范围。

通过应用来自MoE的top-k门控机制，可以使每个Query能够选择性地关注来自不同块的令牌子集，而不是整个上下文： $I=\bigcup_{g_i> 0}I_i$ 为每个查询标记选择最相关的块。具体步骤是首先计算度量Query q和第i个块之间的相关性的亲和度分数si，并在所有块之间应用Top-k门控。