Efficient Content-Based Sparse Attention with Routing Transformers (2021)

Abstract: Self-attention has recently been adopted for a wide range of sequence modeling problems. Despite its effectiveness, self-attention suffers from quadratic computation and memory requirements with respect to sequence length. Successful approaches to reduce this complexity focused on attending to local sliding windows or a small set of locations independent of content. Our work proposes to learn dynamic sparse attention patterns that avoid allocating computation and memory to attend to content unrelated to the query of interest. This work builds upon two lines of research: It combines the modeling flexibility of prior work on content-based sparse attention with the efficiency gains from approaches based on local, temporal sparse attention. Our model, the Routing Transformer, endows self-attention with a sparse routing module based on online k-means while reducing the overall complexity of attention to O(n1.5d) from O(n2d) for sequence length n and hidden dimension d. We show that our model outperforms comparable sparse attention models on language modeling on Wikitext-103 (15.8 vs 18.3 perplexity), as well as on image generation on ImageNet-64 (3.43 vs 3.44 bits/dim) while using fewer self-attention layers. Additionally, we set a new state-of-the-art on the newly released PG-19 data-set, obtaining a test perplexity of 33.2 with a 22 layer Routing Transformer model trained on sequences of length 8192. We open-source the code for Routing Transformer in Tensorflow.1

## introduction generative models of sequences have witnessed rapid progress driven by the application of atten-tion to neural networks. in particular, @xcite , @xcite , and @xcite relied on attention to drastically improve the state-of-the art in machine translation. subsequent research @xcite @xcite demonstrated the power of self-attention in learning powerful representations of language to address several natural language processing tasks. selfattention also brought impressive progress for generative modeling outside of language, for example, image @xcite @xcite and music generation @xcite . self-attention operates over sequences in a step-wise manner: at every time-step, attention assigns an attention weight to each previous input element (representation of past time-steps) and uses these weights to compute the representation of the current time-step as a weighted sum of the past input elements @xcite . selfattention @xcite is a particular case of attention @xcite @xcite . self-attention is commonly used in autoregressive generative models. these models generate observations step-by-step, modeling the probability of the next symbol given the previously generated ones. at every time step, self-attentive generative models can directly focus on any part of the previous context. in contrast, recurrent neural networks (rnns) and convolutional neural networks (cnns) have direct interactions with only a local neighborhood of context around the current time step. this advantage, however, comes at a price: unlike recurrent networks or convolution networks, the time and space complexity of selfattention is quadratic in n, the length of the sequence. specifically, for every position @xmath0, self-attention computes weights for its whole context of length i, which induces a complexity of @xmath1. this makes it difficult to scale attention-based models to modeling long sequences. however, long sequences are the norm in many domains, including music, image, speech, video generation, and document-level machine translation. therefore, an important research direction is to investigate sparse and memory efficient forms of attention in order to scale to tasks with large sequence lengths. previous work has proposed data independent or fixed sparsity patterns bounding temporal dependencies, such as local or strided attention. at each time step, the model attends only to a fixed number of time steps in the past @xcite . extensions to local attention have suggested learning the length of the temporal sparsity for each attention module in the network @xcite . these strategies draw their inspiration from rnns and cnns and bound their complexity by attending only to representations summarizing a local neighborhood of the current time step. their attention matrices (matrices containing the attention weights for every pair of previous, current time-step) are natively sparse and require instantiating only non-zero entries. while these approaches have achieved good results, fixing the sparsity pattern of a content based mechanism such as self-attention can limit its ability to pool in information from large contexts. as an alternative to local attention, @xcite consider content-based sparsity, an approach allowing for arbitrary sparsity patterns. this formulation, however, does require instantiating a full dense attention matrix prior to sparsification through variants of l 0 -sparsity or sparsemax approximations @xcite . the present work builds upon these two lines of research and proposes to retain the modeling flexibility of content-based sparse attention while leveraging the efficiency of natively sparse attention matrices. our formulation avoids sparsemax variants and relies on clustering of attention instead. each attention module considers a clustering of the space: the current time-step only attends to context belonging to the same cluster. in other words, the current time-step query is routed to a limited number of context elements through its cluster assignment. this strategy draws inspiration from the application of spherical kmeans clustering to the maximum inner product search (mips) problem. our proposed model, routing transformer, combines our efficient clustering-based sparse attention with classical local attention to reach excellent performance both for language and image generation. these results are obtained without the need to maintain attention matrices larger than batch length which is the case with the segment level recurrence mechanism used in @xcite and @xcite . we present experimental results on language modeling (enwik-8, wikitext-103, and pg-19) and unconditional image generation (cifar-10 and imagenet-64). routing transformer sets new state-of-the-art, while having comparable or fewer number of self-attention layers and heads, on wikitext-103 (15.8 vs 18.3 perplexity), pg-19 (33.2 vs 33.6 perplexity), and on imagenet-64 (3.43 vs 3.44 bits/dim). we also report competitive results on enwik-8 (0.99 vs 0.98 perplexity) and present ablations on cifar-10. ## related work attention with temporal sparsity: research on efficient attention neural models parallels the advent of attention-based architectures. in the context of speech recognition, @xcite proposed the neural transducer, which segments sequences in non-overlapping chunks and attention is performed in each chunk independently. limiting attention to a fixed temporal context around the current prediction has also been explored in @xcite , while @xcite dynamically segment the sequence into variable sized-chunks. hierarchical attention strategies have also been explored: the model first considers which part of the inputs should be attended to before computing full attention in a contiguous neighborhood of the selected area @xcite @xcite . later, hierarchical attention, simplified by @xcite , alternates coarse layers (attending to the whole sequence at a lower temporal resolution) with local layers (attending to a neighborhood of the current prediction). this alternating strategy is also employed by @xcite , who introduce bounded and strided attention, namely, attending to a fixed context in the past at a sub-sampled temporal resolution. this work formalizes such a strategy using a sparse attention formalism, showing how it relates to full attention with a specific sparsity pattern in the attention matrix. it shows that sparse attention is sufficient to get state-of-the-art results in modeling long sequences over language modeling, image generation and music generation. @xcite build upon this work and show that it is possible to obtain further sparsity by letting the model learn the length of the temporal context for each attention module. this work also makes use of the attention cache introduced in @xcite , a memory mechanism to train models over temporal contexts which extend beyond the length of the training batches. attention with content-based sparsity: the above work mainly relies on two efficient ideas: attending to less elements by only considering a fixed bounded local context in the past, and attending to less elements by decreasing the temporal resolution of context. these ideas do not allow arbitrary sparsity patterns in attention matrices. content-based sparse attention has been introduced to allow for richer patterns and more expressive models. @xcite and @xcite ) propose to compute attention weights with variants of sparsemax. @xcite generalize this approach to every layer in a transformer using entmax which allows for more efficient inference. this line of work allows for learning arbitrary sparsity attention patterns from data, based on the content of the current query and past context. however, sparsity here cannot be leveraged to improve space and time complexity because sparsemax/entmax formulations require instantiating the full attention matrix prior to sparsification. this is a drawback compared with temporal sparsity approaches. our work is motivated by bridging this gap and allows for arbitrary sparsity patterns while avoiding having to instantiate non-zero entries of attention matrices. contemporaneous to our work, @xcite proposed to use locality sensitive hashing (lsh) using random hyperplanes to infer content based sparsity patterns for attention: tokens that fall into the same hash bucket, get to attend to each other. while similar in spirit to our approach, the approach of @xcite keeps the randomly initialized hyperplanes fixed throughout, while we use mini-batch spherical kmeans to learn the space-partitioning centroids. the motivation in both approaches is to approx-imate mips in the context of dot product attention, for which both lsh and spherical kmeans have been used in literature. however, typically spherical k-means is known to outperform lsh for mips (see, e.g., @xcite . this is borne out in the common task of imagenet-64 generation, where reformer gets around 3.65 bits/dim (figure 3 ), while the routing transformer gets 3.43 bits/dim (see sparse computation beyond attention: learning models with sparse representations/ activations for saving time and computation has been addressed in the past in various contexts. previous work often refers to this goal as gating for conditional computation. gating techniques relying on sampling and straight-through gradient estimators are common @xcite @xcite . conditional computation can also be addressed with reinforcement learning @xcite . memory augmented neural networks with sparse reads and writes have also been proposed in @xcite as a way to scale neural turing machines @xcite . in the domain of language modeling, a related work is the sparsely gated mixture-of-experts (moe) @xcite , where sparsity is induced by experts and a trainable gating network controls the routing strategy to each sub-network. another related work is @xcite who use product quantization based key-value lookups to replace the feed forward network in the transformer. our work differs from theirs in that we make use of dynamic key-value pairs to infer sparsity patterns, while their key-value pairs are the same across examples. ## self-attentive auto-regressive sequence modeling in neural models, the conditional distribution p θ (x i |@xmath2. in particular, transformer architectures have shown to reach state-ofthe-art accuracy in several domains, including language modeling @xcite , image generation @xcite , and music generation @xcite . transformer models compose a series of attention modules. each module refines the input representation by taking a weighted average of the representations from the previous modules. for every module, the input representation is a sequence of n vectors @xmath3, . . . , x n ) from a continuous space of dimension d. thus one may actually treat the input sequence as a n × d matrix x. a self-attention layer operates on this representation. it first applies three linear projections, the key and the query matrices determine the n × n attention matrix @xmath4, where the softmax operator over matrices denotes that the softmax function has been applied to each row. in the case of self-attention for autoregressive models, queries attend only over keys from previous time-steps, that is, in practice, transformer @xcite adds several extensions to this basic self-attention mechanism. in particular, the result x of performing self-attention is scaled by 1/ √ d. moreover, each layer relies on multiple attention heads, that is, each layer performs multiple projections onto triplet @xcite and attention is performed for each head. the attention results from all heads are then concatenated. this strategy allows each head to specialize on different aspects of the input sequence. in addition, transformer further processes the result of attention through a learnable nonlinear transformation (multilayer perceptron, mlp) followed by a residual connection and a normalization step, namely, our work is interested in the application of the transformer to long sequences, a challenging problem since space and time complexity of attention is quadratic in sequence length n. we describe various approaches to sparse attention including ours in the next section. ## efficient content-dependent sparse attention attention-based models can be problematic for long sequences. for a sequence of length n, the full attention matrix a, as introduced in section 3, is n × n-dimensional and can be prohibitive to instantiate. this motivates sparse attention models, that is, models relying on attention matrices which have a majority of zero entries. for each query, a sparse attention model defines a set of keys that can be attended to. in the following, we introduce the set s i as the set of key positions that the query at position i can attend to, namely, the set of all such key positions defines a sparsity pattern @xmath5. for example, classical causal self attention can attend to every key prior to the current query, which translates to most previous work on attention sparsity defined such sets purely based on positions, independently of actual query and key vectors. for example, local attention @xcite considers attending only to a klong time window prior to the current query, the approach of @xcite is also a variant of local attention where the cardinality of |s i | is learned from data with an l 1 penalty to trade-off sparsity with modeling accuracy. these local attention sparsity variants are effective in practice since correlation between observations naturally decrease with time for many problems. in our experiments, we actually find that local attention is a surprisingly strong baseline in both image generation and language modeling: for example, a scaled up imagetransformer @xcite gets 3.48 bits/dim compared to the 3.44 bits/dim reported in @xcite . similarly, scaled up versions of transformer with local attention and the relative positional encoding scheme of @xcite are able to get 19.8 perplexity on wikitext-103, 1.10 bits per byte on enwik-8, and 39.3 on pg-19, while transformer-xl @xcite gets 18.3, 0.99, and 36.3 respectively. from an efficiency perspective, local attention is also interesting because sparsity patterns are regular, contiguous in memory, and known in advance. in this work, however, we are interested in a more generic formulation of attention sparsity and would like the sparsity pattern to be informed by the data, namely, @xmath6. this approach has several modeling advantages: it can accommodate data without a clear ordering over observations. for temporal data, it can also discover patterns with greater sparsity if some types of queries have a longer lasting effect on future observations than others. content-based sparse attention should, however, be carefully implemented if we need to avoid instantiating full attention matrices at any point in time. for instance, @xcite infer sparsity from data but their formulation instantiates a full attention matrix before finding its sparse counterpart. the next section explains how a natively sparse approach can actually be devised inspired by the mips problem. ## experiments we evaluate our sparse attention model on various generative modeling tasks including text and image generation. the following sections report our results on cifar-10, wikitext-103 @xcite , enwik-8 @xcite , imagenet-64, as well as pg-19 @xcite . we find that a scaled up version of local attention is a surprisingly strong baseline and that our routing transformer outperforms transformer-xl @xcite and the sparse transformer model of @xcite on all tasks. on the recently released pg-19 data-set, we find that local attention again is a strong baseline, with a slightly worse performance compared to transformer-xl @xcite . we also find that the routing transformer model outperforms both transformer-xl @xcite and compressive transformer @xcite , setting a new state-of-the-art result. in all our models except the one used for pg-19, we allocate half the heads to do local attention and the other half to route attention as in equation 8 . for all our experiments except for pg-19, we use the adam optimizer (kingma and ba, 2015) with learning rate 2 × 10 -4 with β 1 = 0.9 and β 2 = 0.98 following the learning rate schedule described in @xcite . we train all models on 128 tpuv3 cores. the setup used for pg-19 is described in section 5.5. ## analysis 6.1 local vs global as reported in section 5, a scaled up version of local attention is a strong baseline for efficient attention over long sequences. from ## conclusion transformer models constitute the state-of-theart in auto-regressive generative models for sequential data. their space-time complexity is, however, quadratic in sequence length, due to their attention modules. our work proposes a sparse attention model, the routing transformer. it relies on content-based sparse attention motivated by non-negative matrix factorization. compared with local attention models, it does not require fixed attention patterns but enjoys similar spacetime complexity. in contrast with prior work on content-based sparse attention, it does not require computing a full attention matrix but still selects sparsity patterns based on content similarity. our experiments over text and image generation draw two main conclusions. first, we show that a scaled up version of local attention establishes a strong baseline on modern benchmark, even compared to recent state-of-the-art models. second, we show that the routing transformer redefines the state-of-the-art in large long sequence benchmarks of wikitext-103, pg-19 and imagenet-64, while being very close to do so on enwik-8 as well. our analysis also shows that routing attention modules offer complementary attention patterns when compared to local attention. overall, our work contributes an efficient attention mechanism that applies to the modeling of long sequences and redefines the state of the art for auto-regressive generative modeling. our approach could prove useful in domains where the inputs are naturally sparse, such as 3d point clouds, social networks, or protein interactions.