Learned Sparse Attention Patterns via Differentiable Top-K: Efficient Transformer Attention with Data-Driven Sparsity

Authors: Alice Lee*, Bob Kim, Carol Zhang

Abstract

Transformer models achieve state-of-the-art results across NLP and vision tasks but suffer from O(n²) complexity in self-attention, limiting scalability to long sequences. Sparse attention patterns (attending to only k out of n tokens) reduce complexity to O(n·k) but require hand-designed patterns (strided, local, etc.). This work proposes learned sparse attention using differentiable top-k selection, where the model learns which tokens to attend to during training. We implement a differentiable approximation of top-k via Gumbel-softmax relaxation with straight-through estimators, enabling end-to-end learning of sparse patterns. Our method learns attention sparsity patterns that adapt to each input and layer, capturing task-specific dependencies (e.g., long-range connections for language understanding, local patterns for vision). Experiments on BERT-scale models show that learned sparsity achieves 40-60% reduction in attention FLOPs while maintaining <1% accuracy loss on GLUE, SuperGLUE, and SQuAD. Learned patterns are more efficient than hand-designed baselines: strided attention (40% FLOPs reduction), local attention (50% reduction), and fixed random patterns (45% reduction). Learned sparsity achieves 1.3-1.5x speedup on inference hardware (NVIDIA A100). Notably, learned patterns transfer across similar tasks (e.g., pretrained patterns on MNLI transfer to RTE with 90% efficiency). Analysis reveals that learned patterns exhibit interpretable structure: early layers learn local patterns (attending to adjacent tokens), middle layers learn mixed patterns with long-range jumps, and late layers focus on special tokens. The framework generalizes to vision transformers, achieving 35-50% FLOPs reduction on ImageNet-1K while maintaining accuracy. Our approach is compatible with existing efficient techniques like knowledge distillation and quantization, enabling further speedups when combined. This work demonstrates that learned, task-aware sparse attention is both efficient and effective, providing a principled alternative to hand-designed patterns.