Sparser is Sneller en Minder is Meer: Efficiënte Sparse Attention voor Long-Range Transformers
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
Samenvatting
Het efficiënt verwerken van lange sequenties in autoregressieve Transformers, vooral binnen een uitgebreid contextvenster, vormt aanzienlijke uitdagingen vanwege de kwadratische computationele complexiteit en de aanzienlijke geheugenvereisten voor KV (key-value) die inherent zijn aan self-attention-mechanismen. In dit werk introduceren we SPARSEK Attention, een nieuw sparse attention-mechanisme dat ontworpen is om deze computationele en geheugenbelemmeringen te overwinnen terwijl de prestaties behouden blijven. Onze aanpak integreert een scoringsnetwerk en een differentieerbare top-k-maskeroperator, SPARSEK, om een constant aantal KV-paren voor elke query te selecteren, waardoor op gradienten gebaseerde optimalisatie mogelijk wordt. Hierdoor biedt SPARSEK Attention een lineaire tijdcomplexiteit en een constant geheugengebruik tijdens generatie. Experimentele resultaten tonen aan dat SPARSEK Attention eerdere sparse attention-methoden overtreft en aanzienlijke snelheidsverbeteringen biedt tijdens zowel training als inferentie, met name bij taalmodellering en downstream-taken. Bovendien kan onze methode naadloos worden geïntegreerd in vooraf getrainde Large Language Models (LLMs) met minimale fine-tuning, wat een praktische oplossing biedt voor het effectief beheren van langeafstandsafhankelijkheden in diverse toepassingen.
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.