Sparser ist schneller und weniger ist mehr: Effiziente spärliche Aufmerksamkeit für weitreichende Transformer

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

June 24, 2024

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

cs.AI

Zusammenfassung

Die effiziente Berücksichtigung langer Sequenzen in autoregressiven Transformatoren, insbesondere innerhalb eines erweiterten Kontextfensters, stellt aufgrund der quadratischen Rechenkomplexität und des erheblichen KV-Speicherbedarfs, die in Self-Attention-Mechanismen inhärent sind, bedeutende Herausforderungen dar. In dieser Arbeit stellen wir SPARSEK Attention vor, einen neuartigen sparsamen Aufmerksamkeitsmechanismus, der entwickelt wurde, um diese Rechen- und Speicherhindernisse zu überwinden, während die Leistung beibehalten wird. Unser Ansatz integriert ein Bewertungsnetzwerk und einen differenzierbaren Top-k-Maskenoperator, SPARSEK, um eine konstante Anzahl von KV-Paaren für jede Abfrage auszuwählen, wodurch eine optimierungsbasierte Gradientenberechnung ermöglicht wird. Somit bietet SPARSEK Attention eine lineare Zeitkomplexität und einen konstanten Speicherplatzbedarf während der Generierung. Experimentelle Ergebnisse zeigen, dass SPARSEK Attention frühere sparsame Aufmerksamkeitsmethoden übertrifft und signifikante Geschwindigkeitsverbesserungen sowohl während des Trainings als auch der Inferenz bietet, insbesondere bei der Sprachmodellierung und nachgelagerten Aufgaben. Darüber hinaus kann unsere Methode nahtlos in vorab trainierte große Sprachmodelle (LLMs) mit minimalem Feintuning integriert werden und bietet somit eine praktische Lösung zur effektiven Verwaltung von Langstreckenabhängigkeiten in verschiedenen Anwendungen.

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.