МоА: Смесь Разреженного Внимания для Автоматической Компрессии Больших Языковых Моделей

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

June 21, 2024

Авторы: 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

Аннотация

Разреженное внимание может эффективно смягчить значительные требования к памяти и пропускной способности крупных языковых моделей (LLM) в длинных контекстах. Существующие методы обычно используют равномерную маску разреженного внимания, применяя один и тот же разреженный шаблон для различных голов внимания и длин ввода. Однако такой равномерный подход не учитывает разнообразные образцы внимания, присущие LLM, игнорируя их различные компромиссы между точностью и задержкой. Для решения этой проблемы мы предлагаем Микс внимания (MoA), который автоматически настраивает различные конфигурации разреженного внимания для различных голов и слоев. MoA создает и исследует пространство поиска различных образцов внимания и их правил масштабирования относительно длины входной последовательности. Он профилирует модель, оценивает потенциальные конфигурации и определяет оптимальный план сжатия разреженного внимания. MoA адаптируется к различным размерам ввода, показывая, что некоторые головы внимания расширяют свое внимание, чтобы вместить более длинные последовательности, в то время как другие головы постоянно сосредотачиваются на фиксированных локальных контекстах. Эксперименты показывают, что MoA увеличивает эффективную длину контекста в 3,9 раза при том же среднем объеме внимания, увеличивая точность извлечения на 1,5-7,1 раза по сравнению с базовой моделью с равномерным вниманием для моделей Vicuna-7B, Vicuna-13B и Llama3-8B. Более того, MoA сокращает разрыв в возможностях между разреженными и плотными моделями, уменьшая максимальное относительное снижение производительности с 9% до 36% до значения в пределах 5% на двух длинных контекстных бенчмарках понимания. MoA достигает сокращения использования памяти GPU на 1,2-1,4 раза и увеличивает пропускную способность декодирования на 5,5-6,7 раз для плотных моделей 7B и 13B на одном GPU, с минимальным влиянием на производительность.

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.