Attention又升级！Moonshot | 提出MoE注意力架构：MoBA，提升LLM长文本推理效率

源自: AINLPer（每日干货分享！！）
编辑: ShuYini
校稿: ShuYini
时间: 2025-3-13

更多：>>>>大模型/AIGC、学术前沿的知识分享！

引言

对于大模型来说，有效扩展上下文长度对于实现通用人工智能至关重要。我们知道传统注意力机制的计算复杂度呈二次方增长，这极大限制了模型的扩展性。而现有方法要么引入任务特定的强偏差结构，要么是将注意力机制近似修改为线性。而本文作者将MoE的思想应用到了Attention机制上，提出了MoBA方法，该方法能够无缝切换全注意力和稀疏注意力，不仅提升了LLM推理效率，且保证了模型性能。

论文：https://arxiv.org/pdf/2502.13189

代码：https://github.com/MoonshotAI/MoBA

背景介绍

人工通用智能(AGI)的追求推动了大型语言模型(LLMs)发展到前所未有的规模，有望处理模仿人类认知的复杂任务。实现AGI的一个关键能力是处理、理解和生成长序列的能力，这对于历史数据分析、复杂逻辑推理以及规划决策等相关应用来说至关重要。扩展上下文长度不仅体现在长文本提示词输入方面，也体现在对长思维链的输出方面。

然而，要扩展大模型的上下文长度并非易事。就于传统注意力机制来说，它的计算复杂度会随着序列长度的增加而呈二次增长，如果要保证原有Attention机制效果同时要保证其性能一直是学术研究的重点。有研究人员发现注意力分数存在内在稀疏性，这种稀疏性既来自数学上的计算方式，也与生物学上的脑区记忆存储机制相似。（这里其实和LoRA的原理类似）

针对该问题，目前主流的解决方法主要有三种：一是通过预定义的结构约束来利用注意力分数的稀疏性，如sink-based注意力或滑动窗口注意力，这种结构化的方法比较适用于特定任务，会影响模型的泛化能力；二是采用动态稀疏注意力机制方法，在推理时选择部分重要的信息，会大大减少推理计算量，但不会减轻训练成本；三是线性注意力模型，用线性近似替代传统的注意力计算方式，这会极大的减少计算量，但是目前学术领域还没有明确的成果来验证其复杂推理的有效性。

那么如何设计一种既保留原始Transformer架构优势，又不依赖预定义结构的适应性强的注意力机制？理想的解决方案应能在完全注意力和稀疏注意力模式之间无缝切换，最大限度地兼容现有预训练模型，同时实现高效推理和训练而不影响性能。

受MoE架构的启发，本文作者提出了MoBA架构，将MoE应用于Transformer模型的注意力机制，MoBA在长上下文任务中表现出色，且能够无缝切换全注意力和稀疏注意力，提升效率的同时不损害性能。

MoBA架构介绍

MoBA（Mixture of Block Attention）是一种新型架构，借鉴了 Mixture of Experts (MoE) 的思想，并将其应用于 Transformer 的注意力机制。MoBA 通过动态选择历史相关的 Key-Value（KV）块，实现高效长序列处理，同时兼顾计算效率和模型性能。架构图如下所示：

1、标准的Attention

在介绍MoBA之前，我们先看一下标准Attention，其中单个请求 token q 就是通过以下公式计算注意力：

$Attn(q, K, V) = Softmax(qK^T) V$

其中：

q：查询 token
K, V：键值对矩阵
Softmax(qK^T)：计算注意力权重

这种方法要求 q 关注所有 N 个 KV token，计算复杂度为 $O(N^2)$ ，对长序列不够高效。

2、MoBA动态选择块

而MoBA则是通过 动态块选择 来优化注意力计算，每个请求 token 只关注部分 K, V 块：
$MoBA(q, K, V) = Softmax(qK[I]^T) V[I]$

其中，I ⊆ [N]为所选的 KV 块索引集合

3、块划分与选择

你可能好奇，块的划分和选择是如何实现的呢？

1）划分 KV 块：将上下文 N 划分为 n 个块，每个块大小 B = N/n，块 i 的索引范围为，这里可以将序列划分为更小的块（即减小 B），从而让请求 token 在更精细的范围内选择相关信息，从而减少冗余计算，提高注意力的有效性。
$I_i = [(i-1) \times B + 1, i \times B]$
2）基于 MoE 的KV块选择：

计算请求 q 与每个块的相关性得分 $s_i$ ，这里向量计算选择的是内积的形式：
$s_i = \langle q, \text{mean pool}(K[I_i]) \rangle$
选取得分最高的 k 个块：
$g_i = \begin{cases} 1, & s_i \in \text{Topk}(\{s_j | j \in [n] \},k) \\ 0, & \text{otherwise} \end{cases}$

3）确保因果性：

仅允许 q 关注 当前块及之前的块，避免信息泄漏：
$s_i = -\infty, \quad g_i = 0, \quad \text{if } pos(q) < i \times B$
强制 q 关注所在的当前块，并对其应用因果掩码。

4、算法实现

从这里MoBA可以看作是 滑动窗口注意力（SWA）和 Attention Sink 机制的一种推广。其中在 SWA 机制中，每个查询 token 仅关注其附近的窗口，而在 Attention Sink 机制中，查询 token 主要关注初始 token 和最近的 token 组合。MoBA 通过灵活的门控机制，使模型能够动态地选择关键块，从而在不同任务场景下灵活近似 SWA 或 Attention Sink。此外，MoBA 具备更强的表达能力，不仅可以模拟已有的静态稀疏注意力方法，还能根据任务需求灵活调整关注的范围，实现更高效的信息提取和计算优化。

实验结果

架构介绍完了，让我们来看一下实验结果咋样。

如下图所示，左图在 1M 序列长度计算效率提升评估中，随着序列长度从 8K 增加到 1M，MoBA 的计算时间扩展显著优于全注意力（使用Flash Attention实现）；右图在固定稀疏比率的计算时间扩展对比中，序列长度从 8K 增加到 10M，保持 95.31% 的稀疏率（固定 64 个 MoBA 块，块大小可变，top-k=3），MoBA 的效率优势进一步凸显

为了评估MoBA的有效性，通过比较使用全注意力或MoBA训练的语言模型的验证损失，进行了缩放定律实验。左图对比全注意力和MoBA训练的语言模型，发现二者缩放趋势相似，MoBA在高达75%稀疏度下性能与全注意力相当。右图在五次实验中，MoBA在最后一块的语言模型损失上略高于全注意力，但损失差距正在逐渐缩小。这一实验表明MoBA具有长上下文的可扩展性。

在多个长上下文benchmark上，比较了MoBA与全注意力模型的性能表现。