MoBA: Mistura de Atenção em Blocos para LLMs de Contexto Longo

MoBA: Mixture of Block Attention for Long-Context LLMs

February 18, 2025

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

Resumo

Aumentar o comprimento efetivo do contexto é essencial para o avanço dos grandes modelos de linguagem (LLMs) em direção à inteligência artificial geral (AGI). No entanto, o aumento quadrático na complexidade computacional inerente aos mecanismos tradicionais de atenção representa uma sobrecarga proibitiva. As abordagens existentes ou impõem estruturas fortemente tendenciosas, como a atenção de sumidouro ou janela, que são específicas para tarefas, ou modificam radicalmente o mecanismo de atenção em aproximações lineares, cujo desempenho em tarefas complexas de raciocínio permanece insuficientemente explorado. Neste trabalho, propomos uma solução que adere ao princípio de "menos estrutura", permitindo que o modelo determine onde atender de forma autônoma, em vez de introduzir vieses predefinidos. Apresentamos a Mixture of Block Attention (MoBA), uma abordagem inovadora que aplica os princípios da Mixture of Experts (MoE) ao mecanismo de atenção. Essa nova arquitetura demonstra desempenho superior em tarefas de contexto longo, oferecendo uma vantagem crucial: a capacidade de transicionar de forma contínua entre atenção completa e esparsa, aumentando a eficiência sem o risco de comprometer o desempenho. A MoBA já foi implantada para atender às solicitações de contexto longo do Kimi e demonstra avanços significativos no cálculo eficiente de atenção para LLMs. Nosso código está disponível em 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.