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: Combiner - a new cheaper computationally attention algorithm. By introducing a new factorization structure, the design an algorithm with low complexity but capability close to full attention. Numerous experiments demonstrate the utility of the proposed approach.