MoBA: Mixture of Block Attention for Long-Context LLMs

• URL: http://arxiv.org/abs/2502.13189v1
• Date: Tue, 18 Feb 2025 14:06:05 GMT
• Title: MoBA: Mixture of Block Attention for Long-Context LLMs
• Authors: 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,
• Abstract summary: Scaling the effective context length is essential for advancing large language models (LLMs) toward artificial general intelligence (AGI) 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. We propose a solution that adheres to the less structure'' principle, allowing the model to determine where to attend autonomously.
• Score: 46.10222520755179
• License:
• Abstract: 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.

Related papers

• DSMoE: Matrix-Partitioned Experts with Dynamic Routing for Computation-Efficient Dense LLMs [70.91804882618243]
  This paper proposes DSMoE, a novel approach that achieves sparsification by partitioning pre-trained FFN layers into computational blocks. We implement adaptive expert routing using sigmoid activation and straight-through estimators, enabling tokens to flexibly access different aspects of model knowledge. Experiments on LLaMA models demonstrate that under equivalent computational constraints, DSMoE achieves superior performance compared to existing pruning and MoE approaches.
  arXiv  Detail & Related papers  (2025-02-18T02:37:26Z)
• Read-ME: Refactorizing LLMs as Router-Decoupled Mixture of Experts with System Co-Design [59.00758127310582]
  We propose a novel framework Read-ME that transforms pre-trained dense LLMs into smaller MoE models. Our approach employs activation sparsity to extract experts. Read-ME outperforms other popular open-source dense models of similar scales.
  arXiv  Detail & Related papers  (2024-10-24T19:48:51Z)
• SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs [10.702409298302547]
  SeerAttention learns the block-level attention sparsity from the Large Language Models itself. Inspired by the gating mechanism in Mixture of Experts (MoE), SeerAttention augments the conventional attention with a learnable gate. Our evaluation results demonstrate that SeerAttention achieves better model accuracy and lower latency for long-context pre-filling.
  arXiv  Detail & Related papers  (2024-10-17T07:07:09Z)
• FactorLLM: Factorizing Knowledge via Mixture of Experts for Large Language Models [50.331708897857574]
  We introduce FactorLLM, a novel approach that decomposes well-trained dense FFNs into sparse sub-networks without requiring any further modifications. FactorLLM achieves comparable performance to the source model securing up to 85% model performance while obtaining over a 30% increase in inference speed.
  arXiv  Detail & Related papers  (2024-08-15T16:45:16Z)
• Short-Long Convolutions Help Hardware-Efficient Linear Attention to Focus on Long Sequences [60.489682735061415]
  We propose CHELA, which replaces state space models with short-long convolutions and implements linear attention in a divide-and-conquer manner. Our experiments on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method.
  arXiv  Detail & Related papers  (2024-06-12T12:12:38Z)
• Efficient and Economic Large Language Model Inference with Attention Offloading [11.698376311689456]
  Transformer-based large language models (LLMs) exhibit impressive performance in generative tasks but introduce significant challenges in real-world serving. This mismatch arises from the autoregressive nature of LLMs, where the generation phase comprises operators with varying resource demands. To enhance the efficiency and cost-effectiveness of LLM serving, we introduce the concept of attention offloading.
  arXiv  Detail & Related papers  (2024-05-03T02:15:15Z)
• Extending Token Computation for LLM Reasoning [5.801044612920816]
  Large Language Models (LLMs) are pivotal in advancing natural language processing. LLMs often struggle with complex reasoning tasks due to inefficient attention distributions. We introduce a novel method for extending computed tokens in the Chain-of-Thought process, utilizing attention mechanism optimization.
  arXiv  Detail & Related papers  (2024-03-22T03:23:58Z)
• Pessimism meets VCG: Learning Dynamic Mechanism Design via Offline Reinforcement Learning [114.36124979578896]
  We design a dynamic mechanism using offline reinforcement learning algorithms. Our algorithm is based on the pessimism principle and only requires a mild assumption on the coverage of the offline data set.
  arXiv  Detail & Related papers  (2022-05-05T05:44:26Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.