Combiner: Full Attention Transformer with Sparse Computation Cost

√ L ) ) . Combiner is a drop-in replacement for attention layers in existing transformers and can be easily implemented in common frameworks . An experimental evaluation on both autoregressive and bidirectional sequence tasks demonstrates the effectiveness of this approach , yielding state-of-the-art results on several image and text modeling tasks . 1 Introduction . The Transformer [ 1 ] is a powerful neural network architecture that has demonstrated state-of-the-art performance in machine translation [ 2 ] and many other natural language processing ( NLP ) tasks via pretraining , using either unidirectional language modeling [ 3 ] or bidirectional language modeling [ 4–8 ] . It has also achieved excellent results in other domains like image recognition [ 9 ] , code understanding [ 10 ] , speech recognition [ 11 ] , protein [ 12 ] , music [ 13 ] and image [ 14 ] generative modeling . The core component of Transformer is the attention mechanism , which computes dependencies between all pairs of positions in a sequence . However , for a sequence of length L , the expressiveness of pairwise attention comes at a quadratic cost O ( L2 ) in both time and memory consumption . This makes the vanilla Transformer [ 1 ] prohibitive for applications that involve long sequences , including high-resolution images , protein sequences , or raw speech signals [ 15 ] , where the sequence length L is often larger than 10 , 000 [ 14 ] . Recently , there have been several attempts to scale up attention to long sequences . A popular class of methods sparsifies the attention matrix with different sparsity patterns , including local ∗indicates equal contribution . The work was completed during HR ’ s internship at Google Brain . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . window [ 16 , 17 ] , local+stride [ 14 ] , log-sparse [ 18 ] , axial [ 19 , 20 ] , or learnable patterns through hashing [ 21 ] or clustering [ 22 ] . Sparse attention enjoys sub-quadratic cost , but is lossy in capturing all-pair relationships . Generally , sparse attention requires more layers [ 14 , 20 , 23 ] to achieve full autoregressive or bidirectional dependencies ( or receptive fields [ 20 ] ) for each location in a long sequence . Alternatively , another line of research has tried to achieve scalability with an explicit low-rank assumption [ 24 , 25 ] on the attention matrix or by using explicit feature maps of some kernels [ 26 ] . However these explicit low dimensional approximations might be too restricted for the potentially full rank attention matrix , which uses exponential kernels that are effectively infinite dimensional [ 27 ] . The Performer [ 28 ] is among the first works that attempts to approximate regular full-rank attention with the random feature trick [ 29 ] . However such random-feature based approaches [ 30 ] require many more bases to better approximate the exponential kernel [ 27 ] , and empirically we found it produces inferior results in some sequence modeling tasks , such as density estimation . In this paper we propose Combiner , a drop-in replacement for the vanilla quadratic attention mechanism with sub-quadratic computation and memory cost . Combiner still achieves full attention capability within each head of Multi-Head Attention , unlike approaches that adopt sparse or low-rank approximations . As we will discuss , the standard attention computed at each location can be seen as the conditional expectation of the value embeddings at all feasible locations given the current location . Based on such an understanding , Combiner explicitly approximates the conditional distribution in through a structured factorization of the probability space . Specifically , given a location x , the probability of attending to location y can be either directly calculated via the query vector of x and key vector of y , or indirectly through a local abstraction where x first attends to the key vector that represents a group of locations containing y , and multiplying the probability of choosing y within that group . We refer to this model as Combiner since the conditional distributions in attention become a combination between several local attentions and direct attentions . This structured decomposition enables Combiner to take existing sparse attention patterns and convert them into corresponding design choices for probability factorizations that achieve full attention . As shown in Figure 1 , Combiner achieves full attention with the same asymptotic complexity as sparse variants . Combiner can be easily implemented in most existing deep learning frameworks without the need for specialized hardware implementation , and is GPU/TPU friendly . In fact , both the fixed and learnable sparse attention patterns from many existing Transformer variants [ 14 , 18 , 20 , 22 ] can be enhanced with such structured factorizations , with the same order of time or memory cost . We validate Combiner on both autoregressive and bidirectional sequence modeling tasks over a variety of domains including text and images . We show that Combiner can achieve better perplexity and accuracy when using the same transformer architectures while being much faster in terms of runtime , and achieves state of the art performance on density estimation on standard datasets CIFAR-10 ( 2.77 bits/dim ) and ImageNet-64 ( 3.42 bits/dim ) , as well as the Long-Range Arena [ 31 ] . The implementation of Combiner can be found at https : //github.com/google-research/googleresearch/tree/master/combiner . 2 Attention as Conditional Expectation . In this section , we revisit the formulation of the standard Transformer [ 1 ] from the perspective of conditional expectation , which inspires the derivation of Combiner . Without loss of generality , we use a single sequence in the self-attention scenario . Given a sequence of L embeddings X = [ x1 , x2 , . . . , xL ] , where X ∈ RL×d and each embedding xi ∈ Rd is a d-dimensional vector , the core component of Transformer is the multi-head attention , where each head h is a scaled dot-product attention : Ah ( X ) = softmax ( Qh√ d K⊤h ) Vh , { Qh = XW Q h , Kh = XW K h , Vh = XW V h } ∈ RL×d , ( 1 ) and the attention vector from each head Ah ( X ) is concatenated and projected : MultiHeadAttn ( X ) = [ A1 ( X ) , A2 ( X ) , . . . , AH ( X ) ] W o , W o ∈ RHd×d . ( 2 ) Here H is the total number of heads per Transformer layer . In this paper , we focus on how to approximate full attention within each head of multi-head attention . For ease of notation , we drop the head index h whenever possible , and use lower-case letters xi , qi , ki , vi ∈ Rd to denote rows in X , Q , K , V respectively , which corresponds to a location i in the original sequence of length L. We use [ n ] to denote the set of positive integers { 1 , 2 , . . . , n } . For a position i ∈ [ L ] , the attention formulation ( 1 ) can be viewed as conditional expectation of rows in V . Specifically , since softmax outputs a probability distribution , we can rewrite ( 1 ) as A ( xi ) = Ep ( j|i ) [ vj ] , p ( j|i ) = 1 Z ( xi ) exp ( qi√ d k⊤j ) , ( 3 ) where p ( j|i ) denotes the conditional probability at position j given the token at position i and the partition function Z ( xi ) = ∑ j∈Ωi exp ( qi√ d k⊤j ) over support Ωi . The support Ωi of p ( j|i ) defines the set of valid locations that the i-th token can attend to . For instance , the support set in autoregressive language modeling ( LM ) consists of all previous tokens , i.e. , ΩLMi = [ i ] 2 ; in masked language modeling ( MLM ) the support consists of all tokens in the sequence , i.e. , ΩMLMi = [ L ] . That is , ΩLMi and Ω MLM i represent the full attention capability respectively in the LM and MLM setting . 3 Combiner : Full Attention via Structured Conditional Expectation . The complexity of p ( j|i ) is the bottleneck of the computation for A ( xi ) . Generally , in existing sparse transformers , the support of p ( j|i ) is sparsified to reduce the computation and memory complexity , e.g. , ΩSparsei ⊊ ΩLMi for LM and Ω Sparse i ⊊ ΩMLMi for MLM , but this can lead to either reduced capacity or limited applicability . We defer detailed discussion of the full capacity of the model to Appendix A . In this section we introduce the Combiner , which achieves ΩCombineri = Ω LM i for LM and ΩCombineri = Ω MLM i for MLM , while still maintaining sub-quadratic computation and memory cost . Below we denote Ωi as the support for full attention if there is no ambiguity or need to distinguish between LM or MLM . We introduce the main design framework in Section 3.1 and possible parameterizations in Section 3.2 . Then in Section 3.3 we analyze the trade-off of Combiner . 3.1 Local Factorization for Conditional Expectation . The main idea of Combiner is to exploit a hierarchical structure for conditional probability modeling in ( 3 ) , which provides the opportunity for reducing computation complexity while maintaining the 2Following the conventional implementation , the input sequence will be “ right-shifted ” so that the position i can attent to itself in LM setting . same support . Specifically , we introduce support variables Ωri , for r = 0 , . . . , ni and i ∈ [ L ] . The support variables are disjoint , i.e. , Ωri ∩ Ωsi = ∅ , ∀r ̸= s , and ∪ ni r=0Ω r i = Ωi . Then we can factorize p ( j|i ) as p ( j|i ) = ni∑ r=0 p ( j , Ωri |i ) = ni∑ r=0 p ( j|Ωri , i ) p ( Ωri |i ) = p ( j|Ω rj i , i ) p ( Ω rj i |i ) , ( 4 ) where rj denotes the index of the support to which j belongs . The last equation arises from the fact that the Ωri are disjoint from each other ( Ω r i ∩ Ωsi = ∅ , ∀r ̸= s ) . Therefore , there is only one support , Ω rj i , containing j . The remaining terms , where j ̸∈ Ωri for r ̸= rj , are all zero since p ( j|Ωri , i ) = 0 . Furthermore , assume Ωrji is a sufficient statistic , i.e. , j and i are independent given Ω rj i , we obtain p ( j|i ) = p ( j|Ωrji ) p ( Ω rj i |i ) . ( 5 ) Given the partition { Ωri } ni r=0 , the attention form in ( 3 ) can be rewritten as A ( xi ) = Ep ( j|i ) [ vj ] = ni∑ r=0 ∑ j∈Ωri p ( j , Ωri |i ) vj ( 6 ) = ∑ j∈Ω0i p̃ ( j|i ) vj ︸ ︷︷ ︸ direct expectation + ∑ni r=1 p ( Ω r i |i ) ( ∑ j∈Ωri p ( j|Ωri ) vj ) ︸ ︷︷ ︸ local expectation , ( 7 ) where we consider direct attention in partition Ω0i and apply the local factorization ( 5 ) to the partition r = 1 , . . . , ni . Here p̃ ( j|i ) ∝ p ( j|i ) but with different normalization constants , which will be explained below . We refer to this model as Combiner since the structured attention ( 7 ) combines the direct expectation of Ω0i and multiple local expectations via p ( j|Ωri ) and p ( Ωri |i ) to form the final conditional expectation . Equivalently , we can also rewrite the structured attention ( 7 ) as A ( xi ) = ∑ j∈Ωi [ I ( j ∈ Ω0i ) p̃ ( j|i ) + ni∑ r=1 I ( j ∈ Ωri ) p ( j|Ωri ) p ( Ωri |i ) ] ︸ ︷︷ ︸ the new effective conditional probability q ( j|i ) vj , ( 8 ) where I ( · ) is a binary indicator function . After reordering , one can see from ( 8 ) that we obtain the effective conditional probability q ( j|i ) that tries to approximate the original p ( j|i ) . Each probability term depends on both current location i and other location j , and the expectation is still obtained with respect to a valid conditional probability ( non-negative and sums up to 1 over Ωi ) . Requirement for Sub-quadratic Cost . We can immediately see the benefit of this formulation from the fact that the local expectation in ( 7 ) is independent of the position i . The full dependence is achieved via the multiplier p ( Ωri |i ) where j ∈ Ωri . If we can design the local factorization such that : 1. the order of number of terms in ( 7 ) for p ( ·|i ) , ∀i ∈ [ L ] : ∑L i=1 ( ni + |Ω0i | ) is sub-quadratic ; and 2. let U = { Ωri } i∈ [ L ] , r∈ [ 1 , ni ] be the unique set of partitions used for local expectation calculation , then the order of |U| ( i.e. , the number of unique partitions in U ) is sub-quadratic ; 3. the order of total number of unique calculations of local expectation across all locations in ( 7 ) , ∑ Ω∈U |Ω| is sub-quadratic ; then one can see that the overall computation and memory cost will be sub-quadratic with full attention support ΩCombineri = Ωi , ∀i ∈ [ L ] . We will discuss in detail in Section 4 how to instantiate such a principle by drawing inspiration from existing sparse transformers , and how to convert them into a full attention model almost for free with identical asymptotic complexity . Remark ( Further Hierarchical Decomposition ) : We introduce the local decomposition with a one layer partition of support of p ( ·|i ) for simplicity . In fact , such local decompositions can be stacked further , which introduces a partition tree . Specifically , we can further partition Ωri with disjoint subsets { Ωrki } nr k=1 , and consider local decomposition p ( j , Ωri |i ) = p ( j|Ω rkj i , i ) p ( Ω rkj i |Ωri , i ) p ( Ωri |i ) , where kj is the index of sub-region which j belongs to . Thus , we obtain a hierarchical decomposition of p ( j|i ) , which can also be plugged to ( 6 ) and yield a new full attention formulation .

Summary: This paper focuses on the problem of reducing the computational cost of attention in transformers. In particular, it addresses the lack of expressiveness in existing attention approaches that leverage sparsity. To this end, the authors propose to view attention as a conditional expectation of embeddings for each token which can be approximated with a structured factorization that includes a direct expectation and an indirect/local expectation term.  Existing sparse attention patterns with sub-quadratic cost can then be used to design such factorizations which result in more expressive attention mechanisms with the same theoretical cost. The evaluation shows that this model can also reach competitive performance in several benchmarks.