Sparse Sinkhorn Attention

By Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, Da-Cheng Juan. 2020

399 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Model Architectures

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Model Architectures

Citations in the corpus

Scale Efficiently: Insights from Pre-training and Fine-tuning Transformers (2021, background): “…(Lan et al., 2019; Tay et al., 2019; Zhang et al., 2021), efficient attention (Tay et al., 2020c; Kitaev et al., 2020; Choromanski et al., 2020; Tay et al., 2020b; Ainslie et al., 2020; Jaegle et al., 2021), architecture search (So et al., 2019), alternative non Transformer-based architectures…”
Sparse is Enough in Scaling Transformers (2021, neutral): “Sparse attention-based approaches have made the attention layer more efﬁcient, especially for long sequences, by incorporating additional combinatorial mechanisms, as in [40], or selecting a subset of tokens this layer attends to [31, 5, 19, 37, 15, 4] or other approaches [12].”
Deformable DETR: Deformable Transformers for End-to-End Object Detection (2020, background): “Recently, many efforts have been made to address this problem (Tay et al., 2020b), which can be roughly divided into three categories in practice.”
Long Range Arena: A Benchmark for Efficient Transformers (2020, background): “For the sake of brevity, we refer interested readers to (Tay et al., 2020c) for a detailed survey of this line of research.”
Efficient Transformers: A Survey (2020, methodology): “This section introduces the Sparse Sinkhorn Transformer (Tay et al., 2020b).”
Transformers in Vision: A Survey (2021, methodology): “Numerous methods have been proposed that make special design choices to perform self-attention more ‘eiciently’, for instance employing pooling/downsampling in self-attention [61, 209, 221], local window-based attention [138, 201], axial-attention [55, 83], low-rank projection attention [192, 208, 2”
Synthesizer: Rethinking Self-Attention for Transformer Models (2020, background): “For example, recent studies have investigated improving the layer’s overall efﬁciency via sparsiﬁcation and reducing the complexity of computing the alignment matrix (Child et al., 2019; Kitaev et al., 2020; Huang et al., 2018; Tay et al., 2020; Beltagy et al., 2020).”
Multiscale Vision Transformers (2021, background): “Recent works [100, 60, 17, 94, 20, 16, 67] reduce the quadratic attention complexity to make transformers more efficient for natural language processing applications, which is complementary to our approach.”
Multi-Scale Vision Longformer: A New Vision Transformer for High-Resolution Image Encoding (2021, background): “The first is the sparse attention mechanism, including contentindependent sparsity [28, 8, 29, 14] and content-dependent sparsity [17, 32, 37, 51].”
Random Feature Attention (2021, background): “A large body of research has been devoted to improving their time and memory efﬁciency (Tay et al., 2020c).”
FNet: Mixing Tokens with Fourier Transforms (2021, methodology): “Our Linear model is similar to the MLP-Mixer (Tolstikhin et al., 2021) (for vision) and also the Random Syn-thesizer (Tay et al., 2020a), but simpliﬁes the latter model further by removing the multiple heads and softmax projections, resulting in just two matrix multiplications in the mixing…”
Luna: Linear Unified Nested Attention (2021, methodology): “…Set Transformer is designed for set-input problems, which are problems wherein the input is a set of features and the model is thereby invariant to permutation or ordering of the input features (Tay et al., 2020b), while Luna attention is used as a drop-in replacement for regular softmax attention.”
Long-Short Transformer: Efficient Transformers for Language and Vision (2021, methodology): “vi) Similarity and clustering based methods, including Reformer [31], Routing Transformer [32], and Sinkhorn Transformer [33].”
\infty-former: Infinite Memory Transformer (2021, background): “, 2020), or by learning these patterns from data (Kitaev et al., 2020; Vyas et al., 2020; Tay et al., 2020a; Roy et al., 2021; Wang et al., 2021).”
FEDformer: Frequency Enhanced Decomposed Transformer for Long-term Series Forecasting (2022, methodology): “Sinkhorn (Tay et al., 2020) employs a block sorting method to achieve quasi-global attention with only local windows.”
Learned Token Pruning for Transformers (2021, background): “These can be broadly categorized as follows: (i) efficient architecture design [8, 21, 24, 26, 27, 36, 45, 47, 50, 60]; (ii) knowledge distillation [23, 38, 42–44]; (iii) quantization [2, 4, 13, 15, 25, 40, 58, 59, 61]; and (iv) pruning.”
FiLM: Frequency improved Legendre Memory Model for Long-term Time Series Forecasting (2022, background): “(2) where the definition of A and B can be found in Voelker et al. (2019). Using Legendre Polynomials as a basis allows us to accurately approximate smooth functions, as indicated by the following theorem.”
cosFormer: Rethinking Softmax in Attention (2022, methodology): “To address this issue, numerous methods are proposed recently, such as the sparse attention matrix (Za-heer et al., 2020; Beltagy et al., 2020; Tay et al., 2020a; Kitaev et al., 2019; Child et al., 2019),low-rank representations (Wang et al., 2020) or kernel-based methods (Peng et al., 2020;…”
Attention mechanism in neural networks: where it comes and where it goes (2022, background): “It is based on the idea of differentiable sorting of internal representations within the self-attention module [109].”