SeerAttention: Aprendizaje de Atención Esparsa Intrínseca en sus LLMs

SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs

October 17, 2024

Autores: Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Hayden Kwok-Hay So, Ting Cao, Fan Yang, Mao Yang

cs.AI

Resumen

La atención es la piedra angular de los Modelos de Lenguaje Grandes (LLMs) modernos. Sin embargo, su complejidad cuadrática limita la eficiencia y escalabilidad de los LLMs, especialmente para aquellos con una ventana de contexto larga. Un enfoque prometedor para abordar esta limitación es aprovechar la dispersión en la atención. Sin embargo, las soluciones basadas en dispersión existentes dependen predominantemente de patrones predefinidos o heurísticas para aproximar la dispersión. Esta práctica no logra capturar completamente la naturaleza dinámica de la dispersión de la atención en tareas basadas en lenguaje. Este artículo argumenta que la dispersión de la atención debería ser aprendida en lugar de predefinida. Con este fin, diseñamos SeerAttention, un nuevo mecanismo de Atención que complementa la atención convencional con una compuerta aprendible que selecciona de manera adaptativa bloques significativos en un mapa de atención y considera que el resto de los bloques son dispersos. Esta dispersión a nivel de bloque equilibra eficazmente la precisión y la aceleración. Para permitir el aprendizaje eficiente de la red de compuertas, desarrollamos una implementación FlashAttention personalizada que extrae la verdad fundamental a nivel de bloque del mapa de atención con un mínimo costo adicional. SeerAttention no solo se aplica al post-entrenamiento, sino que también sobresale en el ajuste fino de largo contexto. Nuestros resultados muestran que en las etapas de post-entrenamiento, SeerAttention supera significativamente a los métodos de atención dispersa estáticos o basados en heurísticas de última generación, al mismo tiempo que es más versátil y flexible para adaptarse a diferentes longitudes de contexto y ratios de dispersión. Cuando se aplica al ajuste fino de largo contexto con YaRN, SeerAttention puede lograr un notable ratio de dispersión del 90% en una longitud de contexto de 32k con una pérdida de perplejidad mínima, ofreciendo una aceleración de 5.67 veces sobre FlashAttention-2.

English

Attention is the cornerstone of modern Large Language Models (LLMs). Yet its quadratic complexity limits the efficiency and scalability of LLMs, especially for those with a long-context window. A promising approach addressing this limitation is to leverage the sparsity in attention. However, existing sparsity-based solutions predominantly rely on predefined patterns or heuristics to approximate sparsity. This practice falls short to fully capture the dynamic nature of attention sparsity in language-based tasks. This paper argues that attention sparsity should be learned rather than predefined. To this end, we design SeerAttention, a new Attention mechanism that augments the conventional attention with a learnable gate that adaptively selects significant blocks in an attention map and deems the rest blocks sparse. Such block-level sparsity effectively balances accuracy and speedup. To enable efficient learning of the gating network, we develop a customized FlashAttention implementation that extracts the block-level ground truth of attention map with minimum overhead. SeerAttention not only applies to post-training, but also excels in long-context fine-tuning. Our results show that at post-training stages, SeerAttention significantly outperforms state-of-the-art static or heuristic-based sparse attention methods, while also being more versatile and flexible to adapt to varying context lengths and sparsity ratios. When applied to long-context fine-tuning with YaRN, SeerAttention can achieve a remarkable 90% sparsity ratio at a 32k context length with minimal perplexity loss, offering a 5.67x speedup over FlashAttention-2.