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)<n>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.<n>We propose a solution that adheres to the less structure'' principle, allowing the model to determine where to attend autonomously.
• Score: 46.10222520755179
• License: http://creativecommons.org/licenses/by-nc-nd/4.0/
• 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

• MIA-Mind: A Multidimensional Interactive Attention Mechanism Based on MindSpore [0.0]
  We propose MIA-Mind, a lightweight and modular Multidimensional Interactive Attention Mechanism. MIA-Mind jointly models spatial and channel features through a unified cross-attentive fusion strategy. Experiments are conducted on three representative datasets.
  arXiv  Detail & Related papers  (2025-04-27T02:27:50Z)
• PowerAttention: Exponentially Scaling of Receptive Fields for Effective Sparse Attention [73.26995918610669]
  Large Language Models (LLMs) face efficiency bottlenecks due to the quadratic complexity of the attention mechanism when processing long contexts. We introduce PowerAttention, a novel sparse attention design that facilitates effective and complete context extension. Experiments demonstrate that PowerAttention outperforms existing static sparse attention methods by $5sim 40%$.
  arXiv  Detail & Related papers  (2025-03-05T15:24:11Z)
• 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.<n>We implement adaptive expert routing using sigmoid activation and straight-through estimators, enabling tokens to flexibly access different aspects of model knowledge.<n>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.<n>Inspired by the gating mechanism in Mixture of Experts (MoE), SeerAttention augments the conventional attention with a learnable gate.<n>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)
• How Sparse Attention Approximates Exact Attention? Your Attention is Naturally $n^C$-Sparse [9.552839922307587]
  Sparse Attention is a technique that approximates standard attention computation with sub-quadratic complexity.<n> Variations of this technique, such as pruning KV cache, sparsity-based fast attention, and Sparse Transformer, have been extensively utilized for efficient Large Language Models (LLMs) deployment.
  arXiv  Detail & Related papers  (2024-04-03T12:37:34Z)
• 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)
• Large Language Models to Enhance Bayesian Optimization [57.474613739645605]
  We present LLAMBO, a novel approach that integrates the capabilities of Large Language Models (LLM) within Bayesian optimization. At a high level, we frame the BO problem in natural language, enabling LLMs to iteratively propose and evaluate promising solutions conditioned on historical evaluations. Our findings illustrate that LLAMBO is effective at zero-shot warmstarting, and enhances surrogate modeling and candidate sampling, especially in the early stages of search when observations are sparse.
  arXiv  Detail & Related papers  (2024-02-06T11:44:06Z)
• 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.