MoBA: Mixture of Block Attention for Long-Context LLMs

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.

한 줄 요약

MoBA는 MoE 원리를 기반으로 한 Block Attention 메커니즘으로, LLM의 긴 문맥 처리 효율성을 향상시킨다.

핵심 기여도

MoE 원리를 **Attention 메커니즘**에 적용한 **Mixture of Block Attention (MoBA)** 제안.
**전체 Attention**과 **스파스 Attention** 간의 **무단계 전환**을 가능하게 함.
**32K 컨텍스트 길이**에서 **75%의 Attention 스파시티**를 유지하면서도 성능 저하 없이 처리.
Kimi 모델에 이미 **실제 배포**되어 긴 문맥 요청 처리에 활용됨.

핵심 아이디어

기존 Attention 메커니즘은 긴 문맥 처리 시 계산 복잡도가 제곱 비례로 증가하여 비효율적이다. MoBA는 이 문제를 해결하기 위해 **Mixture of Experts (MoE)**의 아이디어를 Attention에 적용한다. 즉, 입력 문맥을 **블록 단위로 분할**하고, 각 **쿼리 토큰**에 대해 **가장 관련성 높은 KV 블록**을 선택적으로 라우팅한다. 이는 **스파스 Attention**을 유지하면서도, 모델이 **자율적으로 주목할 위치를 결정**할 수 있도록 해준다.

기존의 **Sink Attention**이나 **Sliding Window Attention**은 특정 작업에 최적화된 구조를 강제하기 때문에 일반화 능력이 떨어진다. MoBA는 이러한 **사전 정의된 편향 구조**를 최소화하고, **동적 블록 선택**을 통해 **유연성과 효율성을 동시에 달성**한다.

기술적 접근법

**Block Partitioning**: 32K 길이의 컨텍스트를 **8~128개의 블록**으로 분할.
**Gating Mechanism**: 각 쿼리 토큰에 대해 **가장 관련성 높은 KV 블록**을 선택.
**Top-k Selection**: 블록 수와 선택 수를 조정하여 **75%의 스파시티** 유지.
**Hyperparameter 조정**: 블록 수 (8, 16, 32, 64, 128)와 선택 수 (2, 4, 8, 16, 32)를 실험적으로 조정.
**MoE 기반 Attention**: FFN 레이어에서 사용된 MoE를 **Attention 메커니즘**으로 확장.

주요 결과

**32K 길이의 컨텍스트**에서 **75%의 스파시티**를 유지하면서도, **성능 저하 없이 처리**.
**가장 세분화된 블록 설정**(128개 블록 중 32개 선택)에서 **가장 높은 성능**을 보임.
**가장 거친 블록 설정**(8개 블록 중 2개 선택) 대비 **1e-2의 성능 차이** 발생.
**Kimi 모델**에 이미 **배포**되어 긴 문맥 요청 처리에 활용됨.

의의 및 한계

MoBA는 기존의 **선형 Attention**나 **정적 스파스 Attention**과 달리, **동적 블록 선택**을 통해 **유연성과 효율성을 동시에 달성**한다. 이는 LLM이 **복잡한 추론 작업**에서도 **성능을 유지하면서 긴 문맥을 처리**할 수 있도록 한다. 또한, 기존 모델에 **무단계 전환**이 가능하므로, **지속적인 사전 학습**에 유리하다.

하지만, **블록 세분화 수준**에 따라 성능이 크게 달라지므로, **최적의 블록 크기 결정**이 필요하다. 또한, **모든 작업에 동일하게 적용 가능하다는 보장**은 없으며, **추가 실험과 최적화**가 필요하다.

실용적 활용

MoBA는 **긴 문맥 처리가 필수적인 대형 언어 모델**(예: Kimi)에 적용 가능하며, **복잡한 추론 작업**(예: CoT, 문서 분석)에서 효율성을 높일 수 있다. 또한, **기존 Transformer 기반 모델**에 **무단계 통합**이 가능하므로, **산업 및 연구 분야에서의 확장성**이 높다.