Sparser es más rápido y menos es más: Atención dispersa eficiente para Transformers de largo alcance

Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers

June 24, 2024

Autores: Chao Lou, Zixia Jia, Zilong Zheng, Kewei Tu

cs.AI

Resumen

La gestión eficiente de secuencias largas en Transformers autoregresivos, especialmente dentro de una ventana de contexto extendida, presenta desafíos significativos debido a la complejidad computacional cuadrática y los requisitos sustanciales de memoria KV inherentes a los mecanismos de autoatención. En este trabajo, presentamos SPARSEK Attention, un novedoso mecanismo de atención dispersa diseñado para superar estos obstáculos computacionales y de memoria mientras se mantiene el rendimiento. Nuestro enfoque integra una red de puntuación y un operador de máscara top-k diferenciable, SPARSEK, para seleccionar un número constante de pares KV para cada consulta, permitiendo así la optimización basada en gradientes. Como resultado, SPARSEK Attention ofrece una complejidad temporal lineal y un uso constante de memoria durante la generación. Los resultados experimentales revelan que SPARSEK Attention supera a los métodos de atención dispersa anteriores y proporciona mejoras significativas en velocidad tanto durante el entrenamiento como en la inferencia, particularmente en modelado de lenguaje y tareas posteriores. Además, nuestro método puede integrarse sin problemas en modelos de lenguaje preentrenados (LLMs) con un ajuste mínimo, ofreciendo una solución práctica para gestionar eficazmente dependencias de largo alcance en diversas aplicaciones.

English

Accommodating long sequences efficiently in autoregressive Transformers, especially within an extended context window, poses significant challenges due to the quadratic computational complexity and substantial KV memory requirements inherent in self-attention mechanisms. In this work, we introduce SPARSEK Attention, a novel sparse attention mechanism designed to overcome these computational and memory obstacles while maintaining performance. Our approach integrates a scoring network and a differentiable top-k mask operator, SPARSEK, to select a constant number of KV pairs for each query, thereby enabling gradient-based optimization. As a result, SPARSEK Attention offers linear time complexity and constant memory footprint during generation. Experimental results reveal that SPARSEK Attention outperforms previous sparse attention methods and provides significant speed improvements during both training and inference, particularly in language modeling and downstream tasks. Furthermore, our method can be seamlessly integrated into pre-trained Large Language Models (LLMs) with minimal fine-tuning, offering a practical solution for effectively managing long-range dependencies in diverse applications.