MoBA: 長文脈LLMのためのブロック注意機構の混合

MoBA: Mixture of Block Attention for Long-Context LLMs

February 18, 2025

著者: Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, Zhiqi Huang, Huan Yuan, Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Y. Zhang, Zhilin Yang, Xinyu Zhou, Mingxing Zhang, Jiezhong Qiu

cs.AI

要旨

効果的なコンテキスト長のスケーリングは、大規模言語モデル（LLM）を人工汎用知能（AGI）に向けて進化させる上で不可欠である。しかし、従来のアテンションメカニズムに内在する計算複雑性の二次的な増加は、過大なオーバーヘッドをもたらす。既存のアプローチでは、タスク固有のシンクやウィンドウアテンションなどの強いバイアスを課すか、あるいはアテンションメカニズムを線形近似に大幅に変更するかのいずれかであり、後者の場合、複雑な推論タスクにおける性能は十分に検証されていない。 本研究では、「少ない構造」の原則に従い、モデルが自律的にどこに注意を向けるかを決定できる解決策を提案する。我々は、Mixture of Experts（MoE）の原則をアテンションメカニズムに適用した革新的なアプローチであるMixture of Block Attention（MoBA）を導入する。この新しいアーキテクチャは、長文コンテキストタスクにおいて優れた性能を示すだけでなく、完全なアテンションとスパースアテンションの間をシームレスに切り替える能力を提供し、性能を損なうリスクなく効率を向上させる。MoBAは既にKimiの長文コンテキストリクエストをサポートするために展開されており、LLMにおける効率的なアテンション計算の大幅な進展を示している。我々のコードはhttps://github.com/MoonshotAI/MoBAで公開されている。

English

Scaling the effective context length is essential for advancing large language models (LLMs) toward artificial general intelligence (AGI). However, the quadratic increase in computational complexity inherent in traditional attention mechanisms presents a prohibitive overhead. Existing approaches either impose strongly biased structures, such as sink or window attention which are task-specific, or radically modify the attention mechanism into linear approximations, whose performance in complex reasoning tasks remains inadequately explored. In this work, we propose a solution that adheres to the ``less structure'' principle, allowing the model to determine where to attend autonomously, rather than introducing predefined biases. We introduce Mixture of Block Attention (MoBA), an innovative approach that applies the principles of Mixture of Experts (MoE) to the attention mechanism. This novel architecture demonstrates superior performance on long-context tasks while offering a key advantage: the ability to seamlessly transition between full and sparse attention, enhancing efficiency without the risk of compromising performance. MoBA has already been deployed to support Kimi's long-context requests and demonstrates significant advancements in efficient attention computation for LLMs. Our code is available at https://github.com/MoonshotAI/MoBA.