Sparser est plus rapide et moins c'est plus : une attention parcimonieuse efficace pour les Transformers à longue portée

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

June 24, 2024

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

cs.AI

Résumé

L'intégration efficace de longues séquences dans les Transformers autorégressifs, en particulier dans une fenêtre de contexte étendue, présente des défis majeurs en raison de la complexité computationnelle quadratique et des exigences substantielles en mémoire KV inhérentes aux mécanismes d'auto-attention. Dans ce travail, nous introduisons SPARSEK Attention, un nouveau mécanisme d'attention sparse conçu pour surmonter ces obstacles computationnels et de mémoire tout en maintenant les performances. Notre approche intègre un réseau de scoring et un opérateur de masque top-k différentiable, SPARSEK, pour sélectionner un nombre constant de paires KV pour chaque requête, permettant ainsi une optimisation basée sur les gradients. En conséquence, SPARSEK Attention offre une complexité temporelle linéaire et une empreinte mémoire constante lors de la génération. Les résultats expérimentaux montrent que SPARSEK Attention surpasse les méthodes d'attention sparse précédentes et apporte des améliorations significatives en termes de vitesse lors de l'entraînement et de l'inférence, en particulier dans la modélisation du langage et les tâches en aval. De plus, notre méthode peut être intégrée de manière transparente dans les modèles de langage pré-entraînés (LLMs) avec un ajustement minimal, offrant une solution pratique pour gérer efficacement les dépendances à longue portée dans diverses applications.

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.