Più Sparso è Più Veloce e Meno è Meglio: Attenzione Sparsa Efficiente per Trasformatori a Lungo Raggio
Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers
June 24, 2024
Autori: Chao Lou, Zixia Jia, Zilong Zheng, Kewei Tu
cs.AI
Abstract
Accomodare sequenze lunghe in modo efficiente nei Transformer autoregressivi, specialmente all'interno di una finestra di contesto estesa, presenta sfide significative a causa della complessità computazionale quadratica e dei requisiti di memoria KV sostanziali intrinseci nei meccanismi di self-attention. In questo lavoro, introduciamo SPARSEK Attention, un nuovo meccanismo di attenzione sparsa progettato per superare questi ostacoli computazionali e di memoria mantenendo le prestazioni. Il nostro approccio integra una rete di scoring e un operatore di maschera top-k differenziabile, SPARSEK, per selezionare un numero costante di coppie KV per ogni query, consentendo così l'ottimizzazione basata su gradienti. Di conseguenza, SPARSEK Attention offre una complessità temporale lineare e un'impronta di memoria costante durante la generazione. I risultati sperimentali rivelano che SPARSEK Attention supera i precedenti metodi di attenzione sparsa e fornisce miglioramenti significativi in termini di velocità sia durante l'addestramento che l'inferenza, specialmente nel modellamento del linguaggio e nei task downstream. Inoltre, il nostro metodo può essere integrato senza soluzione di continuità in modelli linguistici di grandi dimensioni (LLM) pre-addestrati con un minimo fine-tuning, offrendo una soluzione pratica per gestire efficacemente le dipendenze a lungo raggio in diverse applicazioni.
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.