MoBA: Смешанное внимание блоков для языковых моделей с длинным контекстом

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). Однако квадратичное увеличение вычислительной сложности, присущее традиционным механизмам внимания, создает непомерные накладные расходы. Существующие подходы либо накладывают сильно предвзятые структуры, такие как sink или window attention, которые специфичны для конкретных задач, либо радикально изменяют механизм внимания на линейные аппроксимации, чья производительность в сложных задачах рассуждения остается недостаточно изученной. В данной работе мы предлагаем решение, которое следует принципу «меньше структуры», позволяя модели самостоятельно определять, куда направлять внимание, вместо введения предопределенных предубеждений. Мы представляем Mixture of Block Attention (MoBA) — инновационный подход, который применяет принципы Mixture of Experts (MoE) к механизму внимания. Эта новая архитектура демонстрирует превосходную производительность в задачах с длинным контекстом, предлагая ключевое преимущество: возможность плавного перехода между полным и разреженным вниманием, повышая эффективность без риска снижения производительности. 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.