Efficient Attention via Control Variates
Authors: Lin Zheng
Jianbo Yuan
Chong Wang
Lingpeng Kong
Venue: ICLR 2023 notable top 5% (2023)
Keywords: attention mechanism
transformers
random features
control variates
importance sampling

Abstract: Random-feature-based attention (RFA) is an efficient approximation of softmax attention with linear runtime and space complexity. However, the approximation gap between RFA and conventional softmax attention is not well studied. Built upon previous progress of RFA, we characterize this gap through the lens of control variates and show that RFA can be decomposed into a sum of multiple control variate estimators for each element in the sequence. This new framework reveals that exact softmax attention can be recovered from RFA by manipulating each control variate. Besides, it allows us to develop a more flexible form of control variates, resulting in a novel attention mechanism that significantly reduces the approximation gap while maintaining linear complexity. Extensive experiments demonstrate that our model outperforms state-of-the-art efficient attention mechanisms on both vision and language tasks.

Published as a conference paper at ICLR 2023
EFFICIENT ATTENTION VIA CONTROL VARIATES
Lin Zheng'* Jianbo Yuan²
Chong Wang"
Lingpeng Kong' Apple Inc.
'Tee University of Hong Kong
BayteDance Inc.
(1 (1zheng2, 1pk] @ess hku. hk j jianbo. yuan@bytedanna com
mr. mr. .t.chonsswanggappleeccom
ABSTRACT
Random-feature- based attention (RFA) an efficient approximation of softmax at- tention with linear runtime and space complexity. However, the approximation gap between RFA and conventional softmax attention is not well studied. Built upon previous progress RFA, we characterize this gap through the lens of control vari- ates and show that RFA can be decomposed into sum of multiple control variate estimators for each element the sequence. This new framework reveals that exact softmax attention can be recovered from RFA by manipulating each control variate. Besides, it allows us to develop a more flexible form of control variates, resulting in novel attention mechanism that significantly reduces the approximation gap while maintaining linear complexity. Extensive experiments demonstrate that our model outperforms state-of-the-art efficient attention mechanisms on both vision and language tasks
INTRODUCTION
Random-feature-basee attention (RFA, also known as Performer; Choromanski et al. 2021 Peng et al. b02Tb) is established fast approximation the conventional softmax attention mechanism (Bahdanau et al. 2014 Vaswani et al. 2017). which successfully scales Transformer models to processing much longer sequences (Choromanski al. 202TT. At its core is the usage of random features (RF: Rahimi & Recht, 2008) to linearize the exponential kernel in softmax attention, which reduces the computational cost from quadratic to linear runtime and space complexity. Despite its efficiency, recent studies have pointed out that such approximation suffers from substantial performance degeneration (Xiong et al. 2021a Zheng et al. 2022b).
In this work, we generalize the formulation of RFA via control variates (Owen) 2013), which characterizes the approximation gap between RFA and softmax attention in theory. We first show that RFA can be decomposed from global approximation over the whole sequence into sum of local control variate estimators, each of which is applied to an individual element in the sequence. Under this formulation, RFA is equivalent employing the same coefficient for all control variate estimators to scale their variance isotropically ($3.I). Besides, we prove that if we optimize the coefficient of each control variate minimize the estimation variance individually, RFA estimation becomes exact, that is, softmax attention is recovered with zero bias and zero variance (3.2p.
Our key observation is that such formulation reveals a localized perspective of the RFA approximation. Instead of directly seeking a better estimate over the entire sequence, we can break down the problem into smaller problems that aim improving the approximation for each subsequence ($4). The control variate estimator for each subsequence can be tuned separately and combined to yield better estimation, which provably reduces approximation error in the global sense (S4.1p. Nevertheless, one caveat is that as the number of sub- nub-problems increases, the approximation gap will be reduced but at the expense of higher computational complexity. For instance, if we optimize the control variate for every single element, softmax attention would be recovered as desired but with quadratic complexity. To attain good trade-off between approximation quality and efficiency, we develop a new Efficient attention via control VAriates (EVA) that implements this divide-and-conee and-cuquuer strategy efficiently. In EVA, the sequence partitioned into a fixed number of disjoint subsets. For the subset
"The majority of this work was done while these authors were at Bytedance. Our code and models are available this link

---

Published as a conference paper at ICLR 2023
that might bear the highest correlations to the query, explicitly optimize the control variate for each element, which recovers exact softmax attention probabilities; while for the others, the control variate coefficient is shared locally among all elements within the same subset. The resulting attention mechanism not only highly effective but also runs with the same computational complexity as RFA (84.2). Extensive experiments both language and vision tasks demonstrate that EVA outperforms the state-of-the-ar efficient attention methods $
2 BACKGROUND
2.1 SOFTMAX ATTENTION MECHANISM
Assume there exist set queries {q, 1.W and M key-value pairs = .k..... kasl and V Vi...., V. where queries, keys and values are all d-dimensional vectors. The softmax attention mechanism (Bahdanau et al. (2014) Vaswani et al. (017)) is defined as an average over the value vectors weighted by the dot-produce similarities of the queries and keys. For the n-th query, the attention mechanism outputs
M
exp (qm km)
SoftmaxAttn(q... K, V)
(1)
m=1 m'=1 exp (q km))
2017). we have M = NN, which results in quadratic computational complexity since we have to compute the similarity for each query-key pair explicitly.
In the case of self-attention (Lin et al. 2017: Vaswani et al.
2.2 RANDOM-FEATURR-BAASD ATTENTION WITH SELF-NORMALIZED IMPORTANCE SAMPLING
Recently, Zheng al. (2022b) identifies that softmax attention (Equation can be written as expectation over an attention-lik like aggregating function,
M
exp (qn km)
SoftmaxAttu(a,. K, V)
Vm Buwpp(() [f[[]]]
(2)
mm=m1 exp (q, km))
where
fn(w) = Em E(q.. ((kk) Vm Pn(w)= = N(w 0,1; D. E(q,, w) T((k,w))
(3)
Em= E(g,, (kkm'
Z
Here E(:,)) is the randomized mapping defined in such way that exp (q. km) Ew-N/o..) (E(q,,w) [(km::]]] (
and 7 = Σm= exp (q, km) denotes the normalizing con- stant of distribution Pn- Throughout this paper, consider the positive randomized mapping {(x,w)= exp (w 1 [ =1x(() (Choromanski et al. 2021) unless otherwise specified.
Random-Feattnn Fee--Feettur-ssseeeeee Attention (RFA) methods (Choromanski et al., 2021 Peng et al. 202 lb) can be interpreted as performing self-normalized importance sampling (SNIS: Hesterberg 1995) to
approximate Equation 2 (Zheng et al. (2022b). In SNIS, one draws Monte Carlo samples from some
proposal distribution (ww) instead of the true distribution (w) and estimates the target expectation
as Esppp (w) [f[[W] = Eung(w)
(w(ww ((w) fm (w)
Prrfe)
(w(w))
where Wh...,,W5 ~ q(w).
De
Vanilla RFA amounts constructing the SNIS estimation with ((w)) N(w,,,,)) The SNIS
representation also turns out equivalent to the more established form of RFA,
E Prr(w.) f(ws) ===l glwa)
Σ (b w) dbkk,, W)mm M Im
RFA( K, V)
(4)
Pm((w)) ==1 ((w,)
Φlan W) of ((km/,w)
where the random fandoufeature, denoted by Φ(x, w) = 1/V3L[(x,,)),..... E(x,ws)) ws)l is proposed to approximate exponential kernels in its original motivation (see Appendix AA for detailed review).
2

---

Published as a conference paper at ICLR 2023
2.3 CONTROL VARIATES
Control variates aim to reduce the estimation variance of an expectation E [g(w)]. Assuming our original RFA estimation is g(w) E IR and there is some control variate h(ww) E R with known expectation E [M(w),, we can employ the control variate l((w) with the coefficient B E R as follows,
(5) Note that the resulting estimator remains unbiased since E [))]] = E [g((و] 3E [k(w]] + 3E [h(w)] E [g(w]]. However, the estimation variance be largely reduced if g(-) and the scaled control variate 3h(w) are positively correlated (Owen, 2013).
g(w) - ((w) - 3b(ww) + 3E [k(W]]
3 DISSECTING RFA WITH CONTROL VARIATES
In this section, we first go through the connections among RFA, importance sampling, and control variates, revealing decomposed formulation of RFA [. and then quantify the approximation gap between RFA and softmax attention $3.2 from these connections.
3.1 RFA AS A SUM OF LOCAL. CONTROL VARTATE ESTIMATORS
As shown in Equation 4 RFA estimation considers all key-value pairs and produces a global approximation over the entire sequence. In contrast, our work develops a decomposed representation of RFA based on the recent advances in SNIS (Vlassis et al. 2021), which indicates that an SNIS estimate is asymptotically equivalent to a control variate estimate (the detailed derivations is deferred Appendix B.2). In particular, we have
Σ s=1 ((wa))
S
Pa(we)
(w))
1 s
Pa (n(we) f(w,) Ex =x=1
Pro Pralue) (wss) alwa)
1
Pn(ws))
-
> aml 1 g(wo,)
(or.
8 s=1 g(ws)
Pa(wwa
8 ==1 g(we)
g(wo)
g(w) B(ww) (h(w) - T. [m(w)]) = g(w),
6
where ((w) 1
Pa palwsc ((ws) is our base estimate, h((w) = 3 > Polwel q(w,)
is the control
glas.)
variate with control coefficient ((w) = (Ei
Palwed S(ws)) (Es ((ws)
Pn (w(w) ((w,)
g(w) h(w)
We now examine the formulation of g(-) and in(.) in the context RFA. According to Equation 3 (wss))
8 f(wa) = Σ alws)
M
((w) = 5 ((Ws)
Σ
E(g,, We ( (km wg )
s=1
m=1
5 (w,) g(wa)
h(w)
Σ a(ws)
Σ E(g,, ( Wa),
m=1
2g(ws) collects terms that is constant w.r.t. queries, keys, and values. Our key observation is that by changing the order of summations, both g(-) and h(-) can be decomposed as g(w) = EM (m) and h(w) = EM Am (w) respectively, where
where a (ws.) N(w...0.I)
S Jm (w) Σ (ll))))(an) Wa (kmm Ws) Vm Iim (w) Σ ((ws)E(( (s) (ickmmmws) wa).
$
As a result, we can decompose the entire RFA estimate in Equation 6 into a summation of M control variate estimates following
g(w) = g(w) - -B(w) (h(w) [A(()))
M
Imn (w)( -B(w) ((2)) M
A
(wm)
Σ
hm hm((w) (w))
- E
Σ m=
hmm (w))
m=1
m==
M
Im (w) -B(w) (Am (ww - E (w)))
gm (w))
m=1
3

---

Published as a conference paper at ICLR 2023
Here (m) 9m (w) B(ww) (w) E [km (w]]) denotes the corresponding control variate estimator of the m-th key-value pain and ((ww) is the coefficient shared across the entire sequence.
3.2 OPTIMIZING COEFFICIENTS IN RFA LOCALLY RECOVERS SOFTMAX ATTENTION
Based on the decomposition of RFA in Equation 7 we have one local control variate attached to each key-value pair. To see the benefit of such decomposition, we demonstrate that softmax attention is equivalent to associating each control variate with a locally optimized coefficient 3m in RFA.
Proposition 1. Let (m((w) 9m (w) 3m (hm (w) E [hm (w)))) We denote the variance of some estimator g(w) as Var [g(w)] = Cov (g(w). g(w)]. Then the optimal Bm, that minimizes Tr (Var [mm (w])) (i.e., the sum variance over all dimensions) is of the form
9m (w) hmm (w)
F = arg min Tr (Var [T V
(8)
B
Furthermore, by letting 3m = B for all m = 1, .......,. we have Tr (Var Im (w)) = 0. As result, Tr (Var [[()]]] 0 and thus RFA(q,. K, V) = [(w) SoftmaxAttn 9n K, V).
The proof is deferred to Appendix B.4
This proposition implies optimizing 3m, for each key- value pair in the decomposed formulation of RFA recovers the exact softmax attention. It not only characterizes the theoretical gap introduced by RFA but also sheds light on how to improve RFA towards softmax attention from a localized perspective. Furthermore, delineates the trade-off
between estimation quality and computational costs. On the one hand, if we use a distinct Brg for each estimator, we could achieve perfect estimation, albeit at the expense of computing exp q, km for every query-key pair explicitly with quadratic time and space complexity. On the other hand, a single shared coefficient is employed, it degrades conventional RFA, where all the control variate estimators can be merged and computed together in linear complexity (Choromanski et al. 2021 Peng et al. 2021b: Zheng al. 2022b).
4 EVA: EFFICIENT ATTENTION VIA CONTROL VARIATES
In this section, we demonstrate that the control variate formulation offers natural way to improve RFA with a finer-grained treatment over control variates. We describe the improved efficient attention mechanism EVA in $4.1 and its practical implementation in S4.2
4.1 CONTROL VARIATES WITH LOCALLY SHARED COEFFICIENTS
We denote [M] = (1,2, M) as the set of all key-value indices. Instead of employing the same co- efficient for all control variates as in RFA, propose to partition [M] into C subsets PI, Pa,..... ,c and allocate a locally shared B_ for each subset Pa. For all B_ and their optimum 3r, for each token, define the weighted mean squared error (weighted MSE) Σ ΣmEr am || 3t,
1. To see the benefit of partitioning, we demonstrate that there always exists some (,,.0_. that achieves lower weighted MSE than any globally shared coefficient (see Appendix B.5 for formal argument). The next question is how to determine (3.) According to Proposition IJ a natural choice to adapt the optimal coefficients (Equation ( the
where am > 0 and Ee= Σmep,
case of partitioned subsets. We justify this choice by proving that it is also optimal in minimizing the MSE above weighted by the true attention probabilities.
Proposition 2. Suppose U is a set of key-value indices, B:, is the optimal coefficient for each m E U as defined Proposition 7 and PI, 1,,22.... . Pc are an arbitrury partition of U, where each subset Pe is associated with distinct B_.. We consider the following weighted mean squared error,
exp (qm km
J(B
3c)=
Σ
||3 -
(9)
Emeeu exp (q,, km))
0=1 MEP.
Note that the expectation of individual control variates (mm( is still in closed form as E [wm[[w]]
=
exp(q. / The derivation can be found in Appendix B.3
4

---

Published as a conference paper at ICLR 2023
Then for each I,.... C we have
E mEp Im [w][w]]] =
3 == arg min Ae
(10)
E EEep, hm (w)]
As a consequence, with 3. = 3r. the partition scheme must achieve lower weighted mean squared error than any globally shared B, that is, J(31 = .3......c = 3c) S J(3 B...., = B).
The proof can be found in Appendix B.6 Apart from measuring the squared errors for all coefficients, Equation
9 also governs the significance of each error by its corresponding softmax weights, which attains closer alignment with true softmax attention. Therefore, this proposition implies that is much easier for the partitioned control variate estimators obtain coefficients closer their optimum while faithfully respecting softmax attention. The optimal coefficients 3 could be estimated via
Monte Carlo samples as B² x (w) = (Emep Im (w)) / mec hm (w)). which is a widely adopted strategy in the control variate literature (Wang et al. 2013 Owen 2013). The resulting estimator for each subset Pe takes the form
9m (m(w) - (.(w) B.(w) exp(q, kx)(())k))
Emep, exp(q km)
Σ mEPc
B.(w).
(11)
Z
Partially Optimized Coefficients. Given the optimality of using a separate coefficient for each
key-value pair, we could further improve the estimation by selecting some subset E [ [M] and
employ 3m = =Bmm = V m for each m E E. Without loss of generality, we assume E n Pa = Ø for
all c = I...., C and [M] = (U Po) U.E. According to Proposition for each m E E we have
Im (w) 9m (w)) Bn,, (m + B, exp(q.,km) exp(q, km)v
(12)
Z
We choose E by running additional sparse attention mechanism (e.g., local window attention (Child et al.) 2019) or Reformer (Kitaev et al.) 2020)), which tend to select tokens that are more relevant to the query in sub-quadratic complexity. Since estimates these critical tokens exact, this strategy not only reduces the overall squared error (Equation 99) but also produces a more informative context for queries, which often translates into better empirical performance. Combining Equations 12 and I1 together, we obtain an improved Efficient attention via control VAriates (EVA),
EVA(q.. K, V) = -(ww) Σ (w) + Σ Fm((w)
MEE
m&E
exp( km)
Emep, exp(a..km) km)
Bw(w).
(13)
Z
Z
MEE
Comparison with Vanilla RFA. EVA and vanilla RFA can be re-written in similar way (see Appendix B.7 for detailed derivation),
Σ MM mm= Im (w) M
RFA(q.. K, V)-
(14)
Em hmm (w)
exp(q,) km)
I. (w) hm (w)
Emer exp(q. T. exp(q..kmm) Emer Im (w)
EVA(q.. K, V)
(15)
Z
Σmep, hm (w)
MER
Intuitively, we can think of EVA as a calibrated version of RFA. Instead of directly computing and aggregating the random feature approximation for all tokens as in RFA (Equation 14), EVA (Equation 15) first constructs local estimation for either a single token (m E E) or a subset (e.g.. Po). and then corrects these approximations by their corresponding true attention scores (e.g., Emep cxp km) for P.). These adjusted local estimates are finally aggregated and globally normalized. Thanks to the decomposed representation of RFA, strategy in a principled manner, which imposes finer- finerngrained control on the whole estimation accuracy and enjoys increased approximation fidelity.
realize this divide-and-conguee
5

---

Published as a conference paper at ICLR 2023
Table Classification accuracy on ImageNet 1 in comparison to different RF-based approxima- tions. vanillla PVT-v2-b3 (Wang et al., 2021b) uses a convolutional kernel to downsample key and value vectors, resulting in fewer FLOPs but with significant performance degradation.
DaiT-Tiay FLOPs
DeiT-Small PLOPs
PVT-x2-b3
Model
# Param.
Top-I Acc.
Param.
Top-1 Acc.
Puram.
FLOPs
Top-1 Acc.
Local
5.7M
LIG
67.10
22.0M
4.3G
74.06
36.0M
7.2G
83,34
Performer
5.7M
1.2G
65.92
22.0M
4.4G
74.29
36.0M
8.2G
82.40
LARA
5.8M
12G
71.48
22.2M
4.5G
79.48
39.9M
7.7G
83,47
EVA (Ouns)
58M
1.2G
73.00
22.2M
4.4G
80.65
36.
7.4G
83,71
Softmax
5.7M
1.3G
72.98
22.0M
4.6G
80.36
45.2M
6.9G*
83.14"
4.2 PRACTICAL IMPLEMENTATION
According to the formulation (Equation 13 of EVA, the terms within E could be computed efficiently due to its limited size; however, the partitioning requires computing Σmep, exp q, km) explicitly for each subset, which again builds up to quadratic computational complexity. As discussed above, Emep, exp(q, km) serves as weight to correct the contribution from each subset Pa. In this
regard, we propose approximate such control by Σmep, exp(q, km) R exp(q, k.), where k, is an adaptive vector summarizing the information of all keys belonging to Pe (see Appendix for more details). Such heuristic not only avoids computing the exponential dot product of each query-key pair explicitly, but also induces a fast approximation of the normalizing constant,
Z Σ exp(q km) + Σ E
expla, km)
exp(q, km) + Σ exp(q,,k.) k.). mEE
Σ
mEE
c=1 mEP.
Equipped with these results, our EVA estimator (Equation 13 can be reduced as follows,
EVA(q,, K, V) EmER exp(q km Vm + Ee
===1 exp(q, (..
(16)
Emer exp(q)) km)
+ > exp(( q k.)
Parameterization Details. We define E in the same way simple block-wise local attention (Xiong et al. 2021a). The input sequence is first chunked into multiple blocks (or 2D windows for images), and each query 9n is associated with specific E" that only contains tokens within the same block as the query. For the remaining indices [M] \ E", we evenly split it into C' contiguous chunks (PT. ... Pe). Note that we add the superscript n here to denote the dependence query position; however, for notational brevity, we omit the notation when there is no ambiguity. The pseudo-code of EVA is provided in Algorithm of Appendix. More implementation details, including the definition of k. and (w) in Equation 16 are deferred to Appendix C
the
Extension to Autoregressive Modeling. The decoder (or causal) self self-attention, where each query can only attend to previous tokens, the key ingredient Transformer- based generative modeling (Vaswani et al. 2017 7 Brown et al. (2020). We demonstrate that it is straightforward extend EVA to support such auto uuto-regressive modeling with few modifications. Thanks to the decomposed formulation of EVA, only need to incorporate two triangular mask matrices into the computation, which climinate the information from future singletons m E E' and entire future subsets P. respectively. Unlike previous RFA methods, which are slow during training due their recurrent computation (Choromanski et al., 2021] Peng et al. (0215). the resulting causal variant remains highly efficient. More details can be found in Appendix D including pseudo-code Algorithm 2
EXPERIMENTAL RESULTS
In this section, evaluate proposed method various tasks, including image classification ($5.T). language tasks [3.22), and Long Range Arena benchmark (Appendix (F). Details experimen- tal protocols and baselines be found in Appendix

---

Published as a conference paper at ICLR 2023
Table 2: Image classification accuracy on Table 3: Masked Language Modeling Perplexity on the Books3 validation dataset.
ImageNet 1k dataset with DeiT--Tiny-77
Model
Param.
FLOPs
Tup-1 Acc.
Model
Y Param.
11.05s
Perplexity
Performer (Cavoomeneii 2021
5.7M
1.5G
67.14
Performer
Cheromands
2001
126M
Nystronfommt Niigg Mi. 1b
Scatterbraia Chea 2041
midcee ealle 201
5.7M
5.7M
5.2G
1.8G
74.20
73.50
LARA (2002 auseu
Local attenter
Retormer Ammeer
Whirerse
136M
126M
120M
126M
205G
1980
183G
4,25
4.27
4.39
LARA Zheng
5.8M
4.6G
75.02
Cambinrr
136M
1870
4.13
Combiskr
Long Short 230
142M
218G
4.01
Long-Shon Zia 0
6.1M
5.0G
76.41
EVA (Orm)
EVA (Ours)
5.8M
4.6G
76.67
RVA-4096 (Chrs
1870
3.73
Softmax
5.7M
7.0G
77,16
126M
3.74
Table 4: BLEU scores on the test set of
Table 5: Validation (Val.) and Test perplexity (PPL) on Wikitext 103. 256/480 indicate evaluation context window sizes. numbers are due to Kasai al. (20TT)
WMT14 4 De."" numbers are taken from Zheng et al. 2022b).
Model
# Param.
BLEU
Performer-128
60.92M
23.5
Model
Params.
Val.
256
Test
Val.
Test
LARA-16"
60.96M
26.4
LARA-32
60.96M
26.8
Softaaxx
449M
17.9
18.5
LARA-64"
60.96M
27.0
ELU"
449M
22.0
22.8
REAT
449M
20.4
21.3
BVA-16
60.96M
27.2
EVA-32
60.96M
27.3
EVA (Ours)
450M
17.9
18.6
17.7
18.3
EVA-64
60.96M
27.5
Softmax
247M
18.8
19.5
18.4
19.1
Softmax
60.92M
27.5
BVA (Ours)
247M
18.8
19.4
18,5
19.1
5.1 IMAGE CLASSIFICATION
We explore the ability learn visual representations for different attention mechanisms in vision transformers (ViTs:: Dosovitskiy et al. (022). In particular, we replace softmax attention used in ViTs with its efficient variants and evaluate their performance on the ImageNet dataset (Deng et al., 2009), which contains over 1,280K and 50K images of 1,000 classes for training and validation splits, respectively. For the transformer model, we consider both a plain ViT (DeiT: Dosovitskiy et al.) 2020 Touvron al. T) and pyramidal VIT (PVT: Wang al. 2021b) to test the performance. The former maintains the same sequence length (which is set to 196 by default) across all transformer layers, while the latter processes much longer sequences (up 3136 tokens) at early layers and progressively reduces the sequence length to form a hierarchical structure. Detailed experimental settings could be found in Appendix E.2
Results. We first compare the performance of EVA against our main baselines on the standard ViT architectures. As shown in Table T EVA significantly improves the performance of previous RFA approaches (including Performer (Choromanski et al. 2021) and LARA (Zheng et al. (0226)) and local attention by large margin, and even outperforms the conventional softmax attention. We then consider a more challenging setting, where the plain architecture DeiT-Tiny is used but the sequence length is scaled up 784 (denoted DeiT-Tiny-784), We compare EVA against other attention variants in this setting and report the classification results in Table EVA outperforms most previous baselines and remains highly competitive with softmax attention, ilustrating its effectiveness.
5.2 MACHINE TRANSLATION AND LANGUAGE MODELING
We further evaluate EVA on the natural language domain. Specifically, consider three tasks:
Masked language modeling (MLM) on a pretraining-sete book corpus Books in the Pile
dataset suite (Presser 2020 Gao al. 2020 consisting of over 196,640 published books.
Machine translation (MT) on WMT1 En De benchmark (Bojar et al. (2014).
Autoregressive language modeling (Autoregressive LM) large-scale token-level LM bench- mark Wikitext - -103 (Merity et al. 2016).
Results. We report MIMM validation perplexity in Table where the sequence length is 2048 by default. EVA substantially improves previous methods based on random features (including Performer and LARA) and outperforms the other efficient attention mechanisms. Thanks the linear
7

---

Published as a conference paper at ICLR 2023
Lock Perfurer
Aitertion An Anriereer
EVR hysidven
EAA
Softres
setorey
2000
3000
4000
19:55 N000
8000
1032
120
200
1003
1300
100
3003
1998
1000
(a) Memory consumption.
(b) Running time.
(c) Training speed-up.
Figure 1: Left and middle: empirical memory consumption and running time comparison respectively of different attention mechanisms under various sequence lengths. Right: a snapshot of MLM validation loss curve versus actual elapsed time during training.
complexity of EVA, it can be scaled further to much longer sequences. With input sequences of length increased to 4096, EVA (denoted as "EVA-4096") attains lower validation perplexity than exact softmax attention, which demonstrates its capability of scaling to much longer sequences. Besides, machine translation results are compared in Table where in this task C = 8 by default and EVA-m denotes EVA with |F] =m. EVA outperforms previous random feature methods by large margin and achieves translation quality on par with full softmax attention even under the setting of small [R] and C. For Autoregressive LMM (Table 5). EVA achieves the same perplexity as softmax attention with much lower computational complexity. Comparing against various random feature methods reported by previous work Kasai et al. (202T, we observe significant performance gain brought from EVA even under a Transformer with half parameters. When further increasing the transformer model size as the setting in Kasai et al. (2021), EVA still scales as effectively as softmax attention with comparable perplexity while outperforming previous random feature methods by a larger margin. These results indicate the substantially enlarged capacity of EVA to approximate softmax attention.
5.3 ANALYSIS
Running Time & Memory Comparison. We con- duct simulation experiment to evaluate the empirical efficiency of various attention methods, which is mea- sured by the running time per iteration and memory footprint under different sequence lengths. The setup can be found in Appendix E..4 As illustrated in Fig- ures La and lb. EVA only incurs little computational overhead compared to Performer and local attention and achieves much better running time speed-up than Long-Short (Zhu et al. 202T). a strong baseline across various tasks albeit with much longer running time and larger memory consumption. In Figure we further visualize the speed-up -u of EVA relative conventional softmax attention by plotting the validation loss curve versus actual elapsed time during training transform- ers (equivalent to 32 GPU days). It can be seen that EVA can achieve a much lower loss after running for the same elapsed time; in contrast, conventional soft- max attention needs to run almost 3x longer to match the loss quantity. Overall, our method attains a good trade-off between quality and empirical efficiency.
Table 6: Classification accuracy Ima Imag 1 dataset.
Model
Mos (G)
Timequnsttrr
[8]
C
Top-1
Perfomer Local
8.1
87
0)
1
67.19
7.8
65
70.62 49
49
74.33
75,83
EVA
49
76.67
12.5
119
89
196
77.10
11.9
108
196
49
77.36
Sectmax
17.7
99
B.8.
n.a
77.16
Table 7: MLM validation perplexity on Books3. "_" indicates fail to converge.
Model
MenatCi
Timelundiee)
[B]
Purplesay
Performer-t0s Local-4055
1
256
0
4.34
IVA-4066
256
3.73
Soflmaas-0966
3.65
Ablation Study. In this section, we conduct an ablation study on image classification and MLM tasks to investigate the effects of main hyper-parameter analysis). In particular, we vary EE| and the partition size C' and evaluate their performance on both image classification and masked language modeling. As presented in Table and Table 7 increasing [E] amounts to obtaining exact estimates for more key-value pairs, which greatly improves empirical performance; besides, increasing would process control variates at a finer scale, also translating into better modeling quality, consistent with our theoretical analysis (4.1)
EVA (see Table for more comprehensive
8

---

Published as a conference paper at ICLR 2023
6 RELATED WORK
Control Variates. Control variates are a widely used variance reduction technique in reinforcement
learning (Greensmith et al. 2004; Grathwohl et al., [2018] Vlassis et al. [202T), stochastic optimization
(Wang et al., (013), variational inference (Paisley et al. 2012 Ranganath et al., 2014 Geffner & Domke 2018 Tucker et al. 2017 Grathwohl et al. [2018). Markov chain Monte Carlo (Baker et al.
2019) and many other topics. Our construction with control variates provides a new perspective on
designing faster yet more accurate attention approximations.
Efficient Attention Mechanisms. A lot of research work has put the focus on reducing the quadratic complexity of conventional softmax attention. A widely used approach is to define a sparse attention pattern so that each query is limited to only attending to a subset of tokens. The sparse pattern could
be either learnable (Kitaev et al.) 20200 Vyas et al. [2020] Tay et al.] 2020; Roy et al. [021] Madaan
iet al. 2022) or simply fixed (Liu et al.,
2018 Parmar et al. 2018 Child et al. 2019 Beltagy et al.
2020 Ainslie al., 22200] Zaheerr et al.
20200 Liu et al. 202 Xiong et al. 202 10 Wang et al. 2022
Chen et al. 2022 Hutchins et al. [2022]. Another paradigm is to adopt low-rank approximations,
including the Nyström method (Xiong al. 2021b), down-sampling with learnable projections
(Wang et al. 2020 Peng et al. (021a)), , or explicitly compressing sequences (Rae et al. 2020 Dai
ctt al., 2020 Ma et al. [021]] Jaegle et al. 202I). There are also studies improving both sparse and
low-rank methods for better attention matrix approximation (Nguyen et al. 2021 Zhu et al.
2021
Chen et al. 2021a Ren et al. 202T, Zhu & Soricut, 2021 Hua et 2022 Zeng et al. 2022).
Instead
of adopting approximate methods, a recent line of work (Rabe & Staats [2021] Dao et al.
2022)
proposes to compute the exact softmax attention in an online manner (Milakov & Gimelshein, 2018)
without materializing the full attention matrix. In this way, softmax attention can be computed in
linear memory complexity, and the runtime can also be greatly improved by further minimizing memory accesses (Dao et al. 2022).
Random-feature-basee methods are a popular alternative that uses random features (Rahimi & Recht, (008) to linearize exponential kernels in softmax attention (Katharopoulos et al., 2020 Choromanski et al., 2021 Peng et al. 2021b). Recent work attempts to
Random- Feature-beed Attention.
improve RFA approximation from several aspects, such designing more accurate random feature
maps (Choromanski et al. 2022 Likhosherstov et al. 2022 Chowdhury et al., 2022), incorporating
relative positional or other task-specific biases (Liutkus et al.] 2021 [uo al. [022 Chen 2021
Zheng et al. 2022a Qin et al. [022b) Wu et al.; 2022| Qin et al. 2002a). or leveraging connections
to fast weight programmers (Peng al.
| 202 closely related to ours includes Zheng et al. (2022b), which reinterprets RFA using self-normalized
|Schlag al.] 2021 Irie al. 2021). Prior work
importance sampling (Hesterberg] 1995) and theoretically extends the random feature approximation
from individual exponential kernels to the whole softmax attention. Our work further generalizes this result via control variates and characterizes the approximation gap caused by RFA. Scatterbrain (Chen et J 2021a) is also similar to our work in that it also refines RF approximation on critical local regions. However, it is developed based on a different motivation that attempts to approximate the attention matrix with combination of sparse and low-rank matrices. Interestingly, we find that Scatterbrain can be cast as a special case under our framework; see Appendix G for detailed discussion about connections between EVA and previous attention mechanisms.
7 CONCLUSION AND LIMITATIONS
In this work, we develop an efficient attention mechanism EVA via control variates. Our framework reveals localized perspective of RFA approximation, which not only bridges the gap between RFA and exact softmax attention but also attains a good trade-off between modeling quality and efficiency. We evaluate our method on both vision and language tasks and demonstrate substantial improvements over previous baselines. There are some limitations of our framework. For instance, the approximation in computing control variate estimation for each partitioned subset is crude and might limit the potential modeling capacity; in addition, we only explore the most straightforward partitioning strategy that evenly splits the sequence into multiple contiguous chunks; while in general, the partition could contain arbitrary subsequences or be adaptive to inputs via clustering methods, which can be guided by spec-fce inductive biases. It is interesting to investigate these limitations to unleash the expressiveness of EVA further, which we leave for future work.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
We would like thank the HKU NLP group, the Shark-NLP group, and the anonymous reviewers for their valuable suggestions that greatly helped improve this work. This work is partially supported by the joint research scheme of the National Natural Science Foundation of China (NSFC) and the Research Grants Council (RGC) under grant number N HKU714/21.
REFERENCES
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. Etc: Encoding long and structured inputs in transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 284, 2020.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint vrXiv:1607.06450 2016.
Alexei Baevski and Michael Auli. Adaptive input representations for neural language modeling. In International Conference Learning Representations, 2019. URL httpss://ppeeeeeiiie
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint 409.v:1409.0473, 2014.
Jack Baker, Paul Fearnhead, Emily Fox, and Christopher Nemeth. Control variates for stochastic gradient meme. Statistics and Computing, 29(3):599-615,
2019. 2019.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.051/ 050000000555,, 2020.
Ondřej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, et al. Findings of the 2014 workshop statistical machine translation. In Proceedings of the ninth workshop statistical machine translation, pp. 12-58, 2014.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Ben- jamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 1877-1901. Curran Associates, Inc., 2020. URL httpps::proceeddigggge neur ips ott.://cc....attt2.......a222222
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. In Thirty-Fifhh -Fitth Conference on Neural Information Processing Systems, 2021a.
Beidi Chen, Tri Dao, Kaizhao Liang, Jiaming Yang. Zhao Song, Atri Rudra, and Christopher Re. Pixelated butterfly: Simple and efficient sparse training for neural network models. In International Conference on Learning Representations, 2022. URL
idd-Ei-1Xa-y77
Chun-Fu Chen, Rameswar Panda, and Quanfu Fan. Regionvit: Regional-to -oooonal-to-locl attention for vision transformers. arXiv preprint arXiv:2106.022 2021b.
Peng Chen. PermuteFormer: Efficient relative position encoding for long sequences. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pp. (0606-10618) Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguistics. doi: 10.18653/v1/200 wew.leemmmmaaaaae main.828. URL https: http:////aaaaaaaaaaaaeaa //t.a/aaannnoo..oorrg.....
2021 emn 1p-ma iimmaan.... 828
10

---

Published as a conference paper at ICLR 2023
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10S000 2019.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations, 2021. URL https //openreview. n forum? id Ua6zukORRH
Krzysztof Marcin Choromanski, Han Lin, Haoxian Chen, Arijit Sehanobish, Yuanzhe Ma, Deepali Jain, Jake Varley, Andy Zeng, Michael Ryoo, Valerii LiKhosherstov, Dmitry Kalashnikov, Vikas Sindhwani, and Adrian Weller. Hybrid random features. In International Conference Learning Representations, 2022. URL
Sankalan Pal Chowdhury, Adamos Solomou, Kumar Avinava Dubey, and Mrinmaya Sachan. Learning the transformer kernel. Transactions on Machine Learning Research, 2022. URL https
/oopenreview. net nett/Forrmmid-LLLBBBEE
Zihang Dai, Guokun Lai, Yiming Yang, and Quoe Le. Funnel-tramee transforr...oo.m.... Filtering out sequential redundancy for efficient language processing, Advances in neural information processing systems, 33:4271-4285 4282,, 2020.
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. arXiv preprint arXiv:2205.141135, 2022.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pp. 245-255. leee, 2009.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volune I (Long and Short Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019, Association for Computational Linguistics. doi: 10. 10.18653/1//111-1444 1423. URL https
//aclanthology wow...ooooooooooooooooooooooo.. org 2org/199-220 1423
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 seale. arXiv preprint arXiv:2010.1 / 1929, 2020.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn. Xiaohua Zhai Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021. URL htttpps://ppeereiiiiiee net / ftt://orum222 id- VicbFaNTTy
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint 1XXi:2101100022 2020.
Tomas Geffner and Justin Domke. Using large ensembles of control variates for variational inference. Advances in Neural Information Processing Systems, 31, 2018.
Will Grathwohl, Dami Choi, Yuhuai Wu, Geoff Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In International Conference on Learning Representations, 2018. URL
forum?
id=SyzKd1bCR
Evan Greensmith, Peter L Bartlett, and Jonathan Baxter. Variance reduction techniques for gradient estimates in reinforcement learning. Journal of Machine Learning Research, 5(9), 2004.
11

---

Published as a conference paper at ICLR 2023
Tim Hesterberg. Weighted average importance sampling and defensive mixture distributions. Techno- metrics, 37(2):185- 194, 1995. ISSN 00401706. URL htt.lhhtp///tt http:://wwwwsss....oo.ooo. wwt.or..rg//ttllleee../.
1269620
Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefler, and Daniel Soudry. Augment your batch: Improving generalization through instance repetition. Proceedings of the IEEE/CVE Conference Computer Vision and Pattern Recognition, pp. 8129-8138, 2020.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoe V Le. Transformer quality in linear time. arXiv preprint arXiv:2202.10444 2022.
DeLesley Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block- Recurrent Transformers. arXiv preprint arXiv:2203.07852, 2022.
Kazuki Irie, Imanol Schlag, Róbert Csordás, and Jürgen Schmidhuber. Going beyond linear trans- formers with recurrent fast weight programmers. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL htttps://openrewee
net htte//orrrmmdidooiiieeeee
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Car- reira. Perceiver: General perception with iterative attention. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning. volume 139 of Proceedings of Machine Learning Research, pp. 4651-4664 1664. PMLR, 18 18-24 Jul 2021. URL m wtt.ll...eess/v..///// jaggle21eelee html
atttps://proceedingg
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A. Smith. Finetuming pretrained transformers into RNNs. In Proceedings of the 202/ Conference on Empirical Methods Natural Language Processing, pp. 10630-10643, Online and Punta Cana, Dominican Republic, November 2021. Association for Computational Linguisties. doi: 10. 12.18655/11/22111 http://22...mmmmmmm.m....m URL httppsss///caaattlllllgg/oo/ org/ 2021 em mapp-maiin inn 830
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pp. 5156-5165. PMLR, 2020.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020. URL/tttps://ppeeeeiiie
net/ Eorum? http:/oormmmdd.kggeeeee
Kunchang Li, Yali Wang, Junhao Zhang, Peng Gao, Guanglu Song, Yu Liu, Hongsheng Li, and Yu Qiao. Uniformer: Unifying convolution and self-attention for visual recognition. arXiv preprint arXiv:2201.094500 2022.
Valerii Likhosherstov, Krzysztof Choromanski, Avinava Dubey, Frederick Liu, Tamas Sarlos, and Adrian Weller. Chefs' random tables: Non-trigonometric random features. arXiv preprint arXiv:2205.15317. 2022.
Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. A structured self-attentive sentence embedding. arXiv preprint orXiv:1703.03111 2017.
Peter J. Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. In International Confer- ence on Learning Representations, 2018. URL https::///// http::::///peneeeeeeeeeeeee////e HygOvbWC-
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 arXiv: 1907. 177.11662.. 2019.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IFFF/CVF International CoaeeaaaCCoooeeee Conference on Computer Vision (ICCV), pp. 10012-10022, October 2021.
12

---

Published as a conference paper at ICLR 2023
Antoine Liutkus, Ondřej Cifka, Shih-Lun Wu, Umut Simsekli, Yi-Hsuan Yang. and Gael Richard. Relative positional encoding for transformers with linear complexity. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 7067-7079. PMLR, 18-24 Jul 2021. URL https:
ngs. m www...pres//..........mmm http://u/////uuuuuuuuusuua html
Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint (00.0:/0008..9838 2016.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang. Soft: Softmax-ffee transformer with linear complexity. In Thirty-Fifth Conference on Newral Information Processing Systems, 2021.
Shengjie Luo, Shanda Li, Tianle Cai, Di He, Dinglan Peng, Shuxin Zheng, Guolin Kc, Liwei Wang, and Tie- Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021.
httttp:://peenrvleew..ee.. wwt.:fettooooaaee
Xd=X7XNPor920 or 93G
Xuezhe Ma, Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettle- moyer. Luna: Linear unified nested attention. arXiv preprint arXiv:2106.00 01540, 2021.
Lovish Madaan, Srinadh Bhojanapalli, Himanshu Jain, and Prateek Jain. Treeformer: Dense gradient trees for efficient attention computation. arXiv preprint arXiv:2208.0900 09015, 2022.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint 160.004609007843 2016.
Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv: 1805.02666.. 2018
Tan Nguyen, Vai Suliafu, Stanley Osher, Long Chen, and Bao Wang. Fmmformer: Efficient and flexible transformer via decomposed near-field and far-field attention. Advances in Neural Information Processing Systems, 34, 2021.
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. In Proceedings of the Third Conference on Machine Translation: Research Papers, pp. 1-9, Brussels, Belgium, October 2018. Association for Computational Linguistics. doi: 1111663/11/W11-6333 URL
6301
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: fast, extensible toolkit for sequence modeling. Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties (Demonstrations)) pp. 48-53, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: (118866/11//119/00
CtL.https://
anthology 2otg/N1994000
Zijing Ou, Tingyang Xu, Qinliang Su, Yingzhen Li, Peilin Zhao, and Yatao Bian. Learning set functions under the optimal subset oracle via equivariant variational inference. arXiv preprint arXiv:2203.01693, 2022.
Art B. Owen. Monte Carlo theory, methods and examples. 2013.
John Paisley, David M. Blei, and Michael I. Jordan, Variational bayesian inference with stochastic search. In Proceedings of the 29th International Coference on International Conference Machine Learning, ICML'12, pp. 1363-1370, Madison, WI, USA, 2012. Omnipress. ISBN 978195011 14503
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference on Machine Learning, pp. 4055-4064. PMLR, 2018.
13

---

Published as a conference paper at ICLR 2023
Hao Peng, Jungo Kasai, Nikolaos Pappas, Dani Yogatama, Zhaofeng Wu, Lingpeng Kong, Roy Schwartz, and Noah A Smith. Abc: Attention with bounded -ennddd-memory control. arXiv preprint arXiv:2110.02488, 2021a.
Hao Peng,
Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah Smith, and Lingpeng Kong. Random feature attention. In International Conference on Learning Representations, 2021b.
Shawn Presser. Books3. 2020.
L320222443327784888
Zhen Qin, XiaoDong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. arXiv preprint arXiv:22 10. 10340, 2022a.
Zhen Qin, Weixuan Sun. Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In International Confer- ence on Learning Representations, 2022b. URL httpp:://ppppeeeeeeeeeeee B18CQrx2Up4]
net/ forum?id=
Markus Rabe and Charles Staats. Self-attention does not need O(n²) memory. arXiv preprint arXiv:2112.05668, 2021.
Iliaa Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, and Piotr Dollár. Designing network design spaces. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 10428-10438) 2020.
Jack Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Com- pressive transformers for long-range sequence modelling. In International Conference Learning Representations, 2020. URL
net/ FowoomiidSSyylllo 1iikSYYD
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. Journal of Machine Learning Research, 21 2020a. URL
(etaperssv21/221 v21/20 074.htm.
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 unified text-to-text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020b. URL
http:://mmr wrt/papers wv11202110000222 002212200044111
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In J. Platt, D. Koller, Y. Singer, and S. Roweis (eds.), Advances in Neural Information Processing Systems, volume 20. Curran Associates, Inc., 2008. URL
meer neur pss cc/
paper/2007/F1 1tt00///11110111000000
Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Artificial intelligence and statistics, pp. 814-822. PMLR, 2014.
Hongyu Ren, Hanjun Dai, Zihang Dai, Mengjiao Yang, Jure Leskovee, Dale Schuurmans, and Bo Dai. Combiner: Full attention transformer with sparse computation cost. Advances in Neural Information Processing Systems, 34, 2021.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics. 9:53 68, 2021. doi: 16.1162/ a 00353. URL peatps://aclantholloo tacl-1.4
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 9355-9366. PMLR, 18-24 Jul 2021. URL http::::/ppooceednngge mllr. wwwpp/ewsssseeeeeeeeeppppp sch 1 ag21 a html
14

---

Published as a conference paper at ICLR 2023
Yi Tay, Dara Bahri, Liu Yang. Donald Metzler, and Da-Cheng Juan. Sparse Sinkhorn attention. In Hal Daumé III and Aarti Singh (eds.), Proceedings of 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9438 9009994477777 PMLR. 13-18 Jul 2020. htt.ttsss//prrocedingg..rr http./mr....eee///////////0. html
Yi Tay, Mostafa Dehghani, 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, 2021. URL https
net. wtr.eoorumiiiirmme
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learming, volume 139 of Proceedings of Machine Learning Research, pp. 10347-10357. PMLR, 18-24 Jul 2021. TRRLattaa ttttss://proceedinng
mlz
html
Zhengzhong Tu, Hossein Talebi, Han Zhang, Feng Yang, Peyman Milanfar, Alan Bovik, and Yinxiao Li. Maxvit: Multi-axis vision transformer. arXiv preprini urXiv:2204.0169)) 2022.
George Tucker, Andriy Mnih, Chris J Maddison, John Lawson, and Jascha Sohl-Dickstein. Rebar: Low-variance, unbiased gradient estimates for discrete latent variable models. In Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. URL htttps:/pprcceeeliee neurips httppppaee/0111//EEE... ebd6d2 f5d.f www.orraaaaaaaaaa a8 a0lffff3e2000
pdf
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 information processing systems, pp. 5998-6008, 2017.
Nikos Vlassis, Ashok Chandrashekar, Fernando Amat, and Nathan Kallus. Control variates for slate off-policy evaluation. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL hnttps://openreeiiew 3__PPM_PPP_
Apoorv Vyas, Angelos Katharopoulos, and François Fleuret. Fast transformers with clustered attention. Advances in Neural Information Processing Systems, 33, 2020.
Chong Wang, Xi Chen, Alexander J Smola, and Eric Xing. Variance reduction for stochastic gradient optimization. In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q Weinberger (eds.), Advances in Neural Information Processing Systems, volume 26. Curran As- sociates, Inc., 2013. URL https: http:ss:///pccpepppeeeeeooeeeees neurips http:lppp.pppppppepeee cc/papprr/013//..e.... Paper. pdf
Sinong Wang. Belinda Z Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint wrXiv:2006.04763 2020.
Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pyramid vision transformer: A versatile backbone for dense prediction without convolutions. arXiv preprint arXiv:2102.12122, 2021a.
Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao. Pvtv2: Improved baselines with pyramid vision transformer. arXiv preprint arXiv:2106. 13797, 2021b.
Yuxin Wang, Chu-Tak Lee, Qipeng Guo, Zhangyue Yin, Yunhua Zhou, Xuanjing Huang, and Xipeng Qiu. What dense graph do you need for sclf-attention? In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 22752-252222222 22768. PMLR, 17-23 Jul 2022. URL btttps:/ppoceeeinngg mlr preas/vl 62 wa .html
/
15

---

Published as a conference paper at ICLR 2023
Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang. Cvt: Introducing convolutions to vision transformers. arXiv preprint arXiv:2103. 15808, 2021.
Haixu Wu, Jialong Wu, Jichui Xu, Jianmin Wang. and Mingsheng Long. Flowformer: Linearizing transformers with conservation flows. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 24226- 24242. PMLR, 17-23 Jul 2022. URLhttpa httpss////pocceeingg..ccmo html
/ttpo/////////ppreeeeeeooo.ooo.. mlr www.s.presssvvvvv / ww2 2m
Tete Xiao, Mannat Singh, Eric Mintun, Trevor Darrell, Piotr Dollár, and Ross Girshick. Early convolutions help transformers see better. arXiv preprint arXiv:2106. 14881, 2021.
Wenhan Xiong, Barlas Oguz, Anchit Gupta, Xilun Chen, Diana Liskovich, Omer Levy, Wen-tau Yih, and Yashar Mehdad. Simple local attentions remain competitive for long-contexs tasks. arXiv preprint arXiv:21 12.07210, 2021a.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströnformer: nyström-based algorithm for approximating self-attention. In Proceedings of the AAAI Conference Artificial Intelligence, pp. 14138-14144 2021b.
Jianwei Yang, Chunyuan Li, Pengchuan Zhang, Xiyang Dai, Bin Xiao, Lu Yuan, and Jianfeng Gao. Focal attention for Long-range interactions in vision transformers. Advances in Neural Information Processing Systems, 34, 2021.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed, Big bird: Trans- formers for longer sequences. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.). Advances in Neural Information Processing Systems, volume 33, pp. 17283-172977 Cur- ran Associates, Inc., 2020. URL URL//ntps://pprrcccc..g.... neuriss........p. httpppps.ccpppppe///////// Eile:ceb222111222245555555555513
Zhanpeng Zeng, Sourav Pal, Jeffery Kline, Glenn M Fung, and Vikas Singh. Multi resolution analysis (MRA) for approximate self-attention. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 25955 25972. PMLR, 17-23 Jul 2022. URL zeng22a. html
wwwinngsss...... www.l..pees/vveee /
Pengchuan Zhang, Xiyang Dai, Jianwei Yang, Bin Xiao, Lu Yuan, Lei Zhang, and Jianfeng Gao. Multi-scale vision longformer: A new vision transformer for high-resolution image encoding. arXiv preprint arXiv:2 10.01.3..33...3 15358, 2021.
Lin Zheng, Huijie Pan, and Lingpeng Kong. Ripple attention for visual perception with sub-quadratic complexity. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 26993-27010. PMLR, 17-23 Jul 2022a. URL httph: httttp:://proceedinnge
www..ppress//eee https//ww62//heng222 zhwng22ann222aaa .hmml
Lin Zheng, Chong Wang, and Lingpeng Kong. Linear complexity randomized self-attention mecha- nism. arXiv preprint arXiv:2204. 04667,0666666 2022b.
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. Advances in Neural Information Processing Systems, 34, 2021.
Zhenhai Zhu and Radu Soricut. H-transformer-ID ID: Fast one-dimensional hierarchical attention for sequences. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 3801-3815, Online, August 2021. Association for Computational Lin- guistics. doi: 10. 1/2.2////1/221... htt.://221........ long.lon..294
URL https: http://s///allaattoooo....oo..o
acl 1 llogong..2 294
16

---

Published as a conference paper at ICLR 2023
Appendix
Brief Review of Vanilla Random Feature Atfention
18
Proofs Derivations
18
B
An Extended Review of Control Variates
18
B.2
Derivation of SNIS Control Variate Estimation
19
B.3
Derivation of the Expectation of Per-term Control Variates
19
B.4
Proof of Proposition]]
19
B.5
A Formal Analysis of the Advantage of Partitioning
20
B.6
Proof of Proposition
21
B.7
Derivation of Equations 14 and
23
More Implementation Details for EVA
23
D
A Causal Variant of EVA
25
Experimental Details
26
Efficient Attention Baselines
26
Image Classification
27
Machine Translation and Language Modeling
29
Experimental Settings of Efficiency Comparison
Experiments Long Range Arena
32
Connections to Other Attention Mechanisms
33
RFA, Softmax Attention, and EVA
33
G.2
Connections to LARA
33
Connections Clustered Attention
33
G.4
4
Connections to Combiner
34
Connections Scatterbrain
34
17

---

Published as a conference paper at ICLR 2023
A A BRIEF REVIEW OF VANILLA RANDOM FEATURE ATTENTION
Vanilla random feature attention methods, such as Performer (Choromanski al. 2021 Peng et 2021b), seek to approximate the softmax attention mechanism through random features RRaimmi & Recht 2008) Φ x, w) 1/v 1/V/[[((x,)) [(x,)s)ss Here, ,,, N/ 1), and {(x. w) is the randomized mapping such that
exp (q, km) ((...) (5(g)) " "(kmm W,)]
(17)
Therefore, we can draw multiple Monte Carlo samples to estimate the exponential kernel,
s exp (q, km) Σ ((gn:ws) E(km, ((kksss)T@((sssiia = ())n: @lkmm W),
and then approximate the attention mechanism as
M exp km) M
Em=1 Φ(qn. w) TK w)vm
M
Vm
(18)
m1=1 exp (q, (q..k, km))
Em'=1 Φ(qm. TT(km',ww)
It is recently generalized as self-normalized importance sampling estimator to approximate softmax attention (Zheng et al., 20226), as described in 22.2 We refer the generalized random feature based approximations RFA.
B PROOFS & DERIVATIONS
B.1 AN EXTENDED REVIEW OF CONTROL VARIATES
The control variate method takes the following form,
g(w) g(w) - -Bh(ww + BE [h(w)],
(19)
Given the particular forms of g(-) and h(-). 3 can be optimized minimize the estimation variance. For notational convenience, we denote the covariance between a scalar and random vector as Cov [h(w), [ww] = E [h((w) E [hw))) (g(w) E [g(w]])), and the variance of a random vector as Var [g(w)] = Cov [g(w), g ((w)]. In particular, we have
Var [(w)] Var (g(w) Bh(ww)
Var [g(w)] Cov (3h(w), [(w)] + Var [Sh(w)]
Var [g(w)] 2 Cov [h(w), [(w) B Var [h(w]] (3
We hope an optimal ,3 would minimize Tr (Var [)())]) that is, the sum of estimating variance for each dimension. By differentiating, obtain
Cov [h(w), ا(س)و Var [h(w)]
3 = arg min Tr (Var [(w)))
(20)
3
Since both the covariance and the variance may be intractable to compute, the optimal (3" is generally not available in closed form. Nevertheless, with the optimal coefficient, the variance such control variate estimate would never be larger than the plain estimator g(.)-
18

---

Published as a conference paper at ICLR 2023
B.2 DERIVATION OF SNIS AS CONTROL VARIATE ESTIMATION
For notational convenience, we denote the importance weight as W (wis) = pm (wWW)))((ws) Then we have
Pnlwel (Wa) 9(wa))
_a=1 (ws))
g(w)
===
S Pa (a
Σ ((we)
((wa)
Σ W (ww))((W)))
((w))
3 [Ww))((ss)
+
1
2 ((ws)
5
5
#=1
#=1
W (ww))
((w)) 1
W W(w,)Swss) 1
Σ
Wwws)((ws)
> ((w))
> W(ws)
S
8
E W(ws)
Σ W(wz)((ws)
1
Σ
W(ws)((ss)
Σ3 (wws)
5
(Wwo)f(Wo)
1
8
1-=1
1
Σ W(ws)
+
W(ws)f(wa)
ΣS 2==1 ((ws)
S
5
s=1
8=1
1 (ww))
ESE W (wws)
1
W W(wa)) -
5 ==1
Σ ((w))
5
= ((w) - ((w) (h(w) - E [h(ww)))
Note that the expectation of importance weights equals I, that is,
1 S
g
s 1 p(ws)
s )(ww)
p(we)
E [h((w)]=B =E EE
W (wa) = B (g =oo(e..
E 5
1.
5 g(ws)
g(wa)
Same as SNIS, this estimator is still biased due to the dependence of 3(w) on w. However, it would asymptotically become unbiased since ((w) is consistent and converges to constant B w.r.t. given a large number of samples.
B(w)= g(w)
E [w()]
(.) [f(W]
= B.
(21)
h(w)
E [[ww]]
constant
B.3 DERIVATION OF THE EXPECTATION PER-TERM CONTROL VARIATES
According to the definition of randomized mappings, we have
M(ws; Zg(wa) [(qm We (kkm ww) wa)
[hm (w)) Evs, (
1 1
E(gn) (() Wa Wa)N(ws 0, I)dws
7
s=1
cxp(q. km)
(22)
Z
B.4 PROOF OF PROPOSITION
Proof. We start with the formulation of g(-) and h(-).
>5 Niw.:0.D {(q,, W3 (km, Wm)
Im (m((w)
Zy(w.)
Vn-
hm (w)
<( W. (km, We)
1=1
2q(wa)
19

---

Published as a conference paper at ICLR 2023
As result, have 9m (w) hm vn and E (w)] = E [hm [w)] Vm- We now investigate the
optimal Bm according to Equation
Fm = arg min Tr (Var [7m (w)))
Cov (w)) Im (w))
Var mm (w)]
E [(h(w) E [h(w)]) (h(w> - E [h(w)]] Vm
E ((h(w) [h(w))) (h(w) [M()))
9m (w)
hm (m((w)
In terms of the variance, again use Im (m) hm (w) to obtain
Im (w) Im (w) 8m (Je((w) E [mm (w)))
Im (w) Vm (w)) + Vm [hm (w)]
vm E [Am (w)]
cxp(q,, km)
Vm
(23)
Z
Since this holds true for every term m 1,.... M, our estimate becomes exactly softmax attention,
M
M
exp(q km)
M
explq km
g(w) = Σ gm((w) (w) =
m=1 Z
M =m==1 exp(q, km)
Vm.
m==1
m=
Since all randomness is eliminated, the estimate is exact with zero bias and variance. That is, RFA q., K, V) g(w) SoftmaxAttn(q, q,, K, V).
B.5 A FORMAL ANALYSIS OF THE ADVANTAGE OF PARTITIONING
In this section, we demonstrate the advantage of partitioning by showing that there always exists some set [3, )) that achieves lower weighted MSE than any globally shared coefficient, as discussed in $4.1
Lemma 3. Suppose 3r, the optimal coefficient for each m E [M] as defined in Proposition and
Pi, P2...., Po arbitrary partition of [M], where each subset Pa is associated with distinct
B. We consider the following weighted mean squared error,
J(B
= Σ
am 1| 3
(24)
C= mEP.
where a'm > 0 for each m [M] and Ee EmEp. am Then for any choice of {am JM
and any globally shared coefficient B3 there exists some {B) so that
J(B = = B...,,
Be = 3) 2 J(B = - B3))
Bc
Emer, amp Emep, for each c = 1,..., C. Then we have
Proof.
Let 3r, =
Σ am (Br - 3r 3
Σ am
Σ
am
meR.
mER.
EmEr am
Σ am
Σ
amBa m
Emep,
mEP.
mePe
am am
3 m
Om am 3
0.
(25)
mER
mEP.
20

---

Published as a conference paper at ICLR 2023
According to Equations 25 and 24 for any B we have the following inequality,
J( B..... :c....c= =g)
am ||3 BTnll c=1 mEP
Om || B² + 3 Imll
MEP.
am 350- 3:11 12(B-33) (B% 3rn) + +11.2 -B 1?²
mE
Σ am 1B- 3-12 +2 Σ
Σ
am (B - (3) T (B Bmm + Σ Σ
am 113 - -3
c=1 mEPe
c=1 mEPe
c=1 mEP
C
Σ am || - 3 +
Om 11 1'
-= mEP.
mEP.
am 11 - Brh 11
mEP
= = .i...... 33)
As result, for any choice of {am 1M and any globally shared coefficient B, there always exists some [B that achieves lower (or equal) weighted MSE, and solution can be simply B_ Emep, am
EmEr.
B.6 PROOF OF PROPOSITION
2
Proof. We first consider the case of partitioned indices, where each subset PP is associated with some specific B.. To see the global minimum of J, differentiate on both sides and obtain
0
exp km)k))))))) Σ Em'cv EU exp (q, km))
||B
03.
03.
c=1 mEP.
exp km)
2 (B. 3rn)
m.P. Emvu exp (q, km))
By setting the partial derivative to zero, we obtain
B = EmEr. exp (q km) B
Emep, exp (q, km) Vm
ΣmEp, exp (q km)
ΣmEp, exp (q (q..km))
Emer, E (0m(w))
E
[EmEp [m(w)]
Emep. E (w)]
E EmEp hmm [w)[w)]
As consequence, with 3, 8", the partition scheme must achieve lower weighted mean squared
emror than any globally shared B, that is, J(B , B (8) < J = 3...., Bc (3).
21

---

Published as a conference paper at ICLR 2023
In fact, with B. G the partition scheme usually enjoys much lower enror than adopting globally shared coefficient. To see the error reduction of using the partitioned strategy, we first have
J = B......o
Ac = =3)
exp (qm km)
B | Fmm l|
Σ
En exp (q,, kmr)
6=1 MEP:
exp km)
3 +. -3
Em' exp (q Kml)
6= mEPe
exp (9 km)
3- Br + B. -Bin
mEP. Σ ey exp (q km/)
exp (q km)
|13- + (3-1.) 3.)
(B. (3)- + -13 -/mm1
a MEP. E E exp (q kmr)
Since
exp (qn km)
B (B-B.) T (Be Bm)
Em exp (q,, km²)
mEP.
C
T
exp (q km)
B
3.)
(Bo - -Bm)
mEP: Em eu exp (qm km))
exp (4)
Emep, exp (q km) Vm
Σ
B (8-))
3
mEP: E m' EU exp (qn kmm)
EmEp, exp (9 km)
ΣmEp, exp (q km) (vm - -Bm))
Emeey exp (q km/)
=0,
plugging this result back we obtain
J(B = = B...., B.. = (3)
C
cxp (q km)
(113
33 sa| + (3-
(3-B.) B,)" (3.- - -Bmm) + 11. - -3/ll
Emeev cxp (q )
1= MEP.
exp (q km)
(133 3.1 +11B 3
MEP. Em'ev exp (q kmm)
C Emep, exp km)
cxp (q km
3-
+ c=1
Emeey exp (q, km)) 113 33,
Emeey exp (q km))
meP.
exp 9 km a=l mEP. Em'er exp (q, km))
||3 3 1
The last inequality holds since the first term is always non-negative. Note that the first term computes the squared error between and each B_, weighted by the sum of attention scores over the corre- sponding subset. As result, it is usually positive and the erтor reduction is significant if each B_ deviates from 3 lot. However, although the optimal coefficient in the partitioning always leads to lower error to the optimal individual coefficient, note that does not necessarily yield lower estimation variance.
22

---

Published as a conference paper at ICLR 2023
Table Classification results on ImageNet k dataset under different hyper-parameter configura- tions of EVA. By default, we set E| = 49 and C 49 across all variants below.
Component
Specification
Top-1 1 Acc.
partition over [M)
E
76.53
Partition Scheme of (P...... Pc)
partition over [
76.39
LN(Linear(-)) 0 Identity(-)
76.67
Parameterization of (r))
75.95
Number of Samples
74.10
Number of Groups (C' 1)
Number of Samples
76.39
Number of Samples
76.67
Number of Groups (C' 49)
Number of Samples
76.75
#e q. Ke
76.67
Proposal Parameterization (c(w) = (ww pe 1)
1 9
76.77
76.24
H. Trainable parameters
76.39
Softmax
77.16
B.7 DERIVATION OF EQUATIONS 14 AND 15
According to the definition of gin. (-) and hm (-) in $3.1 for the vanilla RFA (Equation 4 we have
Σ ==1 g(w.)
(w,(
f(we)
g(w) h(w)
DM Im (w) Em=1
RFA( K, V)
Σ (aws)
(w)
g(w,)
Besides, since Vm 9m (ollw//m and Bc(w) (Emep, Im Am(w)) (Emep, hm (w)), we can re-write EVA as
EVA(q,: K, V) = [(w)
Σ Im (w) + Σ
Im (w)
mER
mg
explq
T km
Emep, exp(a km) km))) Be(w)
7
m
7
mE F
c=1
exp(q, T. km
Im
Emep,
exp(qn km)
Emer, Im bm(w) Emep, Um(ww)
hmm
7
MEE
C MORE IMPLEMENTATION DETAILS FOR EVA
In this section, we provide more details of EVA. We also conduct comprehensive ablation study to test the effect of different components in our implementation and report the results in Table The pseudo-code for EVA is listed in Algorithm
Approximating Emep. exp(q) _ km) and Parameterizing k.. In our implementation, approx- imate the sum of exponentials as EmEp. exp(q,, km) exp(q, ko). Here we provide an informal justification for this approximation.
Our main motivation for such approximation is based on the simple intuition that the sum of exponentials grows as fast as the maximum exponential value, as reflected by the following inequality,
max exp(q,, km) S mEPe
Σ exp(q, km) S (Pal max exp(q.) km).
mEP.
mEP.
This means approximate the sum of exponentials by first computing the group representative k, arg max (km./mpp.) exp(q, km), evaluating the corresponding exponential exp(q,) k.)
23

---

Published as a conference paper at ICLR 2023
and then multiplying it by some scalar. Since computing the argmax operation still needs compare each exponential dot-product, it will still incur quadratic computational costs. To circumvent this, we adopt heuristic strategy that computes learnable group representation, which attempts to compensate for the approximation error while only evaluating one exponential dot product.
Through preliminary experiments, try various choices to compute the representative vector of each subset, such as max and average pooling; however, we found these strategies produce almost equally good performance. As result, we adopt the average pooling by default due its simplicity. To be specific, we implement it as
k =
Σ mEP.
km
(26)
iPcl
where (-) is trainable linear projection with the same hidden dimension size as inputs, followed by a layer normalization operation (Ba et al., (0T6) to stabilize training. We leave further improving the approximation, such deriving tighter error bounds using more expressive pooling methods (Zaheer et al. 2017 Ou et al. 2022) as future work.
Parameterizing (w). As discussed in $4.1 we have
EmEp. 9m(WW) Be(w) = Emep, (w)
Σg
s
Nww,:
EmEr. E(gn, Wa (k Vm Nlw,;0,1) 2y(wa)
Zy(ws))
Emep. E(qn, W. ( (5(km Ws)
e=1
Compared to the SNIS formulation of vanilla RFA Equation
we can express it as
E Pa(ws) a(w,)
f(ws)
Σm M 9m (w)
RFA(q,, K, V)
Σ s= a(w>)
EM (w)
We can think of each coefficient (w) computing the output of localized RFA for each group Pe. From this perspective, we can recast each coefficient (w( as SNIS estimator well, which tries to estimate
exp (q..kmm)
E Bwwpp(w) [J[[w]]
Σ
Emer, exp (qg km))
Vm
(27)
mEP.
where
fe(w) = Emep. Elqm, (Kkm w)
Emee. E(qm. (Kmm', W)
Pelw) = Nww; 0, I) Emep, E(gm: w) E(km, w)
Emve,, exp ppp km))
exp (m.
N(w; qn + km. I).
mEP. Em'Ep, exp (q. km))
This interpretation indicates that a good proposal distribution ge (w) should be specific to each subset P.. To get close to the true distribution (w) while keeping efficient computation, /Zeng et al. (2022b) suggests parameterizing the proposal distribution as
(w) = N(w; I) Nlw; (w:qq." + k,,
(28)
where q. is calculated similarly to Equation 26 We refer readers to Zheng et al. (2022b) for more discussions about the parameterization choice of proposal distributions. We conduct further ablation studies to test the effect of proposal parameterizations in our proposed model, as shown in Table X In particular, we found model is robust to different parameterization approaches.
24

---

Published as a conference paper at ICLR 2023
The essence in making the algorithm memorr--ffficiet is to use only one sample in caleulating (c(w). In this case, we have
B.(w) = Emep Im (m(w)
EmEp. hm
N(w :0.I)
Emep, W") w²) V
Zqelwe)
(golwr) EmEp, ((qn)w*) (k wi)
;0,";
Zyelwr) 5(gn) Emer, E(km) V
:( Zynlwe)
5(gns w*) Emep, S 5kmmaa"
EmEp, E(km; N
we ga(w).
EmEp, w0)
Since this degenerated formulation eliminates the dependence on individual queries 4n we could pre- compute (w) for each P.. and then re-uses them for each query, which takes up (Cd) memory. multiple samples are used instead, the influence of queries needs to be explicitly taken into account and thus we need to compute a distinct 3_(ww) for each query, leading to (NCcd) memory usage, which incurs significant compute overhead.
On the other hand, if we set C = 1, that is, using a shared (w(w) over all m & E, our approach does not suffer from this issue, since the memory usage is at most ((Nd). To investigate the effect of using larger C' or increasing the number of samples, we conduct an ablative analysis as in Table S and find that 1) when C' = the performance degrades lot when using one sample, which can be largely improved by adopting more samples; while when C' > 1, our partitioning strategy dominates and increasing the number of samples only improves performance marginally. This also validates the effectiveness of adopting a finer-grained treatment over control variates.
Partitioning Strategy. EVA significantly improves random feature approximation by trying to locally estimate each subset of tokens, which is much easier task than approximating the whole sequence as in previous REA methods. To achieve this, EVA partitions the whole token sequence into multiple subsets according to the current query position n, which is denoted by (E",Pi P . Po) ) [ For elements in subset E", we optimize the control variate coefficient to give exact estimate for each single token E E". In addition,
impose T5-style relative positional encoding (Raffel et al. 2020a) over elements in E", While for some other subset Pa. we employ the shared coefficient approximate all tokens belonging to Pa.. We assume all F. E. are of the same cardinality K, and |Pan is the same for any e = 1,..., and n 1,...,N:
The partition strategy {L",, Pf. PE...... Pa Pin is decided based on simple criterion:
for E", it contains K local neighbors with respect to each query n. To further simplify implemen- tation and reduce memory usage, chunk the whole sequence into contiguous blocks of size K, and all adjacent queries belonging to the same block will share this block as the subset E"; for PT. Pr.
Po. we follow a similar treatment by splitting the complement [M] \ E" into C contiguous chunks of the same size. For ease of implementation, we simply partition the whole index set [M] into multiple groups instead of [M] \ E",, which circumvents the overload for explicitly performing set difference operations in practical implementation. Although this leads to extra approximation error, this amounts to putting more attention weights on tokens belonging to the subset É and we found this approximation does not lead to performance degradation (Table 8)
D A CAUSAL VARIANT OF EVA
In this section, deseribe the causal variant of EVA, where each query can only attend historical tokens. Thanks to the partitioning scheme, all future information with respect to the current query token can be masked conveniently. Following the formulation of EVA, we partition the whole sequence into C + subsets (E",Pri .?. with respect to each query qn- To fulfill the
'Hre add the superseript reflect the dependence on query position n.
25

---

Published as a conference paper at ICLR 2023
Algorithm Pseudo-code for EVA
Input: the randomized mapping ((;,), queries Q == {q 1/ keys K {km 1M m==1 values M V = {vm m=11 and partitions of the sequence ,P. PP. p2...... PC)
Output: attention output Y := Jynin
for c = ........ do Compute k, according to Equation 26
Compute (e(w) according to Equation 28
Sample we ~ (w);;
D During inference, simply set we = (.((w) [w]
Compute Be(w) EmEpa Emery {(kmo)e)
5(kmiwe)
Min
end for for
1,2..... N, do Compute S EmERn exp (qn km) Vmi Compute attention scores the selected subset E
Compute R = Ec exp an (ank.) B,(w):
b Compute approx. expected control variates
Compute Z = Em exp km) + Ee exp (q. ko)
Compute y, = (S R) /2;
end for Return
yyl.
causal requirement, we design two different types of masking matrices to deal with both E" and PaaC respectively.
For E", we adopt a single lower-triangular matrix with shape K x K (recall that each set En is of size K) to mask future tokens locally, similar to the case of standard decoder softmax attention. Future tokens that do not belong to E" are handled by masking functions for {Pa: Pe as described below.
For [P:)e we make use of the fact n E E", Since any P and E are disjoint, we only need to mask all subsets P that appear after F This amounts to first allocating lower-triangular matrix with shape C' x C, and then conducting future masking at subset level.
The pseudo-code for the causal variant of EVA is listed in Algorithm) 2
EXPERIMENTAL DETAILS
All of our experiments are conducted with at most 16 NVIDIA V100 GPUs.
E. EFFICIENT ATTENTION BASELINES
We compare our proposed attention mechanism EVA against various baselines:
Performer (Choromanski et al. 2021). which uses the plain random features to approximate softmax
attention;
LARA (Zheng et al. 2022b), an advanced RF approximation that makes use of multiple adaptive
proposals construct the SNIS estimator;
Linformer (Wang et al.) 2020), low-rank approximation that uses a learnable matrix to project the
key-value sequence into shorter one;
Nyströmformer (Xiong et al.) (021b), low-rank approximation that adopts the Nyström method to approximate softmax attention map with a sub-sampled matrix;
Local attention (Child et al. 2019), simple sparse approximation that splits the whole sequence into multiple blocks and only allows the query to attend to tokens within the same block; Reformer (Kitaev al. 22220), sparse approximation where hash functions used adaptively distribute sequence tokens into multiple buckets, and each token can only attend to tokens within the same bucket;
26

---

Published as a conference paper at ICLR 2023
Algorithm 2 Pseudo-code for Causal EVA
Input: the randomized mapping +, queries Q == (q LM M V = {vm m=1- and partitions of the sequence { P2:
n= keys K {km 1 values
Pe N.
n=
Output: attention output Y := (ynjn=i
for c = 1....... do Compute k, according to Equation 26
Compute (e(w) according to Equation 28
Sample we ~ (w);;
D During inference, simply set we = (.((w) [w]
Compute Be(w) EmEpa Emery {(kmose)
5(kmiwe)
Vmn
end for Let K < |E Initialize M E {0, [0,1] such that M = ligh
we assume all E are the same in size
> Intra-B E masking matrix
Initialize M E {0, 1] cxC such that M Legsi
b Inter-P masking matrix
for 1,.,.... do
Find index such that P.' the most recent chunk on the left of E; Let 6* < min, E }"); The least position within E"; used for shifting token indices.
The same masking matrix M can be reused across via shifting token positions by b". Compute S EmER M m exp (qm (m Vmi
Compute = = M exp q ko) (.(w):
Compute Em M n- exp (q, km) + Ec M., exp q, k.):
Compute yn = (S R)) 17:
end for Return y
Scatterbrain (Chen et al. a0) , an approach that combines Performer and sparse attention. The details can be found Appendix G Here we implement the sparse module as simple local attention to ensure fair comparison;
Combiner (Ren et al. 2021), a probabilistic approach that constructs a structured factorization over the softmax probability distribution via sparse mechanism. Combiner allows both direct and indirect calculations of conditional probabilities, where the direct probability is implemented as the sparse mechanism while the indirect probability is implemented through local abstraction over group of tokens. Similarly, we implement the sparse mechanism as a simple local attention, which corresponds to the Combiner-Fixed Fixed variant (Ren al., (021):
Transformer-LS. or Long-Shor (Zhu et al., 2021). which is proposed to model long-term and short-term dependencies via low-rank structures and local attention respectively. The low-rank structure is defined as an input-dependent weight matrix that compresses the sequence into a shorter one; while the local attention defined similarly as above.
Note that for all mechanisms that involve local attention, we split the sequence into non-overlapping blocks (or 2D windows in terms of images) and each query can only attend tokens within the
same block. We also use the relative positional embedding (Raffel et al.) 2020b Liu et al. 2021 within the local attention computation. Unlike Transformer-LS (Zhu et al.) 2021 that allows each query to attend to multiple blocks, we do not use this extension as we find greatly increases memory consumption, although does improve the model performance.
F.2
IMAGE CLASSIFICATION
Through the experiments on image classification, we consider four different vision transformer (ViT) architectures:
27

---

Published as a conference paper at ICLR 2023
Table 9; Our hypeer-paaameter configuration for different attention mechanisms on DeiT-Tiny-P7
Attention
Hyper-parameler configuration on image clussification
Local attention
Window
umber random feature samples
96
Scateerbaii Kitaev et al. 2020
Local attention window size
49
Nyströmformer (Xiong et al. lb
Number of landmarks
49
Number of random feature samples
128
Performer CChrromaskii
Type of random feature
Positive
Mode
Fixed
Combiner Ren ct 2021
Span size
49
Conditional distribution purameterization
DeepSets-Max
Dynamic projection dimension
16
Transformer-LS Zhu et al. 2021
Local window size
49
Number of partitioned groups (C)
49
(Ours)
Size of E
49
DeiT-Tiny (Touvron al., 2021). which maintains the sequence length as 196 across all transformer
layers. For the particular tiny variant, the number of transformer layers is set to 12, the embedding dimension set to 196 and the number of heads is
DeiT--mmall (Touvron et al. 2021), which scales the embedding dimension and number of attention heads DeiT-Tiny up to 384 and 6, respectively;
DeiT-Tinn--744 where the architecture is the same as DinT-Tiny but the patch size in the tokenization step decreased from 16 to This effectively increases the sequence length from 196 to 784, which we found consistently improves predictive accuracy at the cost of significantly increased time and memory consumption. Under this setting, we also see clearer differences among these attention variants and it helps better evaluate the ability of different attention models to learn visual representations;
PVT-v2-B3 (Wang et al. 2021b), pyramidal transformer architecture that processes much longer token sequences at early layers and progressively reduces the sequence length to form a hierarchical structure. It patchifies input images into 3136 (56 x 56) tokens, and then processes the sequence through 4 stages. Each stage contains several transformer layers and a down-sampling operation, which reduces the sequence length by a factor of 4 and increases the embedding dimension by
2x. Due to the prohibitively long sequences initially, PVT applies an additional down-sampling module on input sequences to obtain key and value vectors, which are then passed through normal softmax attention mechanism. To evaluate different RF approximations, we remove the down-sane hown-sampligg operation and directly operate on the original sequence length, which results in much fewer model parameters than vanilla PVT-v2-B3. We refer readers to Wang et al. (022 for detailed architecture configurations.
For training, we do not use the [CLS] token for classification (Touvron et al., 202 instead, we pool over the output the last transformer layer to extract features and feed them into the classifier head. We followed the same protocol to train all model variants. Closely following DeiT Touvron cct al. (202T), we employ the AdamW (Loshchilov & Hutter 2019) optimizer to train models for 300 epochs, where the number of warm-up epochs is 10, the learning rate is 0.001 with cosine learning rate decay (Loshchilov & Hutter 2016), and batch size is set to 1024. The adopted augmentation and regularization are the same DeiT, except that we remove repeated augmentation (Hoffer et al. 2020) in DeiT models as it often slows down convergence, as also observed in previous studies (Xiao al.
2021) The specific configurations of each attention mechanism on DeiT-Tiny-785 are listed The hyper-paramenn parameter setup for each attention variant follows previous practices (Wang
in Table
al.
202 b Zheng al. 2022b) closely ensure similar computational cost.
Comparison to State-of-the-Art Model Architectures. We also compare model against recent statc-of-the-art (SOTA) model architectures with similar parameter sizes on ImageNet 1k benchmark. As reported in Table 10, we observe that PVT-v2 (Wang al. 2021b) with EVA greatly
'we retain the repeated augmentation technique in training PVT to be consistent with the original training protocol in Wang al. (02Tbb)
28

---

Published as a conference paper at ICLR 2023
Table 10: Results on ImageNet 1 dataset compared with SOTA model architectures.
Model
# Param.
FLOPs
Top-1 Acc.
PVT-v1-M Wang al. 20
44M
6.7G
81.2
RegNetY-8G Kadosavovic
2020
39M
8.0G
81.7
CvT-21 Wu et al.
2021
32M
7.1G
82.5
SOFT-M (L. al.
2021
45M
7.2G
82.9
RegNetY- Kadosavovic al.
2020
84M
16.0G
82.9
UniFormer-S [t ct
2022)
22M
3.6G
82.9
Swin-S (Liu cf 2021
50M
8.7G
83.0
Swin-B Liu ct al.
2044
88M
15.4G
83.3
Region VT-MM (Chen
2021b
42M
7.9G
83.4
Vil-M (Zhang
2121
40M
9.1G
83.5
Focal-S Yang et al.
202
51M
9.1G
83.5
PVT-v2-63
Zheng al. 2022b
40M
7.7G
83.6
Max (Tu et al.
2022
31M
5.6G
83.6
UniFormer-B (I. et
2022
50M
8.3G
83.9
PVT-v2-b3 (Wang et al. 2021b
45M
6.9G
PVT-v2-b3 EVA
36M
7.4G
improves the predictive accuracy and performs competitively with recent SOTA architectures while using fewer parameters and FLOPs.
E.3 MACHINE TRANSLATION AND LANGUAGE MODELING
Our implementation for all language tasks is based on FairSeq toolkit (Ott et al. 2019). To compare different methods, we report BLEU scores on the test set as the main metric for MT and perplexity for both Autoregressive LM and MLM tasks. For the hyper-parameters |E] and C' in EVA, we set [E] = 2C' by default, as we find that this choice attains good trade-off between performance and computational costs across various tasks; while for C, it is determined based on previous practice for each task. Here we provide the detailed experimental protocol for each task.
Masked Language Modeling. Following the standard pretraining practice as RoBERTa (Liu et al. 2019). in MIM, we aim to reconstruct subset of tokens in the input sequence that are randomly masked out, which is the core element BERT-style natural language pretraining (Devlin et al. 2019). This setting allows us investigate the generalization ability of our model larger model sizes and much more data. The task performance is measured with validation perplexity, which reflects how well the model fits the pretraining corpus and also exhibits good correlations with downstream task metrics. For the used corpus Books 3, we randomly select 100 books without replacement for the validation split, similar the setup in C4 dataset (Raffel al. (020b). For the model, we use the RoBERTa-baas operations (Ba et al. 2016) are placed before attention and FFN blocks (i.e., we adopt the pre-nom architecture), which leads to much more stable training for efficient attention mechanisms, We replace all softmax attention with EVA to test its effectiveness. The training setting and attention-specific parameters, which follow previous studies (Xiong et al. (021aa to ensure similar computational cost, can be found in Table and Table 12 respectively.
architecture (Liu et al. 2019). where all the layer normalization
Machine Translation. We follow Ott al. (2018) to process WMT14 En-De dataset, resulting in around 4.5M/3K/3K English-German sentence pairs for and shared vocabulary obtained between the source and target language of around 32K BPE types. The architecture and training specifics closely follow Vaswani et al. (2017). as listed in Table 13 We encoder self-attention blocks the encoder-decoder Transformer with EVA. For EVA, we find it beneficial to introduce overlapping variant of E, where we allow to be overlapped with each other. Following previous practice in the context of local attention (Xiong et al. aa) E not only contains all elements within the designated chunk but also additionally includes half the tokens in its neighboring chunks. As a result, EVA-32 corresponds to E 32 with contiguous chunk size of 16. During inference, follow the same setup as Zheng et al. (2022b) and average the last 10 model checkpoints to obtain the final model parameters. We apply beam search with size length penalty 0.6, and compound split
splits, respectively.
follow the previous protocol Zheng et al.) (2022b) by replacing
29

---

Published as a conference paper at ICLR 2023
Table 11: Our hyper-parameter configuration for Masked Language Modeling (MLM).
MLM
Number transformer encoder layers
Hidden size
768
hidden size in FFN
3072
Number attention heads
Batch size
256
Sequence length
{2048, 4096)
Number training steps
200K
Number warm-up steps
5K
Weight decay rate
0.01
Peak Learning Rate
Learning rate decay
Linear
Optimizer
Adam
Adam
le-6
Adam (0, .bg)
(0.9, 0.98)
Gradient Clipping
0.0
Dropout
Attention dropout (if applicable)
Table 12:
Our hyper-parameter configuration for different attention mechanisms on MLM task. We used the exact positive random feature map /Choromanski et al. (02T) in our preliminary experiments. However, it failed to converge and exhibited substantial training instability. Therefore, we replace the positive random feature with simple ReLU kernel function for MLM experiments, which yields better training performance.
Attention
Hyper-parameter configuration on MLM
Local attention
Window size
256
Linformer Waag et al. 2020
Projected dimension
256
Number of hashes
Reformer Kitacv al.
2020
Chunk
Number random feature samples
256
Performer Choromanski ct al
2021
Type of random feature
ReLU
LARA Zheng et
20226
Number of landmarks
256
Mode
Fixed
Combiner Ren et al.
2021
Spun size
256
Conditional distribution parameterization
DeepSets-Max
Dynamic projection dimension
128
Transformer- I. Zhu al.
2021
Local window size
256
Number partitioned groups (C)
128
EVA (Ours)
Size E
256
post-processing. Since the input sequences in WMT1 En-De benchmark are much shorter than the other tasks considered in this paper (with an average sequence length around 25 tokens), we start with C' = 8, El 16 and gradually increase El to test the translation performance, similar to the setup in Ma et al. y: Zheng et al.) (2022b). Note that increasing C' also leads to better translation quality, although found the performance gain is slightly less effective than that of increasing E fc.f. T 7).
Autoregressive Language Modeling, We consider Wik iikkiteeee text-11x benchmark in this task, which consists of around 103M/218K/246K tokens for traiming/validatioo//ssinng splits, respectively. We adopt the vanilla transformer decoder architecture (Vaswani et al. 2017). replace all decoder self-attention modules in the Transformer with the causal EVA mechanism, and evaluate EVA under two different setups: 1) standard 16- 16-layer Transformer LM (with model sizes of around 247M) as in Baevski & Auli (2019). ) and 2) a larger 32-layer Transformer LM (with model sizes of around 450M) in Kasai et al. (2021). We follow their hyper-parameter settings to train all models, where
30

---

Published as a conference paper at ICLR 2023
Table 13: Our hyper-parametter configuration for machine translation.
Hyper-parameter
Machine Translation
Number of transformer encoder layers
6
Number of transformer decoder layers
6
Hidden size
512
hidden size FFN
2048
Number of attention heads
8
Maximum number of tokens in a batch
32768
Number of training steps
300K
Number of warm-up steps
6K
Weight decay
0.0
Peak Learning Rate
0.0007
Label Smoothing
0.1
Learning rate decay
Inverse square root
Optimizer
Adam
Adam
le-6
Adam (,1, .82)
(0.9, 0.98)
Gradient Clipping
5.0
Dropout
0.1
Attention dropout (if applicable)
0.1
Table 14: Our hyper-parameter configuration for autoregressive language modeling.
Hyper-parameter
LM
Bacvski Auli 2019
LM in Kasai et
2021
Number of transformer decoder layers
16
32
Hidden size
1024
1024
hidden size FFN
4096
4096
Number attention heads
8
Number tokens in batch
65536
65536
Number training steps
286K
286K
Numher warm-up steps
16K
16K
Weight decay rate
0.0
0.0
Peak Learning Rate
1.0
1.0
Learning rate decay
cosine
cosine
Optimizer
nag
nag
Gradient Clipping
0.1
0.1
Dropout
0.3
0,3
LayerDrop
0.2
Attention dropout
0.1
0.1
the corresponding configurations listed in Table 14 The vocabulary size is 267,744 with adaptive input embeddings (Baevski & Auli 2019). During training, we set the sequence length to 512 and evaluate the validation/test PPL with various context window sizes in {256, 480), aligning with previous work (Baevski & Auli 2019 Kasai et al. 2021). For other random feature baselines, unfortunately. we failed to fully replicate their results as reported in Kasai et al. (20TT), where RFA in our implementation achieved a test perplexity of 29.0 even under a 449M Transformer model. For EVA, we set E 128 and C 64 by default for both 16-layer and 32-layer settings, ensuring similar computational cost to previous work that also evaluates random feature methods (typically with 128 or 256 random- deddure dimension size) on Wikitext-10 language modeling task (Schlag et al. 2021 Kasai et al. 2021).
E.4 EXPERIMENTAL SETTINGS OF EFFICIENCY COMPARISON
For the simulation experiment conducted in $5.3 we adopt the same transformer architecture across all attention variants. In particular, it uses transformer layers, 192 embedding dimensions, and attention heads so that longer sequences can fit into our devices. The batch size is set to 64 across
'The setup in Baevski & Auli (011)) can be found in the corresponding Fairseq train- ing script: https: / con/ model 2/de1/RHAAMM: adapt wwtp:/ottttptc/teeeeeeeee hawpc//aaarrrg/bbba ob/masterr/exmmllee... aanguuge
31

---

Published as a conference paper at ICLR 2023
Linformer
Linformer
Local Attention
Local Attention
Memory Consumption (GB)
Long-short short
Running Time ms!
Long-short shart FlashAttention
FlachAttention
Softmax
Softmax
1000 2000
3000
4000
5000
5000
7000
8000
0
1000
2000
3000
4000
5000
6000
7000
8000
Sequence Length
Sequence Len gth
Memory consumption.
(b) Running time.
Figure 2: Left and right: Additional empirical memory consumption and running time comparison for different attention mechanisms under various sequence lengths.
8 V100 GPUs, and the statistics are computed by averaging the results of 30 runs. Besides, in our ablation study, the efficiency metrics reported in Table and Table 7 are evaluated under the same setup used during training.
Remark on Modeling Short Sequences. Unfortunately, similar to most previous efficient attention baselines, EVA also runs slower than softmax attention under shorter sequences (e.g., length 128 or 256), but soon catches up in rumning speed, and the reduction of memory consumption is still significant. Besides, in short-sequence settings (such as the case DeiT- TirT/Sil/// with sequences of 196 tokens), EVA often performs on par with better than conventional softmax attention (see Table T}. whereas most previous attention variants usually perform much worse. This implies EVA can achieve better trade-off between efficiency and quality: for short sequences, EVA possible to achieve stronger performance competitive with softmax attention (despite in longer running time); while for long sequences, EVA can be run much faster with less memory.
Comparison to Memory-efficient Attention Mechanisms. In this section, we conduct an empirical efficiency comparison between efficient approximate attention methods and PlashAttention one of the memory-efficient attention mechanisms (Rabe & Staats [202T] [ao et al.) 2022) with optimized memory accesses. FlashAttention computes the exact softmax attention in an online manner without materializing the full attention matrix, achieving linear memory complexity with respect to sequence lengths; besides, both runtime and memory usage are further improved by minimizing IO accesses. We benchmark different attention modules on one NVIDIA GeForce RTX 3090 GPU, where we measure the memory usage and runtime of running single attention block, consisting of 8 attention heads with 512 embedding dimension size, for both a forward and backward pass. As shown in Figure we observe that FlashAttention achieves significant memory usage reduction for softmax attention approximation and even consumes much less memory than all considered approximate baselines under all sequence lengths. terms of runtime, we notice that FlashAttention runs faster than most attention baselines under sequence lengths less than 2048 despite scaling quadratically, but EVA, along with other more efficient approximate variants, begin catch up longer sequence lengths. This implies that the quadratic computational costs of softmax attention still bottleneck its runtime performance, aligning with one of the main findings Dao et al. (2022). According this empirical study, we observe that Flash Attentio offers general and effective technique to speed up softmax attention; since many approximate variants (including EVA) exhibit similar formulation to softmax attention (c.g., Equation 16), we expect they can also benefit from the optimized online softmax calculation technique and memory accesses of FlashAttention (Dao et al. 2022)
F EXPERIMENTS ON LONG RANGE ARENA
Long Range Arena (LRA; Tay al. 2021 is lightweight benchmark that assesses the ability of efficient attention methods to model I long sequences in diverse domains. We follow the same hyper-parameter setup as Xiong al. 2021bb to re-evaluate all attention baselines and report the
32

---

Published as a conference paper at ICLR 2023
Table 15: Classification accuracy (%) on LRA benchmark with different efficient attention mecha- nisms.
Model
ListOps
Text
Retrieval
Image
Pathfinder
Avg.
Softmax
38.66
64.91
80.70
40.61
68.29
58.63
Linformer
38.21
53.91
77.66
39.40
66 44
55.12
Performer
29.84
65,30
77.70
38.29
66.39
55,50
Reformer
27.12
63.90
78.08
42.40
51.90
52.69
Scatterbrain
38.21
64.04
77.83
42.51
60.62
56.64
Combiner
38.26
63,98
81.47
42.80
55.94
56.49
LARA
37.10
64.62
80.82
38.99
68.96
58.10
Nyströmformer
38.46
65.28
80.44
39.71
68.98
58.57
Local
38.46
63.70
80.71
42.25
68.46
58.72
Long-short
38.56
63.46
81.73
40.54
71.28
59.11
EVA
38.61
64.31
80.21
43.24
59.45
comparison in Table I3| We observe that EVA largely improves previous RFA methods such as Performer (Choromanski al., (02T) and LARA (Zheng et al. 2022b), and performs competitively with full softmax attention. Notably, EVA even achieves better average results over all tasks, with higher accuracy Image and Pathfinder benchmarks, suggesting its capability of capturing long-term dependencies. For LRA benchmark, we set all attention hyper- parameterss to 128 (c.g., e. the number of landmarks in Nyströmformer (Xiong et al. 2021b) and LARA (Zheng et al. 2022b), the window size in local attention and Combiner (Ren et al. 2021), etc.). We set E|| = 128 and C = 64 by default for EVA without any further tuning and find this setup works well.
G CONNECTIONS TO OTHER ATTENTION MECHANISMS
G.1 RFA, SOFTMAX ATTENTION, AND EVA
As mentioned in our main text, one the main contributions of this work is to develop more general framework that bridges RFA and conventional softmax attention. To see how EVA (Equation [3] achieves this goal formally, note that if either |E| M or C = M, EVA would be equivalent to standard softmax attention; while if we
and C EVA would recover vanilla RFA.
G.2 CONNECTIONS TO LARA
Notably, EVA and LARA Zheng et al.) 2022b two efficient attention mechanisms that are both built upon the self-normalized importance sampling (SNIS) formulation of RFAs. LARA (Zengg al., 2022b) puts the main focus the proposal distribution used in SNIS and tries design importance sampling proposals that are closer to the true underlying distribution. The proposed usage of multiple proposals further improves the estimation quality of SNIS and achieves strong empirical performance while still keeping linear complexity.
In contrast to LARA, in this work we do not focus on the design choice of proposals used in importance sampling but aim to generalize the SNIS formulation further via control variates. As demonstrated in $3.2 our theory clearly delineates how the gap between such SNIS estimation and softmax attention can be closed by manipulating control variates. Since LARA and RFA are both SNIS estimators (their main difference lies the choice proposal distributions), our generalization also applies LARA. To summarize, compared with LARA, EVA is more general framework and improves conventional REA from an orthogonal perspective.
G.3 CONNECTIONS TO CLUSTERED ATTENTION
Clustered attention VVas et al. 2020) is an efficient attention mechanism that first clusters the set of queries into multiple groups, computes the mean centroid of each group, and then performs attention between query centroids and original key-value pairs. This framework is fast and effective and enjoys well-hounded approximation error.
33

---

Published as a conference paper at ICLR 2023
Clustered attention and EVA share some similarities in two aspects. First, both of them adopt the partitioning technique to reduce the computational complexity while remaining effective; and secondly, both observe that the efficient attention mechanism can be improved by refining the approximation over specific elements. For instance, clustered attention can be improved (Vyas et al. 2020) by selecting top-k key-value pairs that most relevant each centroid and then refining the approximation by recomputing attention weights over these keys using original queries; while EVA notices that we can directly employ the optimal control variate coefficient for subset of key-value pairs (m E) while still remaining efficient, which yields a more accurate approximation.
Nevertheless, our main technical contribution is develop control variate formulation the context of RFA and demonstrate that how RFA can be further improved locally. On the other hand, while clustered attention (Vyas et al. 2020) clusters queries, EVA partitions key-value pairs. This property makes EVA more amenable to the case of autoregressive language modeling since we do not impose clustering structures over the query set, and thus the causal relation among queries can be well maintained.
G.4 CONNECTIONS TO COMBINER
Combiner (Ren et al. 2021 is recently proposed attention mechanism that also partitions the sequence into chunks combined with local attention. The key difference between EVA and Combiner is the motivation, where Combiner introduces a structured factorization over the attention probability distribution, while our approach built from the control variate perspective.
G.5 CONNECTIONS TO SCATTERBRAIN
In this section, show that Scatterbrain (Chen al. (021aa can be cast special case of our framework EVA, although they are proposed based on quite different motivations.
Brief Review of Scatterbrain. Scatterbrain (Chen et al. 2021a) notes that sparse attention and RFA can approximate sharp and flat regions of the softmax attention matrix well, respectively. Based on this insight, Scatterbrain proposed to first compute a Performer approximation to softmax attention and then cancel out the approximation error critical regions via sparse mechanism.
Specifically, Scatterbrain (Chen et al., 2021a) defines sparse matrix E I Nx) so that for each (n, m)) E S that indexes non-zero entry. For notational simplicity, we also denote Supp(S) = (i.j))(S * 0} and Supp,, (S) = [m]Sum f 0}. With random features ( defined in Appendix A we let
Sam= = exp (g ())n)))))))))))))))))))))) (k(km, w).
We then add it back to the approximate output:
M
s/ = Σ
p(q,: w) (kkmm w)vm + +V
m=1
M
Σ
p(q,, w) w)kk
Sam V Vm
m=l
m'' ErEpupp (S)
Φ(qn, Φlkm, w)
exp (qTh km)) Vmn.
(29)
mgSupp,, (S)
m (nEEuppp..S)
The sparse mechanism can be thought of as modeling the enror due to RFA and eliminating it on the support S. After the correction step, Scatterbrain further adds a post-hoe normalization step to obtain normalized attention output:
Vn Emgsupp, (S) (b W (bkm, w)v + Em Em''Suup, ESupp,, (S) exp (q, (q,]k,,, kmm) Van
(30)
Emesupp, (S) Φ(9,, 9n, w) (0km,, w) + Em'
ESupp, (S) exp (q, k
Intuitively, Scatterbrain (Chen et al. 2021a) produces accurate approximation the support of the sparse matrix and remains the random feature approximation outside the support.
34

---

Published as a conference paper at ICLR 2023
Scatterbrain is Special Case of EVA. For notational convenience, we denote E == Supp,, According to Proposition 1 suppose we employ optimal coefficients f... for all entries in Supp,, and use the same coefficient / for all the remaining entries (in other words, we let C = 1 and the whole index set is only partitioned into two subsets {F.
[M] E.). Then have
Im (w) = [m((w) (w) - hm (w) + 3m expla, km)
expla, Vm
if m E E, if m & E.
2
[m(w)- - 3hm (w) + 3 egexppa
2
And the resulting estimator overall becomes
[(w)
Im
(w)
Wi
Σ mEE
Fem Fmm(w) (w) + Σ
Fm (w)
mgFF
exp( an km Vm
Im(w)) - (lmm(w) - +33 eepp(a q, km)
Z
Z
mEE
mgth E
exp(q, qn km Vm
exp(q km
Σ
Σ
(m(ww) - 3mm(w)) + Σ
Z
Z
MEE
męE
mEE
exp qn, km
Vm
(gm(w) Blm((w) +
1- Σ MEE
exp km)
Z
Z
MEE
E
Scatterbrain (Chen et al. proposal distribution to ((w) N(w I), and estimate the normalizing constant as follows.
2021a) can be a special case of this estimation algorithm if we set the
Niww; 0,0; 1) (EmEg w) 5(km; ((w)
( W) + EmpE E(qn, ui) E(Kma w))
7 (wwa(())
((w;0,1) me E(gm T w) q(w)
Σ exp(q, km) + (uwng(w)
MER
exp(q km) km) + N(w:0.t)
s
m E(g, g(we)
T (k)
S
MEE
8=1
s
Σ
exp(qx kmn) + Σ
Σ
E(qn, T
E(kms w)
MEE
m
Σ
exp(q, km) +
Σ (().:w)
(kkmm w)
MEF
m&E
Σ
exp(qn T km) + Σ 7mm (w),
mEE
mg
where we define Ti (w) Zhm (w), in this case
1
Pn (ws)
1 s
M
g(w)=
f(we)
Σ
E(g,. Wa (kkm wo) V
S
((ws)
S
Z
I=1
==1
m=1
h(w) s
1
Pn (ws)
1
1
Σ
M
E(g.. (5)
((ws)
S
Z
m=1
35

---

Published as a conference paper at ICLR 2023
With these specifications, obtain
exp(dn kmmmm
exp km)
((w) = Σ
Σ
[m((w) Blmm (w) + 1 Σ
Z
Z
mEB
mg
mEB
exp(qe In kmm/ive Vm
Emer E exp(qn) km)
+
Σ
(m((w) B31m +
Z
Z
mER
m
exp( 4n km Vm
(m(ww) 331m (w))(w)) +3 |Eme E km(lw)
Σ
Z
mEB
mgE E
exp( In kmmmm
Vm
+
(m(ww) 331m (w) + Σ km(w))
Z
MER
mith E.
mg
E EMEE exp(
9m km Vm
Σ gm:(w)
Z
Emer exp(qn km))vm
E E(g Wa (kkm V
Z
męE
EEER explq, km)Vm
(Ln,, w) T "olkms wm
Z
Z
m&E
Emer ((q,, w) 6lkm, W)V + Em EE exp (q (ui km/) Vmn
2 Empe Φ(qn, In: w) (Kk,; w) Em' exp (q, km))
(31)
which is equivalent to Scatterbrain (Fquation 30). Note that this equivalence would hold irrespective of the choice of shared coefficients B, which possibly indicates that the formulation of Scatterbrain limits the potential benefit of optimizing control variates under our framework.
36


Reviews:
[{"review_id": "pPVnSINJ3TH", "reviewer": "Reviewer_xJRA", "summary": "", "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": "", "summary_of_the_paper": "This paper studies the gap between random-feature-based attention (RFA) and standard attention. The paper characterizes the gap using control variates and proposes a novel efficient attention mechanism based on the analysis. The authors evaluate their method EVA in terms of model quality, efficiency, and runtime.", "summary_of_the_review": "Overall I think this paper makes a positive contribution to the literature. The results also show minor improvements over previous efficient attention mechanisms. A few gaps and baselines remain, but I think this paper is above the acceptance threshold.\n\nUpdate after rebuttal: my concerns have been addressed, so I am increasing my score to an 8.", "strength_and_weaknesses": "+ A nice analysis of RFA that provides some insight into what the method is doing.\n+ Interesting method building on the analysis.\n\n- Some weak empirical results suggest that this method may not scale to larger models (this is common for RFA-based efficient attention mechanisms).\n- A couple missing baselines.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VUfWtJ4G-zy", "reviewer": "Reviewer_B8ci", "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": "This work improves Random-feature-based attention through the lens of control variates. The paper develops a more flexible form of control variates, which forms a novel attention mechanism that significantly reduces the approximation gap while maintaining linear complexity. The EVA achieves competitive results on several tasks. ", "summary_of_the_review": "In this paper, the author proposes an efficient attention mechanism called EVA with the help of control variates.", "strength_and_weaknesses": "Strength:\n\n- The paper develops an efficient attention mechanism EVA via control variates. This novel attention mechanism is efficient and can bridge the gap between RFA and exact softmax attention. The EVA attains a good trade-off between modeling quality and efficiency. \n\n\nWeakness\n\n- The author should compare more state-of-the-art linear attention like Flowformer or FLASH.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "OHNftQbYjT", "reviewer": "Reviewer_rNYj", "summary": "", "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": "", "summary_of_the_paper": "In this work, the authors propose an efficient attention EVA, where the work reals that exact softmax attention can be recovered from RFA by manipulating each control variate. The mathematic analysis is provided to prove the dissecting RFA with control variates. And the authors also implemented their EVA and applied it to image classification and natural language tasks, showing performance boosting over SOTA. ", "summary_of_the_review": "Though in terms of accuracy/#param/Flops, the EVA might not beat SoftMax completely, the work with theoretical evidence for approximating the  SoftMax attention can provide insights into the future works of light-weight attention development. The paper is in well-written, and the experiments are solid. Hence, the work reaches the acceptance bar. ", "strength_and_weaknesses": "Strength\n- The idea is easy to follow, and the paper is well-written. \n- The mathematic analysis is given to explain the dissecting RFA with control variates.\n- The EVA attention with simplification achieving higher accuracy than softmax in image classification is impressive. \n\nWeakness\n\n- Overall, I would say the paper is in good shape in terms of detailed mathematic analysis and solid experiments on multiple vision and natural language tasks. My main concern is about the usage in the practical. Though the EVA achieves higher accuracy (very subtle) than softmax in table 1, the softmax is still a better choice in terms of accuracy, #params, and FLOPs. According to numbers reported in all tables, softmax still dominates the performance and has comparable #param and FLOPs to EVA. \n- The not-well-designed approximation in computing control variate could impact the EVA's performance. As no code is provided, it is hard to trace the author's implementation. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "MeJxV4YxuP", "reviewer": "Reviewer_KXnA", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Firstly, as in Zheng 2022, the authors interpret the computationally efficient RFA (randomized feature-based attention) as a SNIPS estimate (self-normalized importance sampling estimate). They then use Vlassis' control variate formulation of SNIPS to link exact attention with RFA via a single global control covariate coefficient. This naturally suggests using more control coefficients to obtain more accurate coefficients while maintaining efficiency in the form of exact coefficient estimates for certain keys most correlated with the query and coefficients per cluster for the remaining keys. This new approach is termed EVA and empirical demonstrations of its efficacy are quite convincing. On ImageNet using the DeiT-Tiny-784 architecture, EVA outperforms Zheng's LARA approach and other forms of efficient attention in Table 2, while in Table 1, using a different architecture, EVA outperforms all efficient attention variants, but not the slower exact attention mechanism itself. In 3 NLP applications (MLM, MT, Autoregressive LM), EVA again outperforms other efficient variants in terms of BLEU score or test perplexity, although there were issues with replicating Kasai's perplexity for exact attention.", "summary_of_the_review": "I strongly recommend acceptance on account of the novel theoretical results and the extensive convincing experimental results.\nIt should be straightforward to address the minor concerns regarding clarity and reproducibility noted above.\n", "strength_and_weaknesses": "Strengths\n\n1.\nThe proposed EVA approach is motivated quite well and the paper is quite clear.\n\n2.\nThe theoretical results (Propositions 1, 2 and ) in the paper are impressive and appear to be quite sound.\n(I didn't check the extensive appendix though.)\n\n3.\nThe experimental validation is quite convincing on a variety of problems.\n\nWeaknesses\n\n1.\nCertain choices are not well-motivated. For example, rather than learning the partitioning scheme via a K-means-like algorithm for the keys receiving clustered coefficients (instead of optimized control coefficients), the clusters (partitions) are chosen as contiguous chunks. This is called out in the final section though.\n\n2.\nIt's not clear how hyper-parameters such as E and C in EVA were tuned.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-5%