MoA: Mischung aus spärlicher Aufmerksamkeit zur automatischen Kompression großer Sprachmodelle

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

June 21, 2024

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

Zusammenfassung

Eine spärliche Aufmerksamkeit kann effektiv die signifikanten Speicher- und Durchsatzanforderungen von Large Language Models (LLMs) in langen Kontexten mildern. Bestehende Methoden verwenden in der Regel eine einheitliche spärliche Aufmerksamkeitsmaske, die dasselbe spärliche Muster über verschiedene Aufmerksamkeitsköpfe und Eingangslängen anwendet. Allerdings versagt dieser einheitliche Ansatz dabei, die vielfältigen Aufmerksamkeitsmuster, die in LLMs inhärent sind, zu erfassen, und ignoriert ihre unterschiedlichen Genauigkeits-Latenz-Abwägungen. Um diese Herausforderung anzugehen, schlagen wir das Mixture of Attention (MoA) vor, das automatisch unterschiedliche spärliche Aufmerksamkeitskonfigurationen für verschiedene Köpfe und Schichten maßschneidert. MoA konstruiert und navigiert durch einen Suchraum verschiedener Aufmerksamkeitsmuster und ihrer Skalierungsregeln im Verhältnis zu Eingabesequenzlängen. Es erstellt ein Profil des Modells, bewertet potenzielle Konfigurationen und ermittelt den optimalen Plan zur spärlichen Aufmerksamkeitskompression. MoA passt sich an unterschiedliche Eingabegrößen an und zeigt, dass einige Aufmerksamkeitsköpfe ihren Fokus erweitern, um längere Sequenzen zu berücksichtigen, während andere Köpfe sich konsequent auf festgelegte lokale Kontexte konzentrieren. Experimente zeigen, dass MoA die effektive Kontextlänge um das 3,9-fache erhöht, bei gleichbleibender durchschnittlicher Aufmerksamkeitsspanne die Abrufgenauigkeit um das 1,5-7,1-fache über dem einheitlichen Aufmerksamkeits-Benchmark bei den Modellen Vicuna-7B, Vicuna-13B und Llama3-8B steigert. Darüber hinaus verringert MoA die Leistungslücken zwischen spärlichen und dichten Modellen, reduziert den maximalen relativen Leistungsabfall von 9%-36% auf weniger als 5% bei zwei Benchmarks für das Verständnis von langen Kontexten. MoA erreicht eine 1,2-1,4-fache Reduzierung des GPU-Speichers und steigert den Decodierdurchsatz um das 5,5-6,7-fache für 7B- und 13B-dichte Modelle auf einer einzelnen GPU, mit minimalem Einfluss auf die Leistung.

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.