月之暗面新作 MoBA: 用于长文本大语言模型的块注意力混合模型
摘要 扩展有效上下文长度对于推动大语言模型 (LLM) 向通用人工智能 (AGI) 发展至关重要。然而,传统注意力机制中固有的计算复杂度呈二次方增长,带来了令人望而却步的开销。现有的方法要么强加了强偏置结构,例如特定于任务的 sink 或窗口注意力,这些方法通常是任务特定的,要么彻底地将注意力机制修改为线性近似,但这些方法在复杂推理任务中的性能仍未得到充分探索。 在这项工作中,我们提出了一种遵循“更少结构”原则的解决方案,该方案允许模型自主决定关注的位置,而无需引入预定义的偏见。我们引入了块注意力混合(MoBA),这是一种创新的方法,它将专家混合(MoE)的原则应用于注意力机制。 这种新颖的架构在长文本任务中表现出卓越的性能,同时提供了一个关键优势:能够在完整注意力和稀疏注意力之间无缝切换,从而提高效率,且不会有损性能。MoBA 已经部署用于支持 Kimi 的长文本请求,并展示了在大语言模型(LLM)高效注意力计算方面的显著进展。 1 引言 现有的方法通常利用预定义的结构约束来利用这种稀疏性,例如基于汇聚节点的方法 (G. Xiao et al. 2023) 或滑动窗口注意力机制 (Beltagy et al. 2020)。虽然这些方法可能有效,但它们往往是高度特定于任务的,可能会阻碍模型的整体泛化能力。 替代方法,例如稀疏注意力机制,以 Quest (Tang et al. 2024)、Minference (H. Jiang et al. 2024) 和 Re-trievalAttention (Di Liu et al. 2024) 为例,在推理时选择 Token 的子集。 尽管这些方法可以减少长序列的计算量,但它们并没有实质性地减轻长上下文模型的高昂训练成本,因此难以有效地将大语言模型扩展到数百万 Token 的上下文。 最近,另一种有希望的替代方案以线性注意力模型的形式出现,例如 Mamba (Dao and Gu 2024)、RKWV (Peng, Alcalá, et al. 2023)、Peng, Goldstein, et al. 2024 和 ReNet (Sun et al. 2023)。 这些方法用线性近似代替传统的基于 softmax 的注意力,从而减少了长序列处理的计算开销。 然而,由于线性和传统注意力之间存在显着差异,因此调整现有的 Transformer 模型通常会产生高昂的转换成本 (Mercat et al. 2024, J. Wang et al. 2024, Bick et al. 2025, M. Zhang et al. 2024) 或需要从头开始训练全新的模型 (A. Li et al. 2025)。 更重要的是,它们在复杂推理任务中的有效性证据仍然有限。 ...