MoA: Mezcla de Atención Dispersa para la Compresión Automática de Modelos de Lenguaje a Gran Escala

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

June 21, 2024

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

Resumen

La atención dispersa puede mitigar eficazmente las demandas significativas de memoria y rendimiento de los Modelos de Lenguaje de Gran Escala (LLMs) en contextos largos. Los métodos existentes suelen emplear una máscara de atención dispersa uniforme, aplicando el mismo patrón disperso en diferentes cabezas de atención y longitudes de entrada. Sin embargo, este enfoque uniforme no logra capturar los diversos patrones de atención inherentes a los LLMs, ignorando sus distintos equilibrios entre precisión y latencia. Para abordar este desafío, proponemos la Mezcla de Atención (MoA, por sus siglas en inglés), que adapta automáticamente configuraciones de atención dispersa distintas a diferentes cabezas y capas. MoA construye y navega un espacio de búsqueda de varios patrones de atención y sus reglas de escalado en relación con las longitudes de las secuencias de entrada. Perfila el modelo, evalúa configuraciones potenciales y determina el plan óptimo de compresión de atención dispersa. MoA se adapta a tamaños de entrada variables, revelando que algunas cabezas de atención amplían su enfoque para acomodar secuencias más largas, mientras que otras cabezas se concentran consistentemente en contextos locales de longitud fija. Los experimentos muestran que MoA aumenta la longitud efectiva del contexto en 3.9 veces con el mismo alcance promedio de atención, mejorando la precisión de recuperación en 1.5-7.1 veces en comparación con la línea base de atención uniforme en los modelos Vicuna-7B, Vicuna-13B y Llama3-8B. Además, MoA reduce las brechas de capacidad entre los modelos dispersos y densos, disminuyendo la caída máxima de rendimiento relativo del 9%-36% a menos del 5% en dos puntos de referencia de comprensión de contexto largo. MoA logra una reducción de memoria GPU de 1.2-1.4 veces y aumenta el rendimiento de decodificación en 5.5-6.7 veces para modelos densos de 7B y 13B en una sola GPU, con un impacto mínimo en el rendimiento.

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.