Combiner: Full Attention Transformer with Sparse Computation Cost
Authors: Hongyu Ren
Hanjun Dai
Zihang Dai
Mengjiao Yang
Jure Leskovec
Dale Schuurmans
Bo Dai
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Transformer
long sequence modeling
conditional expectation factorization
full attention

Abstract: Transformers provide a class of expressive architectures that are extremely effective for sequence modeling. However, the key limitation of transformers is their quadratic memory and time complexity $\mathcal{O}(L^2)$ with respect to the sequence length in attention layers, which restricts application in extremely long sequences. Most existing approaches leverage sparsity or low-rank assumptions in the attention matrix to reduce cost, but sacrifice expressiveness. Instead, we propose Combiner, which provides full attention capability in each attention head while maintaining low computation and memory complexity. The key idea is to treat the self-attention mechanism as a conditional expectation over embeddings at each location, and approximate the conditional distribution with a structured factorization. Each location can attend to all other locations, either via direct attention, or through indirect attention to abstractions, which are again conditional expectations of embeddings from corresponding local regions. We show that most sparse attention patterns used in existing sparse transformers are able to inspire the design of such factorization for full attention, resulting in the same sub-quadratic cost ($\mathcal{O}(L\log(L))$ or $\mathcal{O}(L\sqrt{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.

Combiner: Full Attention Transformer with Sparse Computation Cost
"Hongyu Ren', 'Hanjun Dai "Zihang Dai
Mengjiao Yang". Jure Leskovec', Dale Schuurmans Bo Dai
Stanford University, (hyren.jure) @cs stanfordddd Google Research, Brain Team, http:///hhannnddeeee...cuurrnnn (eeooooddd@ao@ooooeoeoeee @google.com
University of Alberta
Abstract
Transformers provide a class of expressive architectures that are extremely effee- tive for sequence modeling. However, the key limitation of transformers is their quadratic memory and time complexity O(L) with respect to the sequence length in attention layers, which restricts application in extremely long sequences. Most existing approaches leverage sparsity or low-rank assumptions in the attention matrix reduce cost, but sacrifice expressiveness. Instead, we propose Combiner, which provides full attention capability in each attention head while maintaining low computation and memory complexity. The key idea is to treat the self-attention mechanism as a conditional expectation over embeddings at each location, and approximate the conditional distribution with structured factorization. Each loca- tion can attend to all other locations, either via direct attention, or through indirect attention abstractions, which again conditional expectations of embeddings from corresponding local regions. We show that most sparse attention patterns used in existing sparse transformers are able inspire the design of such factor- ization for full attention, resulting in the same sub-quadratic cost (O log(L) or (LVTT). Combiner is a drop-in replacement for attention layers in existing trans- experimental
formers and can be easily implemented common frameworks. evaluation on both autoregressive and bidirectional sequence tasks demonstrates the effectiveness of this approach, yielding sfate-of-the-art results on several image and text modeling tasks.
Introduction
The Transformer [1] 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 model- ing [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 mod- eling. The core component of Transformer is the attention mechanism, which computes dependencies between all pairs of positions
sequence. However, for a sequence length L, the expressiveness of pairwise attention comes at quadratic cost O(L²) in both time and memory consumption. This makes the vanilla Transformer [l] prohibitive for applications that involve long sequences, including high-ress reghuresoltiin images, protein sequences, raw speech signals [15], where the sequence length L 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 completed during HR's intersship Google Brain.
35th Conference on Neural Information Processing Systems (NcurIPS 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 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 achieve scalability with an explicit low-rank assumption [24, 25J 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 kemnlss 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-featie
based approaches [30] require
many more bases to better approximate the exponential kernel [27], and empirically we found produces inferior results in some sequence modeling tasks, such as density estimation.
In this paper we propose Combiner, drop-in replacement for the vanilla quadratic attention mech- anism 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 æ, the probability of attending to location can be either directly calculated via the query vector of x and key vector of y. or indirectly through a local abstraction where first attends the key vector that represents 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, Com- biner 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/TPÚ 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 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), well as the Long- Range Arena [31]. The implementation of Combiner can be found at http://hhwwb.ooooooogoooeeeeaaae ressarcc/ttee/mmattrrccommbbe
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 Combiner.
Without loss of generality, we use a single sequence in the self-attention scenario. Given a sequence of embeddings X [F1,.22.... .2...... 2 L1,, where X Lxd and each embedding 2, E Rd is a
d- d-dimensional vector, the core component of Transformer is the multi-head attention, where each head h is a scaled dot-product attention:
A (X) softmax
Qn
Vn, {en = XW K, = XWA, V, = XWX E Lxd
(1)
and the attention vector from each head Ax(X) is concatenated and projected:
MultiHeadAttm Ax)),
A (X)] W", w E l Haxd
(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 i whenever possible, and use lower-case letters 24, qi Ki, Us Rd to denote rows in

---

(A) Fixed
(B) Logsparse
(C) Axial
VI
Dimott Expoctaton
Local Exportatoo
(D) Combiner Fixed
(E) Combiner- Combiner-Logsparse
(F) Combiner-Ae xial
Figure 1: Attention matrices of several instantiations of Combiner in the autoregressive setting. We
transform several sparse attention patterns: Fixed (A) [14], Logsparse (B) [18] and Axial (C) [20] to Combiner-Fixed (D), Combiner-Logspant
(E) and Combiner-Axial (F), Combiner approximates
the conditional expectation (3) with combination direct expectation (blue) and local expectation (yellow). Our instantiations (D)(E(F) achieves full attention with the same sub-quadratic complexity.
X, Q, K, V respectively, which corresponds to a location in the original sequence of length L. We use [n] denote the of positive integers (1.2....,
For a position E [L]. the attention formulation (1) can be viewed as conditional expectation of rows in V. Specifically, since softmax outputs probability distribution,
rewrite (1) as
A(x,) Ep(()) [uy],
p(j()) Z (x²)
exp
28
K
Vd"
where p(jlt) denotes the conditional probability at position j given the token at position i and (
over support Ω, The support Ω4 of p (jli) defines the of valid locations that the i-th token can attend For instance, the support set in autoregressive language modeling (LM) consists of all previous tokens, i.e., Ω = [2]2 ; in masked language modeling (MLM) the support consists of all tokens in the sequence, i... S.MMIM [L]. That is, 12..M and MLML represent the full attention capability respectively in the LM and MLM setting.
the partition function 7 (a::) = Ejen, exp (V
3 Combiner: Full Attention via Structured Conditional Expectation
The complexity of (jl() is the bottleneck of the computation for A (x;).. Generally, in existing sparse transformers, the support of (jli) is sparsified to reduce the computation and memory complexity, e.g., 52,Sparse ç S2I. for LM and S2Spars Ç S2ML 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 D.Combiner Ω for LM and Commiiner = M.M. for MLM, while still maintaining sub- uub-quadratic computation and memory cost. Below we denote Q2 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 we analyze the trade-off of Combiner.
3.1 Local Factorization for Conditional Expectation
The main idea of Combiner exploit a hierarchical structure for conditional probability modeling in (3), which provides the opportunity for reducing computation complexity while maintaining the
Following the conventional inplementation, the input sequence will be tright-shiften so that the position i can attent itself in LM setting.

---

same support. Specifically, we introduce support variables 22 for r = -...... .. and i E [2]. The support variables are disjoint, i.e., 22 22 Vr o s, and U 2 22 S2,. Then we can factorize p(jlk) as
p() = Σ P
= Σ p 1 1)p(1 ==0
p(11 (1 p)p(()))))
(4)
where ry denotes the index of the support to which j belongs. The last equation arises from the fact that the 22 are disjoint from each other (?2 n no = 0, Vr + 8). Therefore, there is only one support,
2 containing j. The remaining terms, where j $ 22 for r # rj are all zero since p (ji i) = 0.
Furthermore, assume 2 is a sufficient statistic, i.e., j and are independent given 22, we obtain p(jli) p (p( (t).
(5)
Given the partition {S) 1 the attention form in (3) can be rewritten as
A (2r)
[1y]
Σ Σ
(6)
Σ JEn
E
Σ pli p()()))))) JEMN
direct expectation
local expectation
where we consider direct attention in partition 22? and apply the local factorization (5) to the partition r I..... Ry.. Here p(jiit x p(ili) 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 Ω6 and multiple local expectations via p(j( and p(22))) to form the final conditional expectation.
Equivalently, we can also rewrite the structured attention (7) as
A(z,) Ejen.
I(j P(())) + Σ I(j e ())(1))((( Vyy
(8)
the new effective conditional probability o(ils)
where T(.) is binary indicator function. After reordering, one can see from (8) that we obtain the effective conditional probability q(ili) that tries to approximate the original p(jli). Each probability term depends on both current location and other location j. and the expectation is still obtained with respect to a valid conditional probability (non-negative and sumn up to over S22).
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 The full dependence is achieved via the multiplier p(?) i) where E O/. If we can design the local factorization such that:
1. the order of number of terms in (7) for p(-1),, Vi E [L]: Σt=, (n, + (QPl) is sub-quadratic; and let U4 = {?} [e(.). be the unique set partitions used for local expectation calculation, then the order of [2]] (i.e, the number of unique partitions in 24) is sub-quadratic;
3. the order of total number of unique calculations of local expectation across all locations in (7),
Eneu [22] is sub-quadratic;
then one can see that the overall computation and memory cost will be sub-quadratic with full attention support O [Combiner Di Vi E [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(-li) for simplicity. In fact, such local decompositions can be stacked further, which introduces a partition tree. Specifically, we can further partition ? with disjoint
Th subsets {12F*) k= and consider local decomposition p(j. (t) p(j] 1)p The | (),()))))) 1
where ky is the indx of sub-region which j belongs to. Thus, we obtain a hierarchical decomposition
of p (), which can also be plugged to (6) and yield new full attention formulation.

---

3.2 Parameterizing Conditional Probabilities
While we obtained a possible way to speed up the standard Transformer via a combination of direct expectation and local expectations, it is also important to have efficient design choice for the probability terms in (7), namely p(i|t) from direct expectation, p[] from local expectation and p( (i) for r E [1, ni]- For simplicity
the scaled dot-product, which means that will associate positions i and variable sets O with the corresponding embedding representation, and thus the probability is proportional to the exponential of the embedding inner products. Specifically:
p(]():: : As this term for the direct expectation, we can let p(]b) x exp( "a,, ). which is the same as vanilla attention (3) but with different normalizations, which will be explained in Equation 9. p((]))); This term aims to capture the joint event probability, i.e., p( x exp (aa S )r) Thus the design choice of kay should make abstraction of the corresponding support 27. We find kny = max poolingjen;" ky already provides good empirical results without introducing additional parameters; we can also use DeepSets [32] to obtain such abstraction.
p(j()); This term is the probability of getting within this local span 52 We make p(j x exp
where
max pooling DeepSets over {oi)yen; obtain gog similarly.
Normalizing Probability Terms. The terms in each local expectation p(j(2), Vj E 22 can be normalized within the local span; the direct expectation p(jli) and the terms in p(22) (t) should be normalized together,
((xi) = Σ
exp
4
+
exp
qi
TT
(9)
Vd
JEO
and (xx)) is the normalizing constant when calculating P(jli) and ((
3.3 Trade-offs in Combiner
Combiner achieves full attention with reduced cost without making explicit sparsity low-rank assumptions over the attention matrix. However this efficiency gain is not free. In this section we discuss the limitations of the simplification made by Combiner, and provide simple workaround.
Structured Attention Approximation. We obtain the local decomposition (5) under the conditional independence assumption. Therefore, the local expectation in (7) is independent of the position i, this suggests that any two locations i and i2 with 9 22 S would have linearly dependent attention scores over the region Q. Formally, the probabilities formed by the effective conditional
p(O,, (+))
distribution a(2), = ((il)))) 9(22)))),...... (Wer 111) = p(D: [52] a()) In other words, the
rank of the sub-matrix over the same partition in the resulting attention matrix 1, therefore, the attention matrix is locally low-rank based on the partition. On the other hand, the direct expectation fully attends to each position in sub-support Po which ensures the full-rank block. These two attention schemes make the attention matrix of Combiner structured. Compared with the low-rank approximation for attention [26, 28, 30], which is inspired from random features [29] in the kernel community, structured approximation that exploits both the locally low-rank and full-rank blocks has been proved more powerful theoretically and empirically in large-scale kernel machines [27].
Improving Expressiveness Using Mixture Model. One way to further improve the expressive-
ness of the local factorization is to use a mixture model. This idea adapted from the mixture of softmaxs [33] to obtain high-rank softmax layer in language modeling, Let w be certain partition
of the support (i.e,, collection of (2 of Ωü then to compute the attention, where each component of the mixture A wm.) is the term (7) using a specific factorization plan Wm. Empirically we find two components are already sufficient to improve performance.
can easily use A(x;) M Em=1 A(x;; wm)
1 M
4 Combiner Instantiations
In this section we show several local factorization schemes satisfying the requirements in Section 3.1. As we will see, Combiner is able to convert several sparse transformers [14, 18, 20-22] into full
5

---

attention, with the same order of computation and memory consumption. One can also design other factorization patterns, which can be easily instantiated in Combiner.
4.1 Combiner-Fixett
The Sparse Transformer [14] is one of the most representative variants that can achieve O L) computation and memory cost with sparse attention. Here we show how convert this fixed pattern proposed in [14] (Figure I(A)) into factorization plan, and instantiate a full attention variant named the Combiner-Fizee (Figure (D)).
In the fixed-sparse attention, the support is S2Paare MLM (j j mod = 03 ز = (div s) where s is a hyper-parameter, div is integer division, and j = i (div s) denotes that the quotients of i and j w.r.t. are the same. In the autoregressive case, 52, 1.M 2 MI.M n [i]. Please refer
to Figure ((A) for an illustration of the LM version.
Our design of MLM has the following form: Wixeed
22 [j= j = i (div s}} 2. = {j:] div 8 r. & {.] Vr E [T. div s), E
(10)
where each local expectation is performed in each span of size s, and there are totally L div s spans across all locations. For each position i E [Z]. there are (s + (L div s)) terms in (7); the local expectation has (L div 8) terms. The overall complexity is OL- (s 2 div 8))) The optimal is (VL)) and we can achieve OLLVL) computation and memory complexity, which is the same as [14] but here we gain full attention capability in each attention head. For the LM case, we can simply have J.M. {92 n [d] 22 E MMMM which has the same OLLVL) optimal complexity.
4.2 Combiner-Logsparee
The Logsparse Transformer is proposed in [18] and can theoretically achieve O(L log L) cost. The general idea is make the size of support 12,00r no larger than [log, i7. For the ease of notation, we first define bit)(n) = [b1. bg. brog mi] to be the binary representation of integer n, with
by E (0, 1} the coefficient of basis 2 Thus we have n = Σ 1= by 22. One of the possible design
m]
[log, i-1]
choices to make Logsparse in the LM case Q I.M = {ufff suff, := >! [logg
b, * 2
1=1
(1), ie., attend the location indices that equal to the suffix sum of the weighted bits(i - 1), as well as location itself. This serves as our base sparse version as shown in Figure I(B).
To exploit this scheme in the Combiner framework, we can define [log2 n non-overlapping supports,
where ? = [suff,] \. [suff with the boundary case [suff i i-11+1] = D. Note that for the ease of notation, some of the S2 are empty which will be ignored. In this case, the direct attention set S2 includes {i}, as well as {i - 1} when i is an even number. Such factorization leads to Combiner Logspprre,,,,,see,, as shown in Figure (EE). From the Figure, we observe that in total we will
have span summaries for every 2,4,8, (L) summaries. Each location will select at most (log(()) non-overlapping spans cover the full support Ω4, and thus, the total cost will be O log L). We leave the design of MLM case to Appendix B.
2 logga LJ locations, resulting in total Σ [logg z.) L
L or
4.3 Combiner-Axial
The Axial Transformer [20] builds the attention along each axis of the input data. Without loss of generality, we focus on 2D case where the input sequence reshaped into a matrix of size n x m = L. Specifically, the location i in original sequence will be in row; = (i (i-D) 1) div m + 1 and col, = (i 1) mod m + We show how to simply enable full attention with factorization on 2D matrix, hence Combiner-Axial.
The sparse axial has S2Sparse MLM Sa (uppurse MLM n [i]. which all have at most O(m + n) entries for each i, as illustrated in Figure ICC). We propose several factorization schemes to make it an attention with full support.
1(div m)), and [j:j - 1 1mod Uf
LM w verticlerticall S2. = 52sparse LM and 22 = (j : j r(mod m)) n [i- colo), for r E [m] \ (col.) As depicted in Figure 2(A), 22 corresponds to the column r above row,, where we use max pooling to

---

Attention matrix
Reshaped sequence
Attenticn matrix
(A) Combiner A cial-Vertical
(B) Combiner Axial- Horizontal
Figure 2: Attention matrices and sequence being attended (e.g., a 3x4 image) of vertical and horizontal variants of Combiner-Axial. Blue and yellow correspond direct and local attention respectively for location (purple). Locations connected by arrows correspond to the same support Ω
obtain the abstraction. To obtain such abstraction for all the locations, can leverage the cummax operator for each column to efficiently obtain the prefix-max.
"axial-horizomaal similar Waxial-vertical except that each 22 summarizes the before row; and excludes col; (Figure 2(B)).
w WLM
.LM "xxial-rowmajor" 22 = - - (div m)} n [i], i.e., elements in the same row are directly attended, while !2g = j j= r [i captures the rows before row; This structure is similar Combiner-Fixed except for the way that the abstraction (and thus the local expectation) is computed. Combiner-Fixxe computes the abstraction only based on of partition S where depends on both and the column col, (Figure 1(F)).
In all cases above, the cost is similar to the Axial Transformer [20], which is O(LVL) if we reshape the sequence to 2D matrix with n, OVVL). We defer the MLM case to Appendix C.
4.4 Combiner-Learnable
Inspired by the Reformer [21] and Routing Transformer [22], we can also learn the factorization plan w from the data. We illustrate this with Routing Transformer and provide way enable full attention in Routing Transformer following the Combiner principle.
For specific layer, suppose we have learned disjoint region (or cluster in Routing Transformer) {2""]= where U, = [L]. In Routing Transformer, simply have Quparse MLM
MLMM
S2 where Q
denotes the region where position i belongs to. To define the Combiner factorization, we let
Wrouting MLM 22 Ω S2 S2" \ 22, Vr [ng].
(11)
Note that n, (i.e,, number of learned clusters) for all locations. The above factorization can only work for MLM. LM requires the following definition:
Wrouting IM 2 Ω n [i], Ω (Ω" ni), Vr [ni]-
(12)
In general, both LM and MLM can have sub-quadratic cost when O(VL). However, routing variants (including the Routing Transformer) require a gather operation, which can slow on TPUs (see illustration in Appendix D).
Experimental Evaluation
We evaluate Combiner with different full attention patterns both autoregressive and bidirectional sequence modeling tasks, covering wide range of input data from images to texts. All tasks considered involve long sequences for up 12,000 in length, some which prevent the applicability of the vanilla transformer. We compare Combiner with state-of-the-an Transformers. We also perform a series of ablation studies where all the models being compared use the exact same architecture that only differ in the attention module, avoiding individual tricks employed in the original works (e.g., using both learnable and fixed patterns Routing Transformer [22]). Details to reproducing all experimental results can be found in Appendix E.

---

Table Ablation results in Bits per Dimension (Bits/Dim) on CIFAR- 10 and ImageNet-64.
Model
Layers
CIFAR-10
ImageNet-64
Reformer [21]
3.740
Performer [28]
3.335
3,719
Logsparse [18]
4.253
4.351
Combiner Logsparse (Ours)
3.366
3.795
Fixed [14]
3.408
3.696
Combiner- Fixed (Ours)
3.321
3.654
Axial [20]
3.666
4.032
Combiner- Axial (Ours)
3.050
3.585
Combiner-Mixture (Ours)
3,040
3.585
Reformer [21]
12
3.710
Performer [28]
12
3.310
3.636
Routing Transformer [22]
2.950
Membbiner-Mixture (Ours)
2.885
3.504
5.1
Autoregressive Sequence Modeling
In this subsection, we first perform density estimation on text and image using Combiner.
5.1.1 Language Modeling
Table 2: LM Perplexity on Wiki-40B (Main).
Model
Perplexity
For language modeling, we focus on the Wiki-40B- En dataset [34], which consists of clean Wikipedia pages in English. We use a sentence piece model with vocabulary size 32K tokenize the text and measure the perplexity at the sentence piece level. To ensure fair comparison, all models being com- pared again have the same number of layers and hidden sizes, are are implemented under the same code base
Transformer-22 2k [1]
17.26
Performer-2k [28]
19.66
Routing- 2k [22]
20.85
Fixed-2k [ [14]
18.04
Combiner-Fixed-22 (Ours)
17.70
Axial-2k [20] Combiner-Axial-2k (Ours)
20.82
17.56
Table 2 shows the results of the comparison. As can see, under 2k sequence length, Combiner vari- ants consistently better than their correspond- ing baselines, and are very close to the standard Transformer. When sequence length goes 8k, the standard Transformer runs out of memory, whereas Combiner continues to achieve improved perplex- ity, surpassing the result of Transformer-2k. If we further use DeepSets to calculate the summariza- tion terms qny and kny, we may further achieve lower perplexity as shown in Table 3.
Combiner-Fixed- -Kkk (Ours)
16.60
Combiner-Axial-8k (Ours)
16.49
Table 3: LM Perplexity Wiki-40B (Ablation).
Model
Perplexity
Transformer-24 []]
17.26
Combiner- Deeeineer-eeSSSss----Seee
PrrppSets-Ma--88 Max-8k (Ours)
16.29
Cumbiner-DeepSets Meennes---Mn-KK (Ours)
16.48
Combiner- (Ours)
16.60
Combiner-Mean-e (Ours)
16.54
5.1.2 Image Generative Models
CIFAR-10. We first perform a sanity check where we compare sparse attention baselines against Combiner with full attention under the same architecture on the CIFAR-10 dataset. The sequence length is 3072. For all the methods, we use same 6-layer transformer with 8 attention heads and 512 embedding dimensions. We train all models for 500k iterations using batch size 32 TPU v2. As shown in Table 1, given the same model architecture, Combiner-X performs significantly better than the base model X under the bits per dimension (BPD) metric on the 10,000 test images. In particular, Combiner significantly decreases BPD by 0.887, 0.087, and 0.626 compared to the base models Logsparse, Fixed and Axial, respectively. Note that all of the Combiner variants achieve better performance than the best of the base models. This demonstrates the advantage of Combiner over the baselines given the same l-yer architecture. We observe similar trend under a 12-layer architecture.

---

Table 4: Bits per Dimension (Bits/Dim) on CIFAR- 10 and ImageNet-64.
ImageNet 64x64
Bits/Dim
CIFAR-10
Bits/Dim
PixeICNN [15]
3.57
PixeICNN [15]
3.03
Parallel Multiscale [38]
3.70
PixcICN++ [36]
2.92
Glow [39]
3.81
Image Transformer [16]
2.90
SPN [40]
3.52
PixeISNAIL [37]
2.85
Sparse Transformer [14]
3.44
Sparse Transformer [14]
2.80
Axial Transformer [20]
3.44
Combiner-Axin (ours)
2.77
Routing Transformer [22]
3.43
Combiner-Axial (ours)
3.42
Following the 128-layer architecture in Child et al. [14], we apply Combiner-Axial and achieve state-of-the-arr performance, 2.77 BPD on CIFAR-10, as listed in Table We run all of the models in Table 4 without data augmentation [35].
ImageNet-64. We also evaluate performance under the autoregressive setting on ImageNet-64, where sequence length 12,288. We first perform the same analysis as CIFAR-10 and compare Combiner-X X with the baselines using the same model architecture. As shown in Table Combiner consistently outperforms the baselines with the same attention pattern. We further apply Combiner- Axial 30-layer Transformer, which achieves state-of-the-an performance on density estimation on ImageNet-64, demonstrating the effectiveness of full attention achieved by Combiner.
5.2 Bidirectional Sequence Modeling
Besides autoregressive tasks, we also evaluate Combiner set of standard bidirectional tasks to show the general applicability of the method.
5.2.1
Long-Range Arena
Long-Range Arena (LRA) is unified benchmark [31] for probing the capability of efficient trans- formers on handling long sequences. We evaluate our models on five tasks from LRA: ListOps, Text Classification, Retrieval, Image Classification and Pathfinder. All of the tasks are sequence-level multi-class classification. Please refer to the original LRA paper for more details.
Table 5: Experimental results Long-Rane Range Arena benchmark.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Avg
Chance
10.00
50.00
50.00
10.00
50.00
34,00
Transformer
36.38
64.27
57.46
42.44
88.81
57.87
Local Attention
15.95
52.98
53.39
41.46
84.64
49.68
Sparse TRRas.
35.78
63.58
59.59
44.24
83.90
$7.42
Longformer
36.03
62.85
56.89
42.22
86.68
56.93
Linformer
35.49
53.94
52.27
86.17
$3.28
Reformer
36.30
56.10
53.40
38.07
79.18
$2.61
Sinkhorn Trans.
34.20
61.20
53.83
41.23
73.36
52.76
Synthesizer
36.50
61.68
54.67
41.61
81.61
55.21
BigBird
37.08
64.02
$9.29
40.83
86.75
57.59
Linear Trans.
17.15
65.90
53.09
42.34
88.13
53.32
Performer
36.00
65.40
53.82
42.77
88.76
57.35
Combiner- Fixed
36.65
64.99
59.81
41.67
88.59
58.34
Combiner-Axial 36.15
64.36
56.10
41.33
88.43
57.27
As shown in Table 5, Combiner is able match the performance vanilla Transformer and achieves even better performance in some tasks. Following the protocol of LRA, all methods use the same architecture and hyperparameters for controllable comparison. We use the numbers from Tay et al. [31] for all tasks except for Pathfinder. Since we were unable to reproduce the original Pathfinder results using the default setup LRA Github repository, we rerun the baselines using Pathfinder- inter configuration to conduct fair comparison. However, as the benchmark is still of small-scale and the LRA official website discourages hyperparameter tuning, Table 5 should be treated as results for the test bench of expressiveness compared to vanilla Transformer.

---

Table 6: MLM perplexity C4 dataset.
Trandfomer
BigBrd Combiner-Aotal
Combiner
Sparse Fixed
Performer
Sparse-letil
Model
Perplexity
201
Transformer-22 2k[ [1]
4.552
Mileconts berion
BigBird-2k [41]
4.696
%% %
Numerr roa)
27
Performer-2K [28]
10.940
Fixed-2k [14]
5.279
Combiner- Combiner-Fixed-25 22 (Ours)
5.170
Axial-2k 2k [20]
5,370
21
21-
201 2
21
212
21
212
Combiner-Axial-22 (Ours)
4.809
Scisetice Longeh
Sesuetice Lengtt
Routing-2k 2k [22]
6.703
Figure 3: We measure the inference runtime and memory usage for eight models. Overall Combiner has similar speed with Performer and its sparse counterpart but Vanilla Transformer
Combiner-Routing-KK (Ours)
6.539
BigBird-8k [41]
4.542
Combincr-Axial-8 (Ours)
4.190
quickly goes OOM when sequence length grows.
Combingr-Fixed-8L (Ours)
4.139
5.2.2 Masked Language Modeling
As the core element of BERT langauge pretraining [5], masked language modeling (MLM) refers to the task of reconstructing tokens that are randomly masked out in the input sequence. As with the LM task, we use perplexity as the main metric, which correlates relatively well with down-stream task performance. Specifically, we use the large scale C4 dataset [8] for training and evaluation, and consider different sequence lengths. Following the original BERT setup, we mask out 155% of the tokens in each input sequence. The comparison is summarized Table 6. Similar to the LM result, different Combiner variants consistently outperform their corresponding baselines under 2k sequence length. However, apart from the standard Transformer, Combiner-2k also falls behind BigBird-2k. We conjecture that this is related to the special design in BigBird such all tokens always attend to the <cls> token directly, which only applicable in con-causal problems. That said, when further increase sequence length to 8k, the standard Transformer runs into OOM issue, whereas Combiner not only outperforms BigBird but also substantially surpasses Transformer-2k. This suggests that Combiner can truly benefit from scaling learning longer sequence lengths.
5.3 Runtime and Memory Usage Combiner
Here we evaluate the inference runtime and memory usage of five baselines Transformer, Performer, BigBird, Sparse-Fixed and Sparse-Axial, well three variants of Combiner- Combiner-Fixe Fixed, Combiner-Axial and Combiner-Mixture. We run inference of all the models on a TPU v3-16 (16 cores 16GB) with batch size 16, and test sequences of length from 20 to 2 As shown in Figure 3, Combiner instantiations achieve comparable runtime and memory usage with their sparse counterpart and Performer. Note Combiner achieves much better empirical performance than the sparse models and Performer. Combiner-Mixture has the same asymptotic complexity with Combiner-Fixed and Combiner-Axial, however, since it requires running two partition plans, is slower than Combiner-Fixee and Combiner-Axidl Due to the gather operation required by the random attention which is not very TPU/GPU friendly, BigBird is very computationally expensive. And the Transformer model quickly runs out of memory when sequence length increases.
6 Conclusion
Inspired by the conditional expectation view of attention mechanism, we propose Combiner, drop-in replacement of the attention module. By introducing structured decomposition to the conditional probability, Combiner achieves full attention capability while maintaining sub-quadratic computational and memory cost. We instantiate several Combiner variants converting existing sparse transformers to full attention. Combiner achieves state-of-the-ar performance on both autoregressive and bidirectional tasks for image and text modeling, showing benefits in both modeling effectiveness and runtime efficiency. Future work includes additional factorization pattern designs, as well as applications of Combiner in domains like bioinformatics and speech.
10

---

Acknowledgments and Disclosure of Funding
We would like to thank Richard Song and David Dohan for the help on introducing Performer codebase and experiment configurations, Yi Tay and Mostafa Dehghani for clarifications on the LRA bench- mark, James Lee-Thorp. Joshua Ainslie, and Ilya Eckstein for clarification on their LRA experiment results, Adams Yu for performing internal paper review and helpful suggestions. We also grate- fully acknowledge the support of DARPA under Nos. HR001 19190039 (TAMI), N66001 1924033 (MCS); ARO under Nos. W91 INN- 1011--10344 1-342 (MURI), W91 INF-16- 112-11--1--11 1-0171 (DURIP); NSF under Nos. OAC- 1835598 (CINES), DAC-193457, (HDR), CCF-1918940 (Expeditions), IIS-2030477 (RAPID), NIH under No. R56LM013365; Stanford Data Science Initiative, Wu Tsai Neurosciences Institute, Chan Zuckerberg Biohub, Amazon, JPMorgan Chase, Docomo, Hitachi, Intel, JD. JD.com, KDDI, NVIDIA, Dell, Toshiba, Visa, and UnitedHealth Group. Hongyu Ren is supported by the Masason Foundation Fellowship and the Apple PhD Fellowship. Jure Leskovec a Chan Zuckerberg Biohub investigator.
References
[1] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Informa- tion Processing Systems (NeurIPS), 2017.
[2] Mia Xu Chen, Orhan Firat, Ankur Bapna, Melvin Johnson, Wolfgang Macherey, George Foster, Llion Jones, Niki Parmar, Mike Schuster, Zhifeng Chen, et al. The best of both worlds: Combining recent advances in neural machine translation. In Annual Meeting of the Association for Computational Linguistics (ACL), 2018.
[3] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[4] Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Ruslan Salakhutdinov, and Quoe Le. Xlnet: Generalized autoregressive pretraining for language understanding. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[5] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Annual Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies (NAACL-HI.T)) 2019.
[6] Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut, Albert: A lite bert for self-supervised learning of language representations. In International Conference on Learning Representations (ICLR), 2020.
[7] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 1rXiv:190..11692 2019.
[8] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint
2019.
[9] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani. Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Leurning Representations (ICLR), 2021.
[10] Aditya Kanade, Petros Maniatis, Gogul Balakrishnan, and Kensen Shi. Learning and evaluating contextual embedding of source code. In International Conference on Machine Learning (ICML), 2020.
[11] Linhao Dong, Shuang Xu, and Bo Xu. Speech--ransformmer: a no-recurrence sequence-to- sequence model for speech recognition. IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 2018.

---

[12] Ali Madani, Bryan McCann, Nikhil Naik, Nitish Shirish Keskar, Namrata Anand, Raphael R Eguchi, Po-Ssu Huang, and Richard Socher. Progen: Language modeling for protein generation. arXiv preprini arXiv:2004.03497, 2020.
[13] Cheng-Zhi Anna Huang, Ashish Vaswani, Jakob Uszkoreit, Noam Shazeer, Curtis Hawthorne, AM Dai, MD Hoffman, and D Eck. Music transformer: Generating music with long-term structure (2018). In International Conference Learning Representations (ICLR), 2019.
[14] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904, 10509, 2019.
[15] Aaron Van Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks. In International Conference on Machine Learning (ICML), 2016.
[16] Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz. Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning (ICML), 2018.
[17] Jack W Rac, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. In International Conference on Learning Representations (ICIR), 2020.
[18] Shiyang Li, Xiaoyong Jin, Yao Xuan. Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting, In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[19] Zilong Huang, Xinggang Wang, Lichao Huang, Chang Huang, Yunchao Wei, and Wenyu Liu. Conet: Criss-cross attention for semantic segmentation. In International Conference on Computer Vision (ICCV), 2019.
[20] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. arXiv preprint 1912 2019.
[21] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations (ICLR), 2020.
[22] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguisticss, 9:53-68, 68, 2021.
[23] Zihang Dai, Zhilin Yang. Yiming Yang, Jaime Carbonell, Quoc V Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Annual Meeting of the Association for Computational Linguistics (ACL), 2019.
[24] Zhuoran Shen, Mingyuan Zhang, Shuai Yi, Junjie Yan, and Haiyu Zhao. Factorized attention: Self-attention with linear complexities. CoRR, 2018
[25] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 200X::0000004768, 2020.
[26] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnss: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning (ICML), 2020.
[27] Si Si, Cho-Jui Hsieh, and Inderjit S Dhillon. Memory efficient kernel approximation. The Journal of Machine Learning Research, 2017.
[28] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations (ICLR), 2021.
[29] Ali Rahimi, Benjamin Recht, et al. Random features for large-scale kernel machines. In Advances in Neural Information Processing Systems (NeurIPS), 2007.
12

---

[30] Hao Peng. Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations (ICLR), 2021.
[31] Yi Tay, Mostafa Dehghnni, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. In International Conference on Learning Representations (ICLR), 2021.
[32] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. In Advances in Neural Information Processing Systems (NeurIPS), 2017.
[33] Zhilin Yang, Zihang Dai, Ruslan Salakhutdinov, and William W Cohen. Breaking the soft- max bottleneck: A high-rank mn language model. In International Conference on Learning Representations (ICLR), 2018.
[34] Mandy Guo, Zihang Dai, Denny Vrandečié, and Rami AI-Rfou. Wiki-40b: Multilingual language model dataset. In Proceedings of The 12th Language Resources and Evaluation Conference, 2020.
[35] Heewoo Jun, Rewon Child, Mark Chen, John Schulman, Aditya Ramesh, Alec Radford, and Ilya Sutskever. Distribution augmentation for generative modeling. In International Conference on Machine Learning (ICML), 2020.
[36] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelenn++: Improving the pixelenn with discretized logistic mixture likelihood and other modifications. In International Conference on Learning Representations (ICLR), 2017.
[37] Xi Chen, Nikhil Mishra, Mostafa Rohaninejad, and Pieter Abbeel. Pixelsnail: An improved autoregressive generative model. In International Conference on Machine Learning (ICMI.), 2018.
[38] Scott Reed, Aäron Oord, Nal Kalchbrenner, Sergio Gómez Colmenarejo, Ziyu Wang, Yutian Chen, Dan Belov, and Nando Freitas. Parallel multiscale autoregressive density estimation. In International Conference Machine Learning (ICML), 2017.
[39] Diederik P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible Ixl convolu- tions. In Advances in Neural Information Processing Systems (NeurIPS), 2018.
[40] Jacob Menick and Nal Kalchbrenner. Generating high fidelity images with subscale pixel networks and multidimensional upscaling. In International Conference on Learning Represen- tations (ICLR), 2019.
[41] Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, et al. Big bird: Transformers for longer sequences. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[42] Chulhee Yun, Wen Chang, Srinadh Bhojamapalli, Ankit Singh Rawat, Sashank J Reddi, and Sanjiy Kumar. o(n) connections are expressive enough: Universal approximability of sparse transformers. arXiv preprint arxiv:2006.04862 2020.
13


Reviews:
[{"review_id": "Xrc2SWoKap", "reviewer": "Reviewer_ETpX", "summary": "This paper proposes the combiner, which could achieve full attention capability while maintaining sub-quadratic computational and memory cost.", "questions": "", "limitations": "", "rating": 9, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: This paper is pretty novel.\n\nClarity: This article is very clear and the formula derivation is detailed and rigorous.\n\nSignificance: For a long time, self-attention has been computationally cumbersome, especially when dealing with some very long sequences. This article makes self-attention more efficient while ensuring expressive ability. The theoretical derivation in this article is very sufficient, and the experiments are also very solid.\n\nQuestion: I am curious about the comparison between this paper with the spatial linear layer[1,2]. Could this method be more efficient? \n\n[1]. Synthesizer: Rethinking Self-Attention for Transformer Models\n[2]. Pay attention to MLPs", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, this method keeps the full attention capability while maintaining low computation and memory complexity.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2 hours"}, "extra_text": {}}, {"review_id": "S2ciB2JlTn", "reviewer": "Reviewer_4mpH", "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.\nNumerous experiments demonstrate the utility of the proposed approach.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The topic of approximate attention is a valid question.\nThe current paper seems to have a very performant approach to that, and their experimental results seem solid.\nHowever, I am concerned that some of the mechanics of the method are not fully explained, and too strong claims are made. Some revisions could make it better.\n\n- l. 5 \"full attention capability\" l.132 \"full attention model almost for free\" - I am puzzled by the claims that the method is equivalent to full attention, while other sections of the paper say it is an approximation. Please clarify this exactly.\n\n- l.123 subquadratic cost. I do not understand fully the claim. In general, if A(x_i) is attention for one query position i, and there are L such positions i. To calculate each A(x_i) in equation (8) one sums over all j positions, which are also L --> it seems to me still quadratic.\n\n- also, please discuss how the choice of n_i, the partition size and count, affects the computation costs and model capacity of the Combiner\n\n- section 3.3 tradeoffs - please improve and  clarify that section, my above comments are also concerned with better understanding of the factorization and its expressiveness\n\n- l.150 I do not understand how exactly max pooling is applied, is it per dimension d of the keys k_j, or the most active key overall, or some other way to do it?\n\n- reference [40] BigBird seems very related to the Combiner. While it is good there is an empirical comparison with it, would be good to mention it also in the introduction related works.  The claims of both papers (efficient attention that keep the properties of full attention due to complicated patterns) seem related to me, so would be good to elaborate on that. I saw that the appendix has some mentions on both papers and the \"universal approximation property\", but it did not answer fully my questions for the comparison of the two methods.\n\n- Experiments, in table 4 the Combiner is even better than full attention, please discuss this, e.g. mention whether this is due to some inductive bias or other reason?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "see above comments for approximation properties of the Combiner", "review_text": "", "extra_scores": {"ethical_concerns": " ", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "OrVhIAV4OFE", "reviewer": "Reviewer_GEmp", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper proposes a new way to factorize attention in transformers using sparse patterns without sacrificing expressiveness.  Even though prior work has achieved linear time complexity attention without sacrificing expressiveness using exponential kernel approximations, the conditional expectation view of attention is a novel contribution that should be of interest in the research community that focuses on efficient transformers.\n\nThe proposed framework is theoretically sound and provides a unified view of existing sparse attention patterns and how they can be converted to full attention approximations which is a solid contribution.  The exposition of the structured factorization and its instantiations were clear and comprehensive for the most part. One limitation that was not addressed in the discussion is the computational cost which is bounded by that of the sparse pattern methods, namely sub-quadratic, which is still higher than linear complexity methods. Prior work has shown that both low-rank and full-rank blocks in the approximation are more powerful in theory but it is not clear if it holds up for the proposed approximation since it uses a mixture model to improve expressiveness.  Also, being more powerful also comes at a computational cost; the main point of efficient attention is to achieve a competitive speed-quality tradeoff.  This is an aspect that was not addressed in the model or the evaluation sections (see details below).  \n\nThe evaluation which is mainly quality-driven shows that the proposed models provide consistent improvements compared to sparse attention methods; the results support the main claim of the paper about improving expressiveness with sparse computational cost. However, the differences between them in terms of speed in practice have not been explored. The proposed models have some additional components that potentially add computational cost in practice. The rest of the evaluation mostly focuses on achieving state-of-the-art results which is great but I think it lacks emphasis on the speed-quality tradeoffs and it lacks an ablation analysis of the components used in the design of the architectures such as Deep Sets for representing the abstraction and the mixture of softmax to obtain a high-rank approximation.  There are also some details that would be helpful for better assessing the results e.g. including the number of parameters for each attention method and reporting the speed achieved by each method in the Long-Range arena benchmark. The only part that involves speed experiments is in Figure 3 which is not thoroughly discussed in the text.  \n\nQuestions: \n\n- What is the impact of deep sets and the mixture of softmaxes on the performance of the model? An ablation analysis would help to show which component is most important and how good is the approximation without adding any additional parameters. \n\n- The experiment in Figure 3 lacks a detailed description.  One would expect Performer to be faster since it has lower complexity than the proposed models but it's not the case. Could the authors elaborate on that? Does the experiment measure the time it takes to encode or generate sequences of X length? Ideally, we would want to see what happens when generating sequences which is a  challenging setting where parallelization is not possible.\n\n- For the long-range arena, it would be informative to show the runtime for each of the methods so that it is easy to compare methods in terms of speed-quality tradeoffs. Currently, it's not possible to know how good tradeoffs the proposed models achieve in practical tasks.  ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The papers discuss theoretical limitations regarding the expressiveness of the proposed factorization which is quite informative. It would be worth including some discussion about the reliance on hand-designed sparsity patterns, the computational cost compared to linear time and space attention methods, and investigate how good speed-quality tradeoffs the proposed models can achieve in practice compared to alternatives. There is no discussion provided about the potential positive or negative societal impact of the work. Although there is no obvious concern, it would be good to see that the authors have thought about it.  ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.5"}, "extra_text": {}}, {"review_id": "4QcQbXghwjY", "reviewer": "Reviewer_u2iY", "summary": "This paper presents combiner, a new variant of sparse Transformer, which is inspired by viewing self-attention as a conditional expectation over embeddings at each location. The authors show that combiners achieve better performance than existing sparse Transformer variants.", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. The proposed method is novel and simple. It is surprising to see that simple max-pooling already produces strong performance.\n2. The paper is well written and easy to follow.\n3. It is great to see some significant improvement over the previous methods on a wide range of tasks.\n\nWeaknesses:\n1. It would be great if we could see a table containing the memory footprint and inference speed results following the setup in Table 2 in the Long Range Arena paper. Otherwise, it is unclear whether the proposed method saves the memory as well. Also, having such a table provides the numbers for future works to compare with.\n2. The authors only time the speed of Combiner-Axial and Combiner-Fixed. The speed overhead of Combiner-Mixture is unrevealed.\n3. It is unclear to me how many partitions and what they are in the Combiner-Mixture model in Table 1.\n4. The proposed Combiner-Logsparse is not evaluated in the experiment section. I wonder what the reason is.\n5. On some tasks Combiner-Axial outperforms Combiner-Fixed, but on some others, it doesn't. This makes it hard for the practitioners to decide which to use.\n6. It would be better to see an ablation study on different alternatives to max-pooling when combining queries and keys. For example, how do DeepSets mentioned in subsection 3.2 perform?\n\n\nOverall, the proposed method is original, interesting, and effective, so I recommend accepting it. Admittedly, if the authors can address the weaknesses above (at least the first four) in their rebuttal, I will raise my rating to 8.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The main limitation of the paper is that the real overhead using the proposed methods is not fully revealed (see the weaknesses part above).", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3.5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a novel efficient Transformer variant, called Combiner, which is a drop-in replacement of attention, achieving full attention with sub-quadratic cost using structured factorization. The reviewers have several concerns regarding the experimental analysis in the initial review. During the rebuttal period, the authors provided intensive analysis and addressed most of the concerns from the reviewers.

All reviewers recommend acceptance, and I hope the authors could include the new experimental results and ablation studies in the camera-ready version.
Decision: Accept (Spotlight)