MoBA: Mezcla de Atención por Bloques para LLMs de Contexto Largo

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

Resumen

Escalar la longitud efectiva del contexto es esencial para avanzar hacia la inteligencia general artificial (AGI) en los modelos de lenguaje grandes (LLMs). Sin embargo, el aumento cuadrático en la complejidad computacional inherente a los mecanismos de atención tradicionales representa una sobrecarga prohibitiva. Los enfoques existentes imponen estructuras fuertemente sesgadas, como la atención de sumidero o ventana, que son específicas para tareas, o modifican radicalmente el mecanismo de atención en aproximaciones lineales, cuyo rendimiento en tareas de razonamiento complejo sigue siendo insuficientemente explorado. En este trabajo, proponemos una solución que se adhiere al principio de "menos estructura", permitiendo que el modelo determine autónomamente dónde atender, en lugar de introducir sesgos predefinidos. Introducimos Mixture of Block Attention (MoBA), un enfoque innovador que aplica los principios de Mixture of Experts (MoE) al mecanismo de atención. Esta arquitectura novedosa demuestra un rendimiento superior en tareas de contexto largo, ofreciendo una ventaja clave: la capacidad de transicionar sin problemas entre atención completa y dispersa, mejorando la eficiencia sin comprometer el rendimiento. MoBA ya se ha implementado para gestionar las solicitudes de contexto largo de Kimi y muestra avances significativos en el cálculo eficiente de la atención para LLMs. Nuestro código está disponible en 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.