Efficient Content-Based Sparse Attention with Routing Transformers

1 INTRODUCTION . Generative models of sequences have witnessed rapid progress driven by the application of attention to neural networks . In particular , Bahdanau et al . ( 2014 ) ; Cho et al . ( 2014 ) ; Vaswani et al . ( 2017 ) relied on attention to drastically improve the state-of-the art in machine translation . Subsequent research ( Radford et al. , 2018 ; Devlin et al. , 2018 ; Liu et al. , 2019 ; Yang et al. , 2019 ) demonstrated the power of self-attention in learning powerful representations of language to address several natural language processing tasks . Self-attention also brought impressive progress for generative modeling outside of language , e.g . image ( Parmar et al. , 2018 ; Menick and Kalchbrenner , 2018 ; Child et al. , 2019 ) and music generation ( Huang et al. , 2018 ; Child et al. , 2019 ) . 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 ( Vaswani et al. , 2017 ) . Self-attention ( Shaw et al. , 2018 ) is a particular case of attention ( Bahdanau et al. , 2014 ; Chorowski et al. , 2015 ; Luong et al. , 2015 ) . Self-attention is commonly used in auto-regressive 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 self-attention is quadratic in n , the length of the sequence . Specifically , for every position i ≤ n , self-attention computes weights for its whole context of length i , which induces a complexity of ∑ i≤n i = n ( n−1 ) /2 . 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 or video generation . Therefore , an important research direction is to investigate sparse and memory efficient forms of attention in order to scale to tasks with long 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 fix number of time steps in the past ( Child et al. , 2019 ) . Extensions to local attention have suggested learning the length of the temporal sparsity for each attention module in the network ( Sukhbaatar et al. , 2019 ) . 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 requires 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 , Correia et al . ( 2019 ) considers 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 L0-sparsity or sparsemax approximations ( Blondel et al. , 2019 ) . 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 word , the current time-step query is routed to a limited number of context through its cluster assignment . This strategy draws inspiration from the application of k-means clustering to Non-negative Matrix Factorization ( NMF ) ( Lee and Seung , 2001 ; Ding et al. , 2005 ; Kim and Park , 2008 ) , which is relevant to the sparsification of non-negative matrices like attention matrices . Our proposed model , Routing Transformer , combines our efficient clustered-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 Dai et al . ( 2019 ) ; Sukhbaatar et al . ( 2019 ) . We present experimental results on language modeling ( Wikitext-103 and enwik-8 ) and unconditional image generation ( ImageNet-64 ) . Routing Transformer sets new state-of-the-art while having comparable or fewer number of self-attention layers and heads , both on Wikitext-103 ( 15.8 vs 18.3 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 ) . 2 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 , Jaitly et al . ( 2015 ) 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 Chorowski et al . ( 2015 ) , while Chiu and Raffel ( 2017 ) 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 ( Gregor et al. , 2015 ; Xu et al. , 2015 ; Luong et al. , 2015 ) . Later , hierarchical attention has been simplified by Liu et al . ( 2018 ) that 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 Child et al . ( 2019 ) , which introduces bounded and strided attention , i.e . attending to a fixed context in the past at a subsampled 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 . Sukhbaatar et al . ( 2019 ) builds upon this work and shows that is 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 Dai et al . ( 2019 ) , 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 . Martins and Kreutzer ( 2017 ) ; Malaviya et al . ( 2018 ) propose to compute attention weights with variants of sparsemax . Correia et al . ( 2019 ) generalizes 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 can not be leveraged to improve space and time complexity since sparsemax/entmax formulations require instantiating the full attention matrix prior to sparsification . This is a drawback compared to temporal sparsity approaches . Our work is motivated by bridging this gap and allows for arbitrary sparsity patterns while avoiding to instantiate non-zero entries of attention matrices . Sparse Computation beyond Attention : Learning models with sparse representations/activations for saving time and computation has addressed in the past in various context . Previous work often refers to this goal as gating for conditional computation . Gating techniques relying on sampling and straight-through gradient estimators are common ( Bengio et al. , 2013 ; Eigen et al. , 2013 ; Cho and Bengio , 2014 ) . Conditional computation can also be addressed with reinforcement learning ( Denoyer and Gallinari , 2014 ; Indurthi et al. , 2019 ) . In the domain of language modeling , a related work is the sparsely gated Mixture-of-experts ( MOE ) ( Shazeer et al. , 2017 ) where sparsity is induced by experts and a trainable gating network controls the routing strategy to each sub-network . 3 SELF-ATTENTIVE AUTO-REGRESSIVE SEQUENCE MODELING . Auto-regressive sequence models decompose the probability of a sequence x = ( x1 , . . . , xn ) as p ( x ) = n∏ i=1 pθ ( xi+1|x≤i ) . ( 1 ) In neural models , the conditional distribution pθ ( xi+1|x≤i ) is modeled by a neural network with learned parameters θ and these parameters are typically learned to maximize the likelihood of the training data . In particular , Transformer architectures have shown to reach state-of-the-art accuracy in several domains , including language modeling ( Vaswani et al. , 2017 ; Radford et al. , 2018 ) , image generation ( Parmar et al. , 2018 ) and music generation ( Huang et al. , 2018 ) . 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 x = ( x1 , . . . , xn ) 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 , Q = XWQ , K = XWK , V = XWV , ( 2 ) where Q , K and V are referred to as keys , queries and values , while WQ , WK , WV are learned projection matrices . The key and the query matrices determine the n× n attention matrix A = softmax ( QK > ) , where the softmax operator over matrices denotes that the softmax function has been applied to each row . A may be interpreted as a matrix of weights in [ 0 , 1 ] where Aij denotes how much query position i at the next layer must pay attention to key position j at the previous layer . In the case of self-attention for auto-regressive models , queries attend only over keys from previous time-steps , i.e . A = softmax ( ltr ( QK > ) ) ( 3 ) where ltr denotes the lower triangular operator . Given the attention matrix A , the next layer representation X ′ is computed simply as AV . In summary , X ′i = n∑ j≤i AijVj , ( 4 ) In practice , Transformer ( Vaswani et al. , 2017 ) 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 , i.e . each layer performs multiple projections onto triplet ( queries , keys , values ) 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 non-linear transformation ( multi-layer perceptron , mlp ) followed by a residual connection and a normalization step , i.e . X ′ = layernorm ( X ′ +X ) ( 5 ) X ′′ = layernorm ( mlp ( X ′ ) +X ) , ( 6 ) where layernorm denotes the parameterized normalization step from Ba et al . ( 2016 ) . A full Transformer model is therefore a chain of attention modules ( Eq . 6 ) preceded by an embedding module ( learnable representation for symbols and their positions ) and followed by a logistic classification module ( learnable linear classifier to predict the next symbol ) . 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 .

Summary: This paper proposes a new mechanism for computing the attention scores efficiently in Transformers to avoid the quadratic computational cost in conventional Transformers (i.e., when computing the QK^T for self-attention). With a content-based "routing" technique ("group/cluster attention" may be a more accurate term to differentiate with the dynamics of Capsule Networks), the time steps of each layer form different clusters, and the self-attention mechanism is only performed within each such cluster. The authors verify their approach on very large-scale and challenging sequence benchmarks such as WikiText-103 (WT103) and enwik-8.