MoBA: Mixture of Block Attention für Large Language Models mit langem Kontext

MoBA: Mixture of Block Attention for Long-Context LLMs

February 18, 2025

Autoren: 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

Zusammenfassung

Die Skalierung der effektiven Kontextlänge ist entscheidend, um große Sprachmodelle (LLMs) in Richtung künstlicher allgemeiner Intelligenz (AGI) voranzubringen. Allerdings stellt der quadratische Anstieg der Rechenkomplexität, der in traditionellen Aufmerksamkeitsmechanismen inhärent ist, ein prohibitives Hindernis dar. Bestehende Ansätze führen entweder stark voreingenommene Strukturen ein, wie beispielsweise Sink- oder Fensteraufmerksamkeit, die aufgaben spezifisch sind, oder modifizieren den Aufmerksamkeitsmechanismus radikal in lineare Approximationen, deren Leistung bei komplexen Denkaufgaben noch unzureichend erforscht ist. In dieser Arbeit schlagen wir eine Lösung vor, die sich an das Prinzip „weniger Struktur“ hält und es dem Modell ermöglicht, autonom zu entscheiden, wo es seine Aufmerksamkeit richtet, anstatt vordefinierte Voreingenommenheiten einzuführen. Wir stellen Mixture of Block Attention (MoBA) vor, einen innovativen Ansatz, der die Prinzipien von Mixture of Experts (MoE) auf den Aufmerksamkeitsmechanismus anwendet. Diese neuartige Architektur zeigt überlegene Leistung bei Aufgaben mit langem Kontext und bietet einen entscheidenden Vorteil: die Fähigkeit, nahtlos zwischen vollständiger und spärlicher Aufmerksamkeit zu wechseln, wodurch die Effizienz gesteigert wird, ohne das Risiko einer Leistungseinbuße einzugehen. MoBA wurde bereits eingesetzt, um die langen Kontextanfragen von Kimi zu unterstützen, und zeigt bedeutende Fortschritte in der effizienten Aufmerksamkeitsberechnung für LLMs. Unser Code ist verfügbar unter 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.