An Attention Free Transformer
Authors: Shuangfei Zhai
Walter Talbott
Nitish Srivastava
Chen Huang
Hanlin Goh
Joshua M. Susskind
Venue:  (2021)
Keywords: Transformers
attention
efficient

Abstract: We introduce Attention Free Transformer (AFT), an efficient variant of Transformers \citep{transformer} that eliminates the need for dot product attention. AFT offers great simplicity and efficiency compared with standard Transformers, where the multi-head attention operation is replaced with the composition of element-wise multiplications/divisions and global/local pooling. During training time, AFT has linear time and space complexity w.r.t. both the sequence length and feature dimension; in the autoregressive decoding mode, AFT has constant memory and time complexity per step. We show that, surprisingly, we are able to train AFT effectively on challenging benchmarks, and also to match or surpass the standard Transformer counterparts and other efficient variants. In particular, AFT achieves the state-of-the-art result on CIFAR10 autoregressive modeling with much reduced complexity, and also outperforms several efficient Transformer variants on Enwik8.

Under review as a conference paper at ICLR 2021
AN ATTENTION FREE TRANSFORMER
A nonyymous authors
Paper under double-blind review
ABSTRACT
We introduce Attention Free Transformer (AFT), an efficient variant of Trans- formers (Vaswani et al. 2017) that eliminates the need for dot product attention. AFT offers great simplicity and efficiency compared with standard Transform- ers, where the multi-heas attention operation is replaced with the composition of element-wise multiplications/diviiioo and global/local pooling. During training time, AFT has linear time and space complexity w.r.t. both the sequence length and feature dimension; in the autoregressive decoding mode, AFT has constant memory and time complexity per step. We show that, surprisingly, to train AFT effectively on challenging benchmarks, and also to match or surpass the standard Transformer counterparts and other efficient variants. In particu- lar, AFT achieves the state-of-the-an art result on CIFAR 10 autoregressive modeling with much reduced complexity, and also outperforms several efficient Transformer variants on Enwik8.
able
INTRODUCTION
Attention mechanisms, represented by Transformers (Vaswani al. 2017). have driven the advance- ment of various machine learning problems, including language modeling (Devlin et al. 2018 Rad- ford et a image modeling (Chen et al.), and set modeling (Lee et al. 2019). Different from other well known model architectures such as Convolutional Neural Nets (CNNs) or Recurrent Neural Nets (RNNs), Transformers enable direct interaction between every pair of elements within a se- quence, which makes them especially powerful at capturing long term dependencies.
However, Transformers require high computational costs. The root cause of this challenge is the need to perform attention operations that have quadratic time and space complexity w.r.t the context size. This makes it especially difficult for Transformers to scale to inputs with large context sizes. A number of recent works have been dedicated to addressing the scalability issue of Transformers (Child et al., 2019 Kitaev et al., 2020) Rae et al. 2020; Wang et al. 2020b; Katharopoulos et al., 2020 Tay et al. [2220a] Choromanski et al. 2020). While the techniques adopted in the literature range from sparsity, locality sensitive hashing, low rank decomposition, kernel approximation and etc., most of them trying to approximate the full attention operation.
In this paper, we take a bolder step towards the same goal, by proposing a computational module that does not use or approximate the standard dot product attention. We hence name our model the attention free transformer (AFT). Similar to dot product attention, AFT is composed of the interaction of three quantities, namely the query, key and value. What's different, however, that AFT operates solely based on element- wise operations. To be more concrete, they key and value are first multiplied element-win wise, the result of which is then pooled over the context dimension (in the causal model, this corresponds to a cumulative sum). The query is then multiplied with the reduced key-value representation element-wiss wise to produce the final output. See Figure Ta for an illustration.
AFT maintains the full advantage of dot product attention, namely direct interaction between any two elements in a sequence (up to proper masking). However, the computational cost is drastically reduced to d( complexity for time and space, where T,, d are the context length and feature dimension, respectively. In the autoregressive decoding mode, AFT also provides constant decoding time and space complexity per step, compared to (TT) for standard transformers. To the best of our knowledge, AFT is the first model that achieves such efficiency in the context of Transformers. See Table for the complexity analysis of AFT in comparison to other variants.

---

Under review as a conference paper at ICLR 2021
Table 1: Complexity comparison with different Transformers: Reformer (Kitaev al. 2020), Syn- thesizer TTay et al. 2020a), Linear Transformer (Katharopoulos et al., 2020) (only variants that support the causal mode are shown). Here T, d denote the sequence length and feature dimension, respectively.
Model
Time @ train
Space @ train
Time/step @ decode
Space/step @ decode
Full Attention
O(T*d)
O + Td)
O(Td)
((Ta)
Reformer
logloTT
O[T Td)
O T +
((Td)
Synthesizer
O(T²d)
0 + Td)
O(Td)
(
Linear Transformer
((Td)
O(Td+
O(d))
(d2))
AFT (ours)
((Td)
((Td)
O(d)
O(d)
We show that we can interpret AFT as an extreme case of multi head dot product attention (MHA). In particular, show that by 1) setting the number of heads equal the feature dimension in MHA and 2) using relu in place of softmaa as the non-linearity, MHA can be decomposed into the summation of two AFT modules (see Equation 6). However, this relationship is not true in general sense, i.e., by varying the non-linearity injected after the query and key in AFT, we can obtain models that do have MHA counterpart. This realization allows us to freely explore the design choices (e.g., nonlinearity) of AFT to achieve the best performance, This philosophy is in direct contrast with previous and concurrent "linearized attention" works (Katharopoulos et al., 2020 Choromanski et al. 2020), which are constrained by the design space of MHA.
We perform experiments with AFT on several benchmarks, including unconditional image model- ing, image super-resolution, language modeling, machine translation and point cloud generation. We show that AFT works very well as an alternative to the standard Transformer, providing competitive results as well as excellent efficiency.
To summarize, our contributions are as follows:
We propose AFT, a new family of Transformer models that achieves O(Td) time and space complexity in training, as well as (d) time and space complexity in autoregressive decod- ing.
We show strong performance of AFT drop in replacement of MHA on various bench- marks, including setting the state-of-the-art result on CIFARIO in the standard setting and outperforming other efficient Transformer variants.
2 MULTI-HEAD ATTENTION
At the core of Transformers is the Multi-Head Attention (MHA) operation. Given three sequences, namely the query Q E xd' key K R2xxd and value V R2 and the number of heads h. MHA performs a scaled dot product attention for each head i, defined as:
F(Q.K.V) = (K)) Vdl.
W., s.t. Q QW, K KW., vV VVWY
(1)
where Wo E Raxd2, W.K. E Rdxds, wy E Raxde are linear transformations for head i, and o is the non- linearity by default set the softmaz, function (subscript indicates softmax is applied to each row of matrix). de. dy are dimensions for key and value, respectively. MHA concatenates the output of attention heads along the channel dimension, resulting feature dimension hd,, Unless otherwise mentioned, we assume dę d, and h d This means the query, key and value are the same dimension within each head, and output dimension matches that of the input.
d,
3 METHODOLOGY
3.1 ATTENTION FREE TRANSFORMER
We now define Attention free transformer (AFT), which provides an alternative to MHA, Given
Q, K, V. AFT first linearly transforms them into Q' QW?, K' KWK V VWV. then
2

---

Under review as a conference paper at ICLR 2021
performs following operation:
T [Q0 @ & (ox(K')V) 1)
(2)
where @ is the element- product, with support for broadcasting when the operands' dimensions og dg are nonlinearities applied to the query and key, respectively. Explained
don't exactly match[| [ in words, the key and value are first combined with element-wise multiplication, the result of which is then pooled over the context dimension, yielding a fixed length context vector E Rd, This context vector then multiplied with each row of the query, which forms the final output of an AFT layer.
One particularly useful variant of MHA is masked attention, oftentimes presented in the form of causal attention. Specifically, in auto-regressive models, queries are constrained to not be able to interact with keys and values beyond the curret position. In standard attention, this is usually implemented with an explicit binary masking matrix of shape Tx T, with non-causal entries masked also straightforward to extend AFT to the causal mode while maintaining
as O. We show that its efficiency. We denote an AFT layer's output as Y, = f(Q Ket, Vst), t= = L...., We formulate the casual AFT as:
Y = (1 ④
Σ
2 V Vedee t= I,..,
(3)
where the subscript X, indexes the tth row of matrix X.
The design philosophy of AFT is to promote extreme efficiency, while keeping the benefit of standard Transformers. Concretely, AFT enables direct interaction of any two elements within the sequence, which is arguably the biggest advantage of Transformers over other types of models such as RNNs and ConvNets. However, AFT gets rid of the need of performing the costly spatial dot product attention, by computing a reduced value representation with the weights only depending the keys. The resulting operation has extremely efficiency of (TTd) w.r.t. both time and space, which is the first model that achieves linear complexity along both context and feature dimensions. Moreover, the causal mode of AFT has additional advantage of a constant decoding cost per step, similar to (Katharopoulos et al. (020). To see this, from Equation 3 we have a simple recursion of Y; = (,,()) (ax (K;) @ V2 + KV--1) with KV, = Σ (ax (K2) @ VT),. assuming Tg: a% both element-wist functions. One thus only need keep KV, in memory, and update it with constant cost per step.
Discussions:
Selecting nonlinearies: Tg: Th provide additional nonlinearity which helps to increase model's capacity. Empirically, we have found that one particularly strong setting is to let a sof tmaa which is normalized along the context dimension. This choice brings an interesting benefit especially in the causal mode, which we can explicitly write as:
Y = =0) & (1(()) O VV Ea(()o:) (c(())
exp(Ki)
(4)
Ei (xp(K)))
Here (t) acts as role similar to that of an input gate in LSTM, and ge (t') is operating like the forget gate, which depends on the input of time f, dynamically downweights the contribution of past time steps. When augmented with standard position embeddings as commonly used in Transformers, this allows the model to be able to learn the notion of recency at the same time of having access to the full context in the history. From this view, og can also be interpreted as the output gate, for which we found that both sigmoid and relu work well, with the former being slightly better. Also note that the same space and time complexity still holds for o softmaz, both in training and decoding. In our experiments, unless otherwise mentioned, we use the sigmoid + sof Imar setting for a and o by default.
Relation to MHA: Although AFT performs fundamentally different operations than standard atten- tion, we show that the two family of models overlap the extreme case. To see this, we explore
We adopt Numpy styled broadcasting convention: We assume here that Y, includes input information at the current position t, the version where the current position is excluded be obtained by shifting the outputs to the right.
http://mmwy.mm////sssssllleeeeer htrp:///eeeeeroorbrrrrsassnnnnn html
3

---

Under review as a conference paper at ICLR 2021
the limit of number of heads in MHA, which amounts to letting dj = 1 for each head. In this case, the dot product operation within each head reduces to scalar product. Next, we set a to be relu instead of softmor in Equation _ In this case, have:
fiQ,, K, V) [Q] [Q;'
((9)++(()))) +--1))))+++-1---)(())))))))))))1
(5)
= [Q]]+ (()))VV))) + +1-00-
where [.]+ denotes the relu operator, and Q, K", V E RT by definition. The concatenated
output of the attention heads can then be concisely written as:
s(Q, K. V) [Q')+ O Σ ((K)+OV), o +1 +-1 @ +))
(6)
which consists of two terms, each of which is an AFT operation, with 0g ok H+ and 0g okk = [-, respectively. However, note that this correspondence is not general, i.c., AFT does not need approximate any MHA counterpart and can indeed have very different inductive biases than that of a standard Transformer.
Relation to Linearized Attention: There few recent works proposing linearize the dot product attention (Linear Attention) from the view of kernel approximation, first proposed in Katharopoulos (2020) and also in concurrent work (Choromanski al. 2020). (Katharopoulos et al. 2020) proposes the linear attention operation in the form:
p(2)) E (o(K)) Y
6(Q)) O(Q) (x ((K))"))"
where Qú Ky, V/ are all row vectors of Rd. Equation similar AFT, the sense that the key and value are first combined and reduced in both cases. However AFT differs in two aspects: 1) the time complexity of Linear Attention OTdd)) which is linear in the sequence length but has difficulty scaling to wide networks 2) Linear Attention is designed to approximate MHA, where the nonlinearity on query and key are shared. AFT however, we show that is beneficial to search for different nonlinearities for both the query and key.
3.2 LOCAL CAUSAL AFT
In autoregressive modeling, locality a strong and effective inductive bias, as has been explored in Chen et al.) Child al. (2019). We similarly propose augmented version of causal AFT, where
we rewrite Equation 3 as
Y = (9))
Σ we (ox(K') o VEL) 1 L..,,, 1'11
(8)
where wrr E R is a locality masking scalar. We consider two strategies of constructing w. the first being the hard local mask where we have wt, = 1 if t - t's and 0 otherwise, with being the desired window size (for 2d inputs such as images, we can similarly construct 2d windows, see Appendix for details). The second one, which works better in practice, is to learn a position based local bias, while still assigning non-zero weights out of window contexts. More concretely, we <sp([[tt='<=)u] vy)
let Wi," Σ... C) where I(-) is an indicator function, and u, RTx are two sets of low dimensional learnable position embeddings, independently learned per layer. Note that in this case, we maintain dense connection between every t, pair, but rather introduce learnable biases for more recent contexts. We typically set d, to be small number (e.g., 64) which greatly reduces the amount additional parameters compared to learning a full matrix. The learned version also adds very little overhead to both the time and space complexity, w is sparse and static, whose memory cost is marginalized out across batches during training. We denote the two versions as AFT-local-hare hard and AFT-local--caande respectively.
4 RELATED WORK
Since the Transformer was introduced, there have been numerous attempts to address the major source of inefficiency in the architecture, the quadratic cost of the attention operation. Improving this
4

---

Under review as a conference paper at ICLR 2021
Qi
K
V
Σ
exp
Σ
Rig
pool
(a) Causal attention free operation.
(b) Local causal attention free operation.
Figure I: AFT blocks require only element-wise and pooling operations.
operation can enable larger context sizes and more efficient implementations. For a comprehensive, recent survey of efficient transformers, see (Tay et al. 2020c).
Approximating the dot product. Katharopoulos et al. (2000; Choromanski et al. (2020) both propose to approximate the exponential kernel with inner product of projections, which leads to linearized attention operation of complexity O(Td)). AFT is similar but offers greater efficiency of O(Td) due to the exclusive use of element seseent-wise operations, well as more design flexibility. Reformers (Kitaev et al.] 2020 apply LSH as an approximation to the dot product, whereas AFT completely gets rid the need of dot product.
Sparse, local aftention. Sparse Transformers (Child et al. 2019) and Image Transformer (Parmar ct al. 2018)) proposes to use fixed sparse or local context patterns. Attention models in vision tasks (often combined with convolutions) use image structure to help handcraft relevant spatial patterns to attend (Wang al.) [2020a] HHaan al. 2019b] Zhu al. 2019 Huang al.) [2019a] RRaacchnndann itt al. 2019). AFT also borrows the locality idea, but we put it as a bias rather than hard constraint (see AFF-local-learned) Also AFT is standalone module, where works plug in replacement of MHA in autoregressive tasks.
Context compression. Other approaches try to learn context patterns. Adaptive-Span Transformers (Sukhbaatar et al., 2019) learn a range for each attention head within which to attend. Routing transformers (Roy et al. 2020) use clustering to compute dot- dot-prode poo-produtt attention only over subset of elements within the same cluster. The Linformer (Wang et al.) 2020b) reduces the length of the context by compressing the keys and values with a linear layer. Compressive Transformers (Rae) et al. 2020) compute and update reduced representations of the input that are far enough back in the input sequence, and attend those compressed representations. AFT is largely complementary to these approaches, as focus is to improve the complexity of any given sequence from the operation level.
Eliminating dot product attention. Instead of limiting the number of comparisons, other methods change the operation used to compute attention. The Synthesizer (Tay et al,, 2020a) uses attention weights predicted from inputs, rather than derived from dot-product interactions. The LightConv module introduced in (Wu et al.) 2019) proposes to replace the dot product self-attention with dy- namic lightweight depthwise convolution, where the weights normalized across temporal dimen- sion. The Sinkhorn Transformer (Tay et al. (2000b) uses differentiable sorting operation to identify relevant comparisons that may not be local in the original sequence order. AFT can be viewed more drastic version in this direction, where the we adopt single global "attention mask" (w) of all ones (vanilla AFT) or with few learnable entries (AFT-local-leam) learned..
Gated RNNs. AFT is also related to the classic line of work on gated RNN variants, including LSTMs Hochreiter & Schmidhuber (1997), GRUs (Chung et al. 2014) and QuasiRNNs (Bradbury et al. 2016). AFT maintains the benefit of RNN models (linear complexity w.r.t. sequence length, constant decoding cost), but offers great parallelism and effectiveness, thanks to the use of simple context reduction operation, which is amendable to fully parallel implementation during training. We believe that AFT also offers new perspectives for rethinking the success and limitations of gated RNNs.
Dynamic Convolution. AFT is also related to dynamic convolution (Wu et al. 2019) when applied to auto-regressive tasks, where the reduced key-value representation can be interpreted as per sequence convolutional kernel. However, AFT operates in an extreme case where the dimension of the kernel is 1 along both the feature and spatial dimensions, again presenting superior efficiency.
5

---

Under review as a conference paper at ICLR 2021
cim 256
dim 1026
4000
Rurime (msttet
attention soth aetlattention
2000
mausol-mnar
1000
reformer
2000
4000
8000
5000
10000
2000
4000
6000
8200
10000
Contect Largth
Langth
dim 258
dim 1024
30
GPU Memmuy (SB)
20
atlanfon sell-athentio
10
causaHineer
reformer
2000
4000
5000
3000
10000
0
2000
4000
6000
8000
10000
Context Length
Longth
Figure 2: Comparisons of efficiency between models for forward and backward pass with batchsize of on a single GPU with 32 GB RAM.
30030
20000
attnttin
solf ottontin
10000
1000
2000
3000
4000
5000
9000
7000
8000
Context
Figure 3: Comparisons of efficiency between models for decoding sequences.
5 EXPERIMENTS
We conduct experiments on five tasks: unconditional image modeling (Sec. 5.2 language modeling (Sec. 5.3 machine translation (Sec. (.4). image super resolution (Sec. B.2 and point cloud generation (Sec. B.3 We focus on the causal mode of AFT, while leaving systematic evaluation of
the non-causal version as future work. Unless otherwise mentioned, all experiments are conducted on 8x V100 GPU machines.
5.1 EFFICIENCY
To support our analysis in Table we benchmarked an implementation of AFT on a single forward and backward pass on random data with batch size of We compared AFT with the self-attention from Transformers, a linear attention mechanism from (Katharopoulos et al. 2020), and Reformer Kitaev et al. 2020). For all of these, we used a base architecture with 12 layers, 8 heads (except for AFT, where there are no heads), feature dimension of either 256 or 1024, and context lengths up to 10,000. We used the code from the fast transformers library to perform the evaluations F] Results in terms of runtime and peak GPU usage are in Figure 2 Where data-points do not exist in the figure, the model exhausted GPU memory. We see from this that compared to Transformers and Reformers, AFT and linear attention require far fewer computational resources context increases. In addition, we see that AFT not as sensitive to the feature dimension as the linear attention, which expands the design space for feasible models. In all settings, AFT performs best. Additionally,

---

Under review as a conference paper at ICLR 2021
5.00
AFT-ral
475
All Firelu tran loes AFTireuu
AFT-relu colusotmma last loss AFTirel sotmar train loss AFT-sigmaideuhiin AF T-eig mongooodetttrooo loss PirelCNN++ baseline
4.25
bitsddin
3.75
325
125
Training epoch
Figure 4: Proof of concept experiments for AFT-relu2, AFT-relu and AFT-softmax, tested on CI- FARIO. All three versions train well with standard optimization settings. AFT-relu2 and AFT-relu perform similarly, while AFT-relu-softma and F--iigmidd-softma are more stable and yields significantly better results.
examine decoding speed in the same benchmark. Figure 3 increases linearly with the context length, and is faster than linear attention and Transformers with self-attention.
we see that the total runtime of AFT
5.2 UNCONDITIONAL IMAGE MODELING
In our first set of experiments, we consider the problem of image modeling by minimizing the negative log likelihood (NLL). Similar to Parmar al. (2018), we represent an RGB image as sequence of length H x W x 3, with H, W being the height and width, respectively. Each sub-pixel represented as a 256-way discrete variable. We use CIFARIO for image density modeling.
Feasibility study and choice of nonlinearities. We first conduct experiments validating the legit- imacy of AFT and its four nonlinearity variants. Our reference Transformer design largely follows that Chen et al. where transformer block consists of an attention layer (AFT layer in our cac) with residual connection and 2 layer MLP with residual connections. Layer Normalization (LN) Ba et al. 2016) is applied in "pre-act" fashion. We adopt learned position embeddings, and use set of shared token embeddings and prediction heads across RGB.
Our base architecture consists of 24 Transformer blocks, each with d-256 dimensional features. The hidden layer of the MLP per block has 4 dimensionality of its input. We use Adam, and follow a standard warmup learning schedule in Vaswani et al. (2017). We initial learning rate of 3 x 10-3 and weight decay of 0.1 applied to all linear transformations weights, and dropout of O.1.
We adopt simple data augmentation. During training, we first randomly flip each image horizontally, then add or subtract value in the range [-10, 10] from all its subpixels, and clip resulting pixel values to [0, 255]. We use cross entropy loss, and default batch size of 128 for 200 training epochs. We train three versions of AFT, namely AFT-relu2 (Equation 6.. AFT-relu (a, o relu), AFT-reluusofftmax: (a, = relu, aı soflmar)), AFT-ssigmid-ssofmaa (a,, sigmoid, ak softmax)), all of which use full contexts. We show the training and test loss curves in Figure 4 All versions of AFT trainable with standard optimization techniques for Transformers. In particular, AFT-relu2 performs slightly worse than AFT-relu, and both are significantly worse than AFT-sigmoid-softmal
also outperforms the strong PixeICNN+ baseline (Salimans et al., (017). Based on this observation, we use AFT-sigmoid- softmax as the default setting for all remaining experiments, unless otherwise mentioned.
AFT-relu-softmax and OFT-sigmoid-softma.
Comparing with the state of the art CIFARIO is a crowded benchmark for image autoregres- sive modeling, and we compare with few competitive baselines, as shown in Table 2 Note that

---

Under review as a conference paper at ICLR 2021
Table 2: NLL results on CIFARIO, evaluated by bits/dim, the lower the better. Speed and memory are measured during training time, with batch size of 32 across 8 VI00 GPUs. AFT achieve the state-of-the-ar result this setting, with much fewer parameters, better speed and significantly less memory.
Method
Train loss
Test loss
#params
Iters/Sec
GB/GPU
PixeICNN
3.08
3.14
PixeICNN+
2.92
PixelSNAIL
2.85
Sparse Transformer strided L=128, d=256
2.80
59M
Image Transformer local2d L=12, d=512
2.90
40M
1.61
22.3
AFT-full L=24, d=256
2.82
2.89
20M
2.15
9.5
AFT-hocal2d-hard L=24, d=256
2.81
2.87
20M
2.15
9.5
AFT-local-leared L=12, d=512
2.78
2.80
49M
1.68
11.4
AFT-local-ceamed L=24, d=256
2.75
2.74
29M
1.67
12,8
CIFARIO has an unrolled sequence length of 3072, which is already prohibiting to train full Trans- former with reasonable size. For example, for standard 12 layer 512 dimension and 8 head config- uration, the maximum batch size we can fit in our V100 node is only 16, which makes infeasible already. Our closest baseline Image Transformer (Parmar al. 2018), which restricts attention to test our AFT-local-learned and aft-local2d-hard variants, with
local2d windows of size of 256. the same window size, and the same architecture well as deeper but narrower one (24 layer and 256 dimensions), which are still fair comparisons. We also compare to Sparse Transformers (Child ct al. 2019), which restrains attention to sparse but global subset of context elements.
From 2 we see that all AFT variants outperform the Image Transformer baseline. Both AFT lo- cal versions better than the full counterpart, with AFT-local-learmed being significantly stronger than others. We also observe that the deeper but narrower architecture is more effective than the shallow but wide baseline. Our best model also achieves the state-of-the-ard result on CIFARIO 10 in this setting, outperforming a much larger Sparse Transformer model. Efficiency wise, we bench- marked Image Transformer against AFT variants on a V100 GPU node All our variants are faster than Image Transformer, while consuming only half of the memory
5.3 LANGUAGE MODELING
We apply AFT character level language modeling on Enwik8 (Mahoney 2011), which is another popular benchmark for auto-regressive modeling. We follow the standard preprocessing procedures and valiii//ivvviiaain/eestet splits in (Dai al. 2019). Our base Transformer reference layer 512 dimensional 8 head architecture with 2048 feed forward dimensions. For the first set of experiments, we use sequence length of 1024. Our training protocol as largely the same as before, except that we increase the weight decay to 0.5 and train for 100 epochs with batch size 128 in all experiments. We evaluate the AFT-local-learned variant with a window size of 32 and d, 256. We also compare to several efficient Transformer baselines, namely Reformer (Kitaev et al.) 2020). Synthesizer (Tay et al.) 2020a) and Linear Transformer (Katharopoulos et al. 2020). From Table we see that with the base L 12, d 512 architecture, AFT achieves the lowest training bits per character (bpc), indicating its high capacity. Its test performance is slightly worse than that of the basic Transformer, but outperforms all other three variants. The deeper and narrower architecture of AFT strikes the best balance across parameter, speed, memory and performance. Its test bpc is only 0.02 away from the full Transformer's, while only consuming third of the memory and provides 44% speedup. In the end, we have also trained the same architecture with a sequence length of 2048, which results in an improved performance both on the training and test set. This suggests AFT's ability to effectively model long range dependencies.
12
We hatch size of 32 which the largest batch size Image Transformer fit $ Fair comparison against Sparse Transformer is infeasible, as it relies on a set of advanced implementation tricks such as mixed precision and gradient checkpointing, whereas APT is implemented with standard Pytorch utilities full precision.
8

---

Under review as a conference paper at ICLR 2021
Table 3: Enwik8 results, measured in bits per character (bpc), the lower the better. Basclines com-
pared are Reformer (Kitaev et al. 2020), Synthesizer (Tay et al. 2020a) (it's best performing dense version) and Linear Transformer (Katharopoulos et al. p Speed and memory are measured during training time, with a batch size of 128 on a 8 V100 GPU node.
Method
T
Train bpc
Test bpc
#params
Iters/Sec
GB/GPU
Transformer L= 12, d=512
1024
0.977
1.137
39M
1.42
29.4
Reformer L=12, d=512
1024
1.04
1.195
35M
1.05
20.9
Synthesizer L=12, d=512
1024
0.994
1.298
42M
1.49
29.9
Linear Transformer L=12, d=512
1024
0.981
1.207
39M
1.46
10.6
AFT-local-leareed L d=512
1024
0.854
1.18
45M
1.85
11.3
AFT-local-leareed L=24, d=256
1024
0.973
1.157
32M
2.04
11.2
AFf-llcalllearredd 2=24, d=256
2048
0.942
1.139
45M
0.97
21.9
On the local window size. In all our experiments, AFT-local-lcareed demonstrates superior per- formance compared to other variants. In order validate its efficacy, we performed additional experiments wit the L 24, d 256 architecture, fixing everything but varying the local window size s. We show the results on where we see that both the training and testing bpc forms a U-shape w.r.t. the window size, with 32 achieving the best performance.
Table 4: Training and testing bpc w.r.t. the local window size for AFT-local-learned
Window size
0
2
4
8
32
64
128
256
512
1024
Train bpe
1.046
1.043
1.009
0.990
0.983
0.973
0.981
0.985
0.986
0.988
0.991
Test bpc
1.209
1.205
1.176
1.165
1.162
1.157
1.160
1.165
1.164
1.171
1.173
5.4 MACHINE TRANSLATION
As machine translation benchmark, show experiments with the WMT 2014 English to German translation task. The training set contains approximately 4.5 million sentence pairs. We compare against Transformer architecture baseline using the OpenNMT implementation (Klein al. 2017). For translation, the standard architecture is an encoder-decoder structure, where the encoder uses non-causal attention encode the input sentence. The decoder two different types of attention. The first, self attention, sequentially attends to the output translation as it is being generated token by token. The second attends to the translation and the context from the encoder.
In our experiments, we replace the multi-headed decoder self-attention blocks. We compare per- plexity (PPL), BLEU score, and efficiency between the Transformer base and AFT in Table In this task, we see that AFT performs on par with the Transformer. As expected for the small context size, typically around tokens, AFT does not show dramatic improvements in speed or memory.
Table WMT 2014 Aiglish-to-Germar Translation.
Method Transformer
Training PPL
Validation PPI.
Test BLEU score
tokens/sec
GB/GPU
4.38
4.06
27.32
54.7K
7.80
AFT (Ours)
4.34
4.10
27.70
54.4K
7.54
6 CONCLUSIONS
We have introduced the Attention Free Transformer that replaces attention with an efficient, easy- to-implement new operation. We have demonstrated strong results on challenging benchmarks, despite of the simplicity of our design. We believe that our model opens new design space for Transformer- like models, and will see impact in various areas where Transformers are applied.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint
1607 06450, 2016.
James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. works, arXiv preprint 1.1:v566 2016.
Quasi-recurrent neural net-
Angel X. Chang, Thomas Funkhouser, Leonidas Guibas, Pat Hanrahan, Qixing Huang, Zimo Li, Silvio Savarese, Manolis Savva, Shuran Song, Hao Su, Jianxiong Xiao, Li Yi. and Fisher Yu. ShapeNet: An Information-Rich 3D Model Repository. Technical Report 1rXiv:1512.0301 [cs, [R) Stanford University Princeton University - Toyota Technological Institute at Chicago, 2015.
Mark Chen, Alec Radford, Rewon Child, Jeff Wu, and Heewoo Jun. Generative pretraining from pixels.
Rewon Child, Scott Gray, Alee Radford, and lyy Sutskever. Generating long sequences with sparse transformers. CoRR, abs/1904. 10509, 2019.
100190....119990 10509
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, and Adrian Weller. Rethinking attention with performers, 2020.
Junyoung Chung, Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint 11X2v:1412.3555 2014.
Ryan Dahl, Mohammad Norouzi, and Jonathon Shlens. Pixel recursive super resolution. In ICCV, 2017.
Zihang Dai, Z. Yang, Yiming Yang, J. Carbonell, Quoe V. Le, and R. Salakhutdinov. Transformer-xi: Attentive language models beyond a fixed-length context. ArXiv, abs/1901 .02860, 2019.
Jacob Devlin, Ming- Wei Chang. Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: arXiv:1810.04805. 2018.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Lang Huang, Yuan, Jianyuan Guo, Chao Zhang, X. Chen, and Jingdong Wang. Interlaced sparse self-attention for semantic segmentation. ArXiv, ths/1907.12273. 2019a.
Zilong Huang, Xinggang Wang, Lichao Huang, C. Huang, Yunchao Wei, and Wenyu Liu. Ccnet: Criss-cross attention for semantic segmentation. 2019 IEEF/CVE International Conference on Computer Vision (ICCV), pp. 603-612, 2019b.
A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are mns: Fast au- toregressive transformers with linear attention. In Proceedings of the International Con- ference on Machine Learning (ICML), 2020.
URL littps://fleuret: pttpuureettog//ppper..ss..
katharopoulos-et-all aawaamlm11000001.111 pdf
Nikita Kitaev, L. Kaiser, and Anselm Levskaya. abs/2001.06 004451, 2020.
Reformer: The efficient transformer. ArXiv,
Guillaume Klein, Yoon Kim, Yuntian Deng, Jean Senellart, and Alexander Rush. OpenNMT: Open- source toolkit for neural machine translation. Proceedings of ACL 2017, System Demonstra- tions, pp. 67-72, Vancouver, Canada, July 2017. Associaton for Computational Linguistics. URL
https: hwtp:://wwww....b.org/ooo/oo
Juho Lec, Yoonho Lec, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invariant neural networks. In Interna- tional Conference on Machine Learning, pp. 3744-3753. PMLR, 2019.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV) December 2015.
10

---

Under review as a conference paper at ICLR 2021
Matt Mahoney. Large text compression benchmark, 2011.
Charlie Nash, Yaroslav Ganin, S. M. Ali Eslami, and Peter W. Battaglia. Polygen: An autoregressive generative model of 3d meshes. ICML, 2020.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. arXiv preprint arXiv: /82.v:/80..05511, 2018.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training.
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, and T. Lillicrap. Compressive transformers for long-range sequence modelling. ArXiv, abs/1911.055000 2020.
Prajit Ramachandran, Niki Parmar, Ashish Vaswani, I. Bello, Anselm Levskaya, and Jonathon Shlens. Stand-alone self-attention in vision models. ArXiv, abs/1906.05909 2019.
Aurko Roy, M. Saffar, Ashish Vaswani, and David Grangier. Efficient content- based sparse attention with routing transformers. ArXiv, abs/2003.05997, 2020.
Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelenn++: Improving the pixelcnn with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv:1701.05517, 2017.
Sainbayar Sukhbaatar, E. Grave, Bojanowski, and Armand Joulin. Adaptive attention span in transformers. In ACL, 2019.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models, 2020a.
Yi Tay, Dara Bahri, L. Yang, Donald Metzler, and D. Juan. Sparse sinkhorn attention. ArXiv, abs/2002.11296.. 2020b.
2020b.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, 2020c.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is you need. Advances in neural information processing systems, pp. 5998-6008, 2017.
Huiyu Wang. Y. Zhu, B. Green, H. Adam, A. Yuille, and Liang-Chich Chen. Axial-deeplab: Stand- alone axial-attention for panoptic segmentation. ArXiv, abs/2003.07883, 2020a.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. ArXiv, abs/2006.04768, 20206.
Felix Wu, Angela Fan, Alexei Baevski, Yann Dauphin, and M. Auli. Pay less attention with lightweight and dynamic convolutions. ArXiv, abs/1901.10430, 2019.
Zhen Zhu, Mengdu Xu, Song Bai, Tengteng Huang, and X. Bai. Asymmetric non-local neural networks for semantic segmentation. 2019 IEEE/CVF International Conference on Computer Vision (ICCV), pp. 593-602, 2019.
11

---

Under review as a conference paper at ICLR 2021
A APPENDIX
B LOCAL AFT
Algorithm 1: Pseudo code an efficient, in-place causal ATT/sottnna//AAFFoooooooollllll
Input: Query, key and value Q', K" V E R"xx, optionally context size {2", E N}.
Output: Causal AFT output Y E R Rd
K' exp(K') // new memory allocation
KV exp(X')) +V // new memory allocation
// it is not provided default to [To] iterations for: 1,, ..., "...,.m.milll wwmmmmm(mmnn [o_ (s))
stride 2
KV[sttidd: :1 = KV[stridd: :1 KV/-T-s :1 // in-place op
// now KV = Σ mm (exp( (x") [*] Vi
KK[stride:, :] = [Ktriidee: :] KK[-T-stride, :] in-place op
// now [ - (exp((()))) Vi
11 normal zze according to softmane and multiply with query relu(Q') KV/KK
Algorithm 2: Pseudo code of an efficient, in-place causal AFT-softmax-loaa22
Input: Query, key and value Q', K', V' E R Rww context sizes Sh, (2", No} Output: Causal AFT output Y E xd
KK exp(K') // new memory allocation
KV exp(K') + new memory allocation
// first aggregate locally across rows; pass if 8 2. do
for j = (og,,,]
stride 2
KV[; stride;, :1 = KV stride; + KV[; W-stride, :I // in-place op
// now KV = Σl (0, 1)) (x)(K") [),,,, Vi
KK[: stride: :] = KK[: stride; KK[: WWstride, :] // in-place op // now [K -0= (exp(K"))
// then aggregate locally across columns
for j = 1, (8g)(8)) do
stride 2
KV[stride;, :. KV]stride, : KV[:H-stride,, a // in-place op now
KV i..] Ei m0, - 21+1) kk -oo
Eiv
(exp(K') V") kwl, Vin dm
10
KK[stride; :. :] KK[ : KK[:H-stride.. // in-place op
// new KK ]..] 'n (0, 1))
E. 2ku=mor0, (exp (oop((K)))[kee kiwls Vin, dn
// incorporate contexts to the right
I idx = min(arange( ft, W-1) // a arange (W) [0, 1, W-1]
KV[I: :.:] KV[I: :.:] KV[:H-I, i, // in-place op
13 KK[I: ::: KK[I: 4]]] KK[:H-I, idx. in-place op
// normal ze according to softmar, and multiply with query
relu(Q') KV/KK
B. CIFARIO VISUALIZATIONS
Here we show the visualizations of our best performing model trained on CIFAR10 (with test bits/dim 2.81). In Figure 5 sample 32 test images and mask out the bottom half for each of them. We then use the model to sample the remaining pixels, one at a time. We see the model provides consistent completions for most cases.
12

---

Under review as a conference paper at ICLR 2021
(a) Masked images.
(b) Completed images.
(c) Full images.
Figure 5: Image completion with test examples.
B.2 IMAGE SUPER RESOLUTION
We also consider super-resolution task based pixel-wise image generation. Following Dahl et al. 2017 Parmar et al. 2018), we enlarge an 8 x sized image to 32 x 32.
We use CelebA dataset (Liu al. 2015) the benchmark. Our baseline model is the Image Trans- former (Parmar et al. 2018) with its encoder and decoder connected through the attention mech- anism. Both the ID and 2D local Image Transformer models have L 12 layers, d 512 and
13

---

Under review as a conference paper at ICLR 2021
8x8
20 local AFTHocaed
32x52
Table 6: Image super resolution results CelebA. Our AFT mod-
Tranaformrr
Truth
els outperform the PixelRecursive baseline (Dahl et al., 2017) in bits/dim (the lower the better), and show clear advantages param- eter efficiency and memory saving over Image Transformers (Pa-- mar al. 2018), with comparable or even better performance.
Train bits/dim
Dev bits/dim
Iters
GB
Method
L
Params
/sec
/GPU
Recerreeursive
2.81
54M
ID local Transformer
12 512
2.54
2.68
39M
2.98
9.5G
Upscaled images from baseline and our 2D lo-
Figure 6:
2D Transformer
12
512
2.43
2.61
1.45
21.2G
AFT-local2d (Ours)
32 256
2.39
2.59
25M
1.43
cal transformers CelebA
attention heads=8, and are trained under the DMOL (discretized mixture of logistics) loss for 200 epochs. We experimented by replacing the standard attention blocks in decoder with our AFT- localld or -local2d, keeping other modules the same. We follow similar training schemes, but with tuned dropout and learning rate. Evaluation performed terms of NLL measured in bits/dim, with the sampling temperature fixed at 0.8.
Table O shows results of our best AFT-local2d model (context size 16 16), in comparison the PixelRecursive baseline and 1D/2D local Image Transformer models. Note the large consumption of model capacity and memory from 2D local Image Transformer, while our AFT-local2d shows clear advantages with loss in model quality (see F for visual comparison).
Figure shows more samples from different models trained on CelebA face images.
8x8 Input
1D local Transformer
2D local Transformer
AFT-local2d
(Ours)
32x32 G. Truth
Figure 7: Upscaled images from baseline 1D/2D local Image Transformers (Parmar et al. 2018 and our AFT-local2D model trained CelebA.
B.3 POINT CLOUD GENERATION
In addition to images and text, explore modeling point clouds randomly sampled from objects the ShapeNetCore v2 dataset (Chang et al. 2015). Each point cloud consists 2048 points. Following Nash et al. 22020), the points uniformly 8-bit quantized based on their positions. Our preliminary results of point cloud generation shown in Table and examples of generated point clouds shown Figure We that model able to generate self consistent objects with fine details and great diversity.
sorted into sequence in the order of z, y and then
14

---

Under review as a conference paper at ICLR 2021
Table 7: Results on ShapeNetCore v2, evaluated by bits/dim, the lower the better.
Class
Train bits/dim
Test bits/dim
All classes
3.67
3.87
Airplane
3.07
3.14
Chair
3.67
3.96
Car
3.64
3.99
Figure 8: Point clouds generated by AFT trained on airplane point clouds.
15


Reviews:
[{"review_id": "ROJLnBrEXuC", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper introduces a method to replace qkv-attention by a simpler, efficient building block. This is done by element-wise multiplication of a query representation with a compressed kv-memory. Per-channel attention pooling is used to compress the kv-memory. The model is derived from a softmax-free version of self-attention. The results show good performance on a couple of standard image and language modeling tasks while occasionally exhibiting favorable training speed. The results are largely on par with transformer baselines.\n\nI have a couple of major concerns with this paper:\n\n1) The derivation: The model is derived by using relus on the QK dot-product andthen simplified in Eq. 4, which basically makes the resulting model different from the starting point (Eq. 2-3). The final model is further changed by a applying a softmax over the keys, channel by channel (Eq. 5). So I am not sure how the resulting model actually still relates to the original formulation in Eq. 2. That leaves me with the impression that the derivation just exists to establish a connection with standard qkv-attention which is a bit hand-wavy. The results (Figure 2) even suggest that without a per-channel softmax over the input elements, the method doesn't work well. The arbitrariness of the derivation is exemplified further by the fact that one could have similarly started from using no non-linearity at all after the dot-product, which directly leads to Eq.4. The non-linearities on K and Q could be arbitrarily applied before the dot-product.\n\n2) The efficiency comes at a cost of a strong memory bottleneck as we basically pool the entire memory into one fixed state on hidden size. That won't scale well to larger inputs. The conducted experiments are mostly on smaller scale settings (small images, standard text lengths) which reinforces my impression that this approach won't help solving the efficiency issues of self-attention. Using larger receptive fields (though costly) typically leads to better performance in standard self-attention. Here, however, due to the strong bottleneck results are getting worse after a point (Table 1). This makes also sense because the memory pooling imposes a very strong information bottleneck, that is, much of the memory has to be forgotten.\n\n3) There should be at least some controlled baselines from related work that tries to eliminate the self-attention bottleneck with similar compression techniques, e.g. Linformer, Sinkhorn Transformer, Compressive transformer, Performers.\n\n4) The potential efficiency gains are never put to practice, that is, the authors don't show any application of the model to very large input sequences.\n\n\nOther comments and questions:\n\n- The work reminds me of dynamic convolutions [1], which compute depth-wise convolution kernels dynamically based on the current context. Here we compute a dynamic depthwise 1x1 convolution on the Qs, based on the context around each Q. I think this connection might be more closely connected to the proposed model than attention.\n\n- Why were the models not trained till convergence in on WikiText-103?\n\n\n[1] https://arxiv.org/abs/1901.10430", "extra_scores": {}, "extra_text": {}}, {"review_id": "Oi66x-ocNlP", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper suggest an alternative to the Multi-Head Attention (MHA)  operation, which is one of the core elements in Transformers models. The proposed alternative is targeting the non-linear soft-max operator (in the MHA) and suggest to replace it with the \"relu\" operator. After doing so they could reformulate the new attention mechanism as a O(T) operator instead of the original O(T^2) operator (where T is the context size).\n\nArguably, the MHA is one of the important components of the transformers architecture and reducing its memory and time complexity is crucial to increasing the training batch-sizes and the usage of more context.\n\nThe paper is nicely written and presents a comprehensive experimentation section, ranging over several machine learning benchmarks in computer vision and NLP.\n\nStrong points:\n- simple solution\n- comparable results with reduced memory and latency\n- the paper is clear and nicely written  \n- comprehensive experimentation section\n\nWeak points:\n- moving from AFT-relu to AFT-softmax is is not sufficiently motivated (only empirically) i would expect more experimentation to clear this point\n- not all experiments show improved or comparable results (for example table 5)", "extra_scores": {}, "extra_text": {}}, {"review_id": "EoZqv3YU_BN", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper introduces the Attention Free Transformer (AFT), an alternative to multi-head attention (MHA) operation in Transformer. While the motivation of the authors is to replace MHA with more cost-efficient operation, it is not clear whether the proposed method is the better alternative.\n\nPros:\n1. AFT shows better asymptotic space and time complexities than MHA.\n2. The implementation of AFT allows for faster training with larger batches.\n\nCons:\n1. Theoretical analysis is conducted for the extreme case of num_heads=hidden_dim and ReLu non-linearity. It is not clear how to generalize them to more practical cases with num_heads<hidden_dim and SoftMax non-linearity. There is a missing link between the theory (and motivation arising from it) and the best-performing implementation (AFT-softmax).\n2. AFT-softmax does not fully complies with the title of the paper as the proposed operation contains aggregation via softmax. Also, despite the claim that AFT can \"be readily adopted as a plug in alternative to Transformers\", the architectures from the experimental section also use vanilla MHA blocks in addition to AFT. Thus, it is an exaggeration to say that the Transformers evaluated in the paper are attention-free.\n3. Language modeling experiments on WikiText-103 draw an ambiguous picture. Baseline Transformer implementation has large positive difference between train and val/test PPL at 70k iterations which decreases as the training progresses. For AFT models, on the other hand, this difference is negative which might suggest that they have already overfit at 70k iterations and they will never reach the resulting performance achievable by the baseline. The plot with train&val PPL / number of iterations for those experiments would be more informative than the table.", "extra_scores": {}, "extra_text": {}}, {"review_id": "7jhtW9oFSd", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes an efficient transformer variant by replacing softmax in the self-attention layer with a RELU activation and arranging the computation using element-wise products and global/local pooling. This reduces complexity to linear complexity in the non-autoregressive case and log-linear complexity in the autoregressive case. The evaluation shows that it can reach the performance of a vanilla transformer in most of the examined tasks while having fewer memory requirements in general. \n\n**Strengths**: \n\nThe paper is reasonably well-written and clear for the most part. The problem of scaling transformers to longer sequences is an important one since transformers cannot deal otherwise with long sequences due to their quadratic complexity. \n\nThe proposed idea is interesting and reminiscent of recent methods that re-arrange self-attention computation using kernels albeit it differs in the way the computation is carried out. This one is simple computation-wise and does not aim to approximate the original computation in any way. \n\nThe evaluation performed on multiple tasks shows that the proposed approach can reach the quality of a vanilla transformer and be more memory-efficient. \n\n**Weaknesses**: \n\n(1) The motivation of AFT and the positioning with respect to prior work were somewhat weak. The introduction does not acknowledge recent efforts towards efficient transformers and what is the unique contributions of this work. What are the benefits of AFT compared to recent established efficient transformers such as Sparse Transformer (Child et al., 2019), Reformer (Kitaev et al., 2019), or Linear transformer (Katharopoulos et al., 2020)? It is unclear why one should prefer the proposed variants over existing ones both from theoretical and practical perspectives.  Related work states some previous efficient transformers without any individual discussion about their merits or limitations in comparison to AFT. \n\n(2) One major limitation that stands out from the experiments, despite their size, is that there is no head-to-head or controlled comparison with a previously established efficient transformer such as the ones mentioned above. The results compared to Sparse Transformer given in Table 3 are not directly comparable since the model size and design are quite different.  In brief, it is not very clear what are the practical benefits compared to previous efficient alternatives. \n\n(3) The memory benefits are not reflected or they are not as important when looking at the quality achieved in the tasks where a speed-quality trade-off was reported.  In language modeling,  AFT has higher perplexity (even when it uses a much larger number of parameters) which makes the memory benefits less interesting. In MT, AFT reaches the performance of the baseline but then the efficiency benefits are not present. So, I am curious is it the same in the two former tasks when comparing to the vanilla transformer? Under what circumstances we should expect AFT to reach vanilla transformer performance and still offer clear efficiency benefits when using the same setup?\n\n(4) In terms of training speed, AFT is generally slower than the vanilla transformer when the form reaches the same quality as the latter. Also, it is especially slower when the depth is small in Table 2 (~30% with 12 layers). Could the authors elaborate a bit on why that happens?  Moreover, it would be useful to show in Table 2 what is the quality (NLL or bits/dim) achieved by each model because it's hard to tell how good the speed-quality tradeoff is. \n\n(5) Recent studies have shown that it is possible to speed up inference time using efficient transformers (see above). What is the benefit of AFT during inference time? ", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The new non linearity proposed in this paper present interesting observations and improvements on image and text datasets.
However, reviewers point out that there should’ve been more comparisons to other efficient transformers and on more datasets.
The speed improvements are also not clear.
I’d encourage the authors to revise and submit in the future.
Decision: Reject