Sparse Sinkhorn Attention

Abstract

We propose Sparse Sinkhorn Attention, a new efficient and sparse method for learning to attend. Our method is based on differentiable sorting of internal representations. Concretely, we introduce a meta sorting network that learns to generate latent permutations over sequences. Given sorted sequences, we are then able to compute quasi-global attention with only local windows, improving the memory efficiency of the attention module. To this end, we propose new algorithmic innovations such as Causal Sinkhorn Balancing and SortCut, a dynamic sequence truncation method for tailoring Sinkhorn Attention for encoding and/or decoding purposes. Via extensive experiments on algorithmic seq2seq sorting, language modeling, pixel-wise image generation, document classification and natural language inference, we demonstrate that our memory efficient Sinkhorn Attention method is competitive with vanilla attention and consistently outperforms recently proposed efficient Transformer models such as Sparse Transformers.

References in corpus (12)

• Show, Attend and Tell: Neural Image Caption Generation with Visual Attention
• Categorical Reparameterization with Gumbel-Softmax
• Effective Approaches to Attention-based Neural Machine Translation
• Generating Long Sequences with Sparse Transformers
• A large annotated corpus for learning natural language inference
• Reformer: The Efficient Transformer
• Learning Latent Permutations with Gumbel-Sinkhorn Networks
• Ranking via Sinkhorn Propagation
• Bi-Directional Block Self-Attention for Fast and Memory-Efficient Sequence Modeling
• Adaptive Input Representations for Neural Language Modeling
• Reinforced Self-Attention Network: a Hybrid of Hard and Soft Attention for Sequence Modeling
• Simple and Effective Curriculum Pointer-Generator Networks for Reading Comprehension over Long Narratives

Cited by in corpus (30)

• Transformers in Vision: A Survey
• Deformable DETR: Deformable Transformers for End-to-End Object Detection
• Synthesizer: Rethinking Self-Attention in Transformer Models
• Long Range Arena: A Benchmark for Efficient Transformers
• A Practical Survey on Faster and Lighter Transformers
• CSWin Transformer: A General Vision Transformer Backbone with Cross-Shaped Windows
• Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers
• Long-Short Transformer: Efficient Transformers for Language and Vision
• An Attention Free Transformer
• Adaptive Fourier Neural Operators: Efficient Token Mixers for Transformers
• Data Movement Is All You Need: A Case Study on Optimizing Transformers
• Summ^N: A Multi-Stage Summarization Framework for Long Input Dialogues and Documents
• Cluster-Former: Clustering-based Sparse Transformer for Long-Range Dependency Encoding
• Conformer-Kernel with Query Term Independence for Document Retrieval
• Sinkformers: Transformers with Doubly Stochastic Attention
• Convolution-enhanced Evolving Attention Networks
• Neural Proof Nets
• The NLP Task Effectiveness of Long-Range Transformers
• Transformer Acceleration with Dynamic Sparse Attention
• Efficient Attentions for Long Document Summarization
• ERNIE-Doc: A Retrospective Long-Document Modeling Transformer
• KVT: k-NN Attention for Boosting Vision Transformers
• SMYRF: Efficient Attention using Asymmetric Clustering
• Concerning Iterative Graph Normalization and Maximum Weight Independent Sets
• The Role of Global and Local Context in Named Entity Recognition
• On Learning the Transformer Kernel
• Token Pooling in Vision Transformers
• $\infty$-former: Infinite Memory Transformer
• T-EMDE: Sketching-based global similarity for cross-modal retrieval
• Pruning Self-Attention for Zero-Shot Multi-Speaker Text-to-Speech