Sparse Sinkhorn Attention
This addresses the computational bottleneck in large-scale sequence modeling for applications like language processing and image generation, representing a novel method rather than an incremental improvement.
The authors tackled the memory inefficiency of attention mechanisms by proposing Sparse Sinkhorn Attention, which uses differentiable sorting to enable quasi-global attention with local windows, and demonstrated that it is competitive with vanilla attention and outperforms other efficient Transformer models like Sparse Transformers in various tasks.
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.