MoA : Mélange d'Attention Sparse pour la Compression Automatique de Modèles de Langage à Grande Échelle

MoA: Mixture of Sparse Attention for Automatic Large Language Model Compression

June 21, 2024

Auteurs: Tianyu Fu, Haofeng Huang, Xuefei Ning, Genghan Zhang, Boju Chen, Tianqi Wu, Hongyi Wang, Zixiao Huang, Shiyao Li, Shengen Yan, Guohao Dai, Huazhong Yang, Yu Wang

cs.AI

Résumé

L'attention parcimonieuse peut efficacement atténuer les importantes demandes en mémoire et en débit des modèles de langage de grande taille (LLMs) dans des contextes longs. Les méthodes existantes utilisent généralement un masque d'attention parcimonieuse uniforme, appliquant le même motif parcimonieux à travers différentes têtes d'attention et longueurs d'entrée. Cependant, cette approche uniforme ne parvient pas à capturer les divers motifs d'attention inhérents aux LLMs, ignorant leurs compromis distincts entre précision et latence. Pour relever ce défi, nous proposons le Mixture of Attention (MoA), qui adapte automatiquement des configurations d'attention parcimonieuse distinctes à différentes têtes et couches. MoA construit et explore un espace de recherche de divers motifs d'attention et leurs règles de mise à l'échelle par rapport aux longueurs des séquences d'entrée. Il profile le modèle, évalue les configurations potentielles et identifie le plan de compression d'attention parcimonieuse optimal. MoA s'adapte à des tailles d'entrée variables, révélant que certaines têtes d'attention élargissent leur focus pour s'adapter à des séquences plus longues, tandis que d'autres têtes se concentrent systématiquement sur des contextes locaux de longueur fixe. Les expériences montrent que MoA augmente la longueur de contexte effective de 3,9 fois avec la même portée d'attention moyenne, améliorant la précision de récupération de 1,5 à 7,1 fois par rapport à la base de référence à attention uniforme sur les modèles Vicuna-7B, Vicuna-13B et Llama3-8B. De plus, MoA réduit les écarts de capacités entre les modèles parcimonieux et denses, diminuant la baisse de performance relative maximale de 9%-36% à moins de 5% sur deux benchmarks de compréhension de contexte long. MoA réalise une réduction de la mémoire GPU de 1,2 à 1,4 fois et augmente le débit de décodage de 5,5 à 6,7 fois pour les modèles denses de 7B et 13B sur un seul GPU, avec un impact minimal sur les performances.

English

Sparse attention can effectively mitigate the significant memory and throughput demands of Large Language Models (LLMs) in long contexts. Existing methods typically employ a uniform sparse attention mask, applying the same sparse pattern across different attention heads and input lengths. However, this uniform approach fails to capture the diverse attention patterns inherent in LLMs, ignoring their distinct accuracy-latency trade-offs. To address this challenge, we propose the Mixture of Attention (MoA), which automatically tailors distinct sparse attention configurations to different heads and layers. MoA constructs and navigates a search space of various attention patterns and their scaling rules relative to input sequence lengths. It profiles the model, evaluates potential configurations, and pinpoints the optimal sparse attention compression plan. MoA adapts to varying input sizes, revealing that some attention heads expand their focus to accommodate longer sequences, while other heads consistently concentrate on fixed-length local contexts. Experiments show that MoA increases the effective context length by 3.9times with the same average attention span, boosting retrieval accuracy by 1.5-7.1times over the uniform-attention baseline across Vicuna-7B, Vicuna-13B, and Llama3-8B models. Moreover, MoA narrows the capability gaps between sparse and dense models, reducing the maximum relative performance drop from 9%-36% to within 5% across two long-context understanding benchmarks. MoA achieves a 1.2-1.4times GPU memory reduction and boosts decode throughput by 5.5-6.7 times for 7B and 13B dense models on a single GPU, with minimal impact on performance.