MoBA : Mélange d'Attention par Blocs pour les LLMs à Contexte Long

MoBA: Mixture of Block Attention for Long-Context LLMs

February 18, 2025

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

Résumé

L'augmentation de la longueur effective de contexte est essentielle pour faire progresser les grands modèles de langage (LLMs) vers l'intelligence artificielle générale (AGI). Cependant, l'augmentation quadratique de la complexité computationnelle inhérente aux mécanismes d'attention traditionnels représente un surcoût prohibitif. Les approches existantes imposent soit des structures fortement biaisées, comme l'attention par puits ou par fenêtre, qui sont spécifiques à des tâches, soit modifient radicalement le mécanisme d'attention en des approximations linéaires, dont les performances dans des tâches de raisonnement complexe restent insuffisamment explorées. Dans ce travail, nous proposons une solution qui adhère au principe de « moins de structure », permettant au modèle de déterminer de manière autonome où porter son attention, plutôt que d'introduire des biais prédéfinis. Nous introduisons le Mixture of Block Attention (MoBA), une approche innovante qui applique les principes du Mixture of Experts (MoE) au mécanisme d'attention. Cette nouvelle architecture démontre des performances supérieures sur des tâches à contexte long tout en offrant un avantage clé : la capacité de passer de manière fluide entre une attention complète et une attention parcimonieuse, améliorant ainsi l'efficacité sans risquer de compromettre les performances. MoBA a déjà été déployé pour prendre en charge les requêtes à contexte long de Kimi et montre des avancées significatives dans le calcul efficace de l'attention pour les LLMs. Notre code est disponible à l'adresse 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.