MoBA: Miscela di Attenzione a Blocchi per LLM con Contesti Lunghi

MoBA: Mixture of Block Attention for Long-Context LLMs

February 18, 2025

Autori: 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

Abstract

Scalare la lunghezza effettiva del contesto è essenziale per far progredire i grandi modelli linguistici (LLM) verso l'intelligenza artificiale generale (AGI). Tuttavia, l'aumento quadratico della complessità computazionale intrinseco nei meccanismi di attenzione tradizionali rappresenta un sovraccarico proibitivo. Gli approcci esistenti impongono strutture fortemente polarizzate, come l'attenzione a finestra o a sink, che sono specifiche per determinati compiti, oppure modificano radicalmente il meccanismo di attenzione in approssimazioni lineari, le cui prestazioni in compiti di ragionamento complesso rimangono insufficientemente esplorate. In questo lavoro, proponiamo una soluzione che aderisce al principio del "minor struttura", consentendo al modello di determinare autonomamente dove focalizzare l'attenzione, piuttosto che introdurre polarizzazioni predefinite. Introduciamo la Mixture of Block Attention (MoBA), un approccio innovativo che applica i principi della Mixture of Experts (MoE) al meccanismo di attenzione. Questa architettura innovativa dimostra prestazioni superiori nei compiti a contesto lungo, offrendo un vantaggio chiave: la capacità di passare senza soluzione di continuità tra attenzione completa e sparsa, migliorando l'efficienza senza compromettere le prestazioni. MoBA è già stato implementato per supportare le richieste a contesto lungo di Kimi e dimostra progressi significativi nel calcolo efficiente dell'attenzione per i LLM. Il nostro codice è disponibile all'indirizzo 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.