Learning the Transformer Kernel
Authors: Sankalan Pal Chowdhury
Adamos Solomou
Kumar Avinava Dubey
Mrinmaya Sachan
Venue: Accepted by TMLR ()

Abstract: In this work we introduce KL-TRANSFORMER, a generic, scalable, data driven framework for learning the kernel function in Transformers. Our framework approximates the Transformer kernel as a dot product between spectral feature maps and learns the kernel by learning the spectral distribution. This not only helps in learning a generic kernel end-to-end, but also reduces the time and space complexity of Transformers from quadratic to linear. We show that KL-TRANSFORMERs achieve performance comparable to existing efficient Transformer architectures, both in terms of accuracy and computational efficiency. Our study also demonstrates that the choice of the kernel has a substantial impact on performance, and kernel learning variants are competitive alternatives to fixed kernel Transformers, both in long as well as short sequence tasks.

Published in Transactions on Machine Learning Research (07/2022)
Learning the Transformer Kernel
Sankalan Pal Chowdhury
spalchovd@inf.cth..dd
Department of Computer Science
ETH Zürich
Adamos Solomou
solomou.adamos@e @mmaill.ommmmmmmmmmmmmmmm
Department of Computer Science
ETH Zürich
Avinava Dubey
avinavadubey@
Google Research
Mountain View, CA
Mrinmaya Sachan Department of Computer Science
mrinmaya.sachan @ inf.ethz..c
ETH Zürich
Reviewed OpenReview:
Abstract
In this work we introduce AILFORANSFORMER, a generic, scalable, data driven framework for learning the kernel function Transformers. Our framework approximates the Transformer kernel as a dot product between spectral feature maps and learns the kernel by learning the spectral distribution. This not only helps in learning generic kernel end-to-end, but also reduces the time and space complexity of Transformers from quadratic to linear. We show that KL-TRANSFORMER. achieve performance comparable to existing efficient Transformer architectures, both in terms of accuracy and computational efficiency. Our study also demonstrates that the choice of the kernel has substantial impact on performance, and kernel learning variants are competitive alternatives to fixed kernel Transformers, both in long as well as short sequence tasks.
1 Introduction
Transformer models (Vaswani al., 2017) have demonstrated impressive results on a variety of tasks dealing with
language understanding (Devlin al. 2019 Radford al. 2018 Raffel al.) 2020 Brown al., 2020), image processing (Parmar et al. 2018 Carion et al.
2020 Lu et al. 2019 well biomedical informatics (Rives et al. 2020). Albeit powerful, due to the global receptive field of self-attention,
2020; Ingraham et al. 2019 Madani et al.
the time and memory complexity of Softmax Transformer models scale quadratically with respect to the sequence length. As a result, the application of Transformers to domains with long contexts is rather limited. This limitation has spawned several efficient Transformer designs (Liu et al. 2018 Parmar et al. 2018 Child et al. 2019: Zaheer et al. 2020; Beltagy al., 22200) Roy al., 2020) Tay al. 2020aa] Kitaev et al. (2020). Kernelization offers one such design. The use of kernel feature maps allows reformulate the computation of attention in way that avoids the explicit
computation of the full attention matrix which is the key bottleneck for Softmax Transformer. This also opens up
directions for more generic attention mechanisms.
Tsai et al. 2019) first proposed kernel-based formulation of the attention mechanism. However, the time and memory complexity their approach remains quadratic with respect to the sequence length. To address this limitation, Katharopoulos et al. (2020) expressed self-attention as the inner product of kernel feature maps and made use of the
Our code models available https::tps/.. /

---

Published in Transactions on Machine Learning Research (07/2022)
associative property to reduce the complexity from quadratic to linear. For their experiments, they used the arbitrarily chosen LinearFlu feature map f(x)) = max(x+1, e."). Performer (Choromanski et al. 2021b) replaces this with feature maps that can directly approximate the softmax kernel, thereby allowing the use of pre-trained Softmax Transformer weights in a linear time model. Concurrently with then et al. (202T) proposed a linear space and time method that added causal and recency based features to random Fourier methods. More recently. (acllgghhhghhh et al. (202Tb) showed the formal equivalence of linearized self-attention mechanisms and fast weight programmers. While the aforementioned approaches provide a significant reduction in computational and memory requirements, this often comes at the cost of performance, as can be seen from Fig. [ In this work, posit that this is partly due to the fact that the similarity functions/kernels, including gagled-dot-produeee
were hand picked and not learnt from data. Thus, we explore whether kereel learning can help to bridge this gap in performance while retaining the scalability of efficient Transformers.
Although, to the best of our knowledge, kernel learning has never been explored within the framework of Transformers, kernel learning methods have been an ubiquitous tool in machine learning. The most notable among them is Random Kitchen Sinks (RSS: Rahimi Recht 2007). a data-independent framework for approximating thift-invariant kemels using an explicit feature map. In RKS, the kernel is approximated by <(2r, y) R (d(x)), ((y),, where the explicit feature map Φ _ + Rx is obtained by sampling from spectral distribution defined by the inverse Fourier transform of the kemnel function K. Willson& & Adams (013) modeled the spectral density as mixture of Gaussians, A la Carte Yang et al. 2015) proposed an optimization based framework for learning the spectral distribution, BaNK (Oliva et al.) 2016 modeled the spectral distribution using an infinite mixture of Gaussians, while Fang et al. (2020) implicitly modeled it using deep generative models. We build on these advances and incorporate them into the Transformer framework.
Contributions: In this work, we propose TRRNRROSFRRRRR a scalable data driven framework for learning the kernel of Transformers and investigate whether fully learnable kernel can help im- prove the performance of linear, fixed kernel Transformers. Thus, we introduce Transformers with learn- able similarity functions, which happen to retain the linear complexity in terms of the sequence length. We motivate our learning method using RKS and learn the kernel by learning the corresponding spectral distribution. In $2.T]we first propose to learn generic Transformer kernel by explicitly approximating the spectral distribution using Mixture of Gaussians (GMM) and propose modifications to scale it further. In an attempt to further explore the trade off between computational complexity and accuracy we also propose to model the spectral frequency distribution of Trans- former kernels implicitly by using deep generative models Goodfellow et al., 2014). Finally, we also propose a novel method to learn the spectral distribution of positive random feature (PRF) maps, which provides better approximation of the softmax kernel (Choromanski al. 20211b.
Transtormers Vaswaui 2017
Memory Consumption (GB)
Symthsiierr Tray 20222
BigBird Zeboct ol al. 2020
Reformmer Kimee et al., 20120
Local Alt.,
Sinkhorn. Tay et al., 201200
Parmar 2018
This Paper
Lizentlln, Kelloropeulos et 20,0
Patsermer, Charomanaki et 2021
0.5
We analyse the expressivity and precision of our proposed models ($2.2) and show that the proposed GMM with posi- tional encodings is Turing-compler
55
Average Accuracy
(Pérez et al. 2019) with controllable variance. We experimentally evaluate our mod- els on LRA (tasks with long context), GLUE (tasks with short context) and a synthetic dataset with controllable sparsity, and analyze the performance of our models ($3 22.2) In our experiments, we find that learnt kernels improve perfor- mance in long-contexs tasks, while staying competitive to the Softmax Transformer of the same size in short-context tasks. We also find that our models learn parameters that reduce the variance in their predictions, and can handle sparsity quite well. We also benchmark the computational efficiency of KL- TRANSFRRMERS and find that each of our proposed ST.-TRANSFORMER scales linearly with respect the sequence length. We conclude with short comparison between Random Kitchen Sinks (RKS, Rahimi & Recht (007)) and Positive Random Features Choromanski al. 2021b) in terms of their performance and provide recommendations on which approach should be chosen under which circumstances.
Figure 1: Peak memory (y-axis), average performance (x-axis) and speed (denoted by area of circle) for var- ious efficient Transformer models (i.e. bigger circles in the bottom right corner are better) across three long sequence tasks (> 1024 tokens) introduced in LRA (Tay et al., 2021b). All values except for "This Paper" are taken from Tay al. (20216).
2

---

Published in Transactions on Machine Learning Research (07/2022)
2 Kernel Learning in Transformers
We begin with the generalized formulation of self-attention proposed by [ssa et al. (2019). Given non-negative kernel function <(;-): R. x dR - 4. the output of the generalized self-attention mechanism at index i operating on an input sequence X = (I1,..., x1 E Lxd is defined as
ATTN(X); = Σ k(g;,ki)) ES n(g,, ky))
V..
(I)
where h; = z,W 9% = x; x,,,,, 8; = I, WV are linear transformations of the input sequence into keys, queries and values of dimension dg = dr and dy respectively and W E IRdxds, Wº E RRdxd, WV E Rdxd. While the formulation in Eq. ① is more generic and defines a larger space of attention functions, suffers from a quadratic time and memory complexity. To reduce the quadratic time and memory, we briefly review the method of random Fourier features for the approximation of kernels (Rahimi & Recht, 2007). The details of the method will help motivate and explain our models.
Random Fourier Features for Kernels: At the heart of this method lies the theorem of Bochner (Rudin) 1990)
which states that a continuous shift invariant kemnel k(q, k) = k(g - k) over arbitrary variables q and k is a positive definite function if and only if R(6) is the Fourier transform of a non-negative measure p(w). Moreover, if (0) = 1, then Bochner's theorem guarantees that p(w) is normalized density function, i.e.
[ p(w) exp (iw" (q - k)) dw Buw [exp(in" q) expp(uw> RR)].
k(g k:)
Rahimi & Recht 2007) proposed to sample from the spectral density p(w) for a Monte Carlo approximation to the integral in Eq. Specifically, for real valued kernels, they define x(q, k)) ss s(a)² (o(2)"((k "(k)), where wi p(w) and
((x) = (K, = "olt....,.M = VM
[cos(w] :)....
(os(wTra)) sin(wT cos ,i)_....ii((()))))))
(3)
To learn a kernel, we can either learn a parametric function k(::.) or learn the corresponding parameterized feature map (-) directly, which corresponds learning the spectral density p(w) (Wilson & Adams 2013 Yang et al. 2015
Oliva et al. 2016). In this paper, we focus on the latter because this helps us in keeping the computational complexity
d(u:)" E
E *(ky)e,)
linear in the sequence length L. This can be achieved by rewriting Eq. (1) as ATTN(X); d(a))" Ei #(ky)) To
the best of our knowledge this is the first attempt to learn the kernel of the generalized self-attention mechanism (Eq.
2.1 Learning Kernels in Spectral Domain
Our objective is to enable learning of any translation invariant kernel. This is realizable if we can leam the spectral distribution. Gaussian Mixture Models (GMMs) are known universal approximators of densities and hence may approximate any spectral distribution. GMMs have been shown to be useful for kernel learning for regression and classification tasks (Wilson & Adams 2013 Oliva et al. (0166). Thus, to learn the kernel of the generalized self-attention mechanism (Eq. we model the spectral distribution of the kernel as a parameterized GMM, i.e.
GMM-RKS:
p(w) Σ NoNe Ee) x(q,)) k)
Σ The meelit (a (0- (( F) k))
(4)
Here {He E IRd, Eg E 12 IC are the learnable parameters of the feature map and C' is the number of components in the Gaussian mixture. It can be shown using Plancherel's Theorem that p(w) can approximate any shift invariant kernel (Silverman 1986). Since we are working with only real valued kernels, the corresponding kernel reduces to x(q, k)) = -Ec Tee " kx(g-k))) cos (ni (q k)).
To speedup learning, we assume that Te J/ and parameterize the feature map with spectral frequency, Ω (w..i....,,C, as:
DKSM-RKS(2 = RK 2), Wem = Em + Pe: n'm ~ I).
(5)
This allows us to sample R'm N(O0I) and learn the parameters of the feature map. (lue E Rdh, E. E R end-to-end along with the other parameters of the Transformer.
3

---

Published in Transactions on Machine Learning Research (07/2022)
FASTFOOD-RKS: GMM-RKS removes the quadratic dependency on context length, but we still need calculate Ω Q and 1 (where Si = [w) .4....., (w)) ,..) ) which takes O ((Md, time and ( + d, + MT. space, which can be too much if M is large. For further gains in scalability, we approximate the spectral frequency matrix S2 using the product of Hadamard matrices (FastFood: Le et al. 2013), such that the computation can be done in time log-lineat in M, i.e.:
OKASTFOOD-RKSSA == RKS(2. V), where V
SHGHHBB. o oddg
(6)
Here, IT E (0, 1)d,xd, is permutation matrix, H is the {+1} entries, G' is a diagonal matrix with Gaussian entries drawn from (0,, 1) and finally S is a random diagonal scaling matrix that makes the row lengths non- uniformo..m The entire multiplication can be carried out in logarithmic time, and the space requirement reduced by storing diagonal matrices instead of full matrices. For M > dg we use multiple blocks, and the only restriction is that we need M (dg- In order to make this learnable, Yang et al. (2015) proposed making S and optionally G and B learnable. For the main paper, we keep all three learnable (the case where only S is learnable is discussed in Appendix D).
matrix, B is diagonal random matrix with
GENERATIVE-RKS: If we increase the number of components (C) in GMM-RKS, the computation and space complexity increases dramatically. Instead, to learn more generic kernel, without blowing up the computational complexity, we use deep generative models (DGMs). DGMs have achieved impressive results in density estimation (Goodfellow et al. 2014 Kingma & Welling 2014 Richardson & Weiss 2018 Ruthotto & Haber 2021) and end-to-end
kernel learning in the spectral domain for classification (Fang et al. 2020).
In GENERATIVE RRKS we use a DGM to replace the Gaussian probability distribution from GMM-RKS with an arbitrary probability distribution. This DGM acts as generator similar to the ones used by variational auto-encoders used in computer vision (Kingma & Welling 2014). In particular, to make the sampling process differentiable, we use the reparameterization trick, where learnable neural network (called the generator, and denoted by in Figure transforms samples from simple noise distribution (po in Fig into samples from the desired distribution.
Wm g(nm),
Pen pa(-)
The generator network is trained end to end with the whole model, allowing to choose the best possible distribution for the given data. These samples (S2 = [WI, (2, in Fig 2) are then used in Random Kitchen Sinks as follows:
ÖGENERATIE--RKS (x) = RKS (x, S2), Wm ~ g(po)
We experimented with various configurations and eventually and chose to learn a generator network which consisted of fully connected layers with batch normalisation and LeakyReLU activation, followed by a single fully connected layer with tanh activation. While this methodology allows us generalise the gaussian distribution to a much larger class of distributions, it also causes blowup in the number of parameters e.g. a a + 1 layer constant width generator, as used by would require 5d + 5d parameters as opposed to the d² + d parameters in GMM-RKS. To counter this and to improve generalization, we share the same generator network across all heads in a layer, which means that the different heads only differ in the Keary/Key/Vallue projection matrix.
Positive Random Features (PRF): Until now, we focused feature maps defined using RKS. While our formulation is very general, recently it was shown that positive random features provide a better approximation to both Gaussian and Softmax kernels (see Lemma 2 in Choromanski et al. 2021b). In particular they showed that k(g, k) = exp(q" k) (o (ex.:/o..njexp(uTTT [ q
|| exp(w>T k [[]][)) ){ and demonstrated that Monte Carlo approximation to this expectation leads to a low variance estimate of the softmax kernel. Moreover, the presence of only positive values within the randomized feature map ensures that kernel estimates remain strictly non-negative. To incorporate this prior knowledge, propose novel kernel learning framework in which we learn the spectral density while using the feature map corresponding to the above expectation. For instance, when we model the spectral distribution of S = (Wi,... WM)
2
using GMM (1b- Σ (He, Fx.) we have that:
4

---

Published in Transactions on Machine Learning Research (07/2022)
dg
RKS
2M
w
Query Q
M x dg
gen (Q)
Lx
Lx dq
Samples
Generator
g(.)
Lx d
w.k
Noise Po
Frequency f
x
Output
Input X
Key K
W
RKS
Linear
Value V
4gnn KK)
In I.
Figure 2: Generalized self-attention with deep generative RKS. Q, V and are linear transformations of input, X. The generator generates spectral frequency (I2) from implicit spectral distribution. Using RKS (Eq. ( create the feature map @ gen (Eq. 8). The numerator of the output is calculated &gen (Q) (K)V) while the denominator is Ögen (q.)) Ej'=1 gen (ky) making attention linear in sequence length
Model
Space Complexity
Time Complexity
Softmax Transformer
D(L'(1+d//D) O(L(d, +M Md,/L)) O[L(d,+d//))
O(L?dg)
Performer
OLLMd,)
LinearElu
O(Ld)
GMM-RKS
O[L(d, C(d//L+ C +d./L+
OMM + Ld,)) O(LMMd,)
GMM-PRI
+ + CM))
FASTFOOD/RKS/PREF)
O(L,d, +M+ ML)/L))
O(LMd,)
GENEBATIVE/RKSS/PRTT)
O M
O(M Ldg))
Table 1: Space and time complexity of self-attention kernel of KL- T.-TRRANSFORMRS compared with Softmax Trans- former Vaswani el al. 2017), Performer (Choromanski et al. 2021b), and LinearElu (Katharopoulos et al., 2020).
I. refers to length of context, d,,=d, is the query/key/value dimension, while M the number of samples (where applicable).
e11121|
k(g, k)) = of 9-11)1
y- (( (-W11111)))) PRF(x, = VM le
c
(9)
PRMM- = PRF(I, (We, (w..
Wom = Em + Me: n'm ~ N(O.1)
(10)
Similarly we redefine the other two methods as:
@GENERATIVE-PRE PRF(x, Wm = g((m)) 1mm pol.)
(11)
(RF())) = PRF(x, V), where V SHGHHBB.
(12)
odd
To the best of our knowledge the first to explore kernel learning with positive random features.
2.2 Analysis
In this section, we explore what can be said about the expressivity of the proposed linear time KL--RANNFORRMERS While our understanding of the properties of Transformers is rudimentary, we would still like to know whether the known properties extend our models. For example, it has been shown that Softmax Transformers and its sparse This also Turing complete?
counterparts are Turing complete (Pérez et al., 2019 Zaheer al. 2020). This raises the question as whether the proposed linear
5

---

Published in Transactions on Machine Learning Research (07/2022)
It is easy to that the generalized kernel self-attention (Eq. includes the softmax kernel and hence should satisfy the properties of Softmax Transformer. Interestingly, we can also show that this property holds for GMM-R Transformers with number of components C' = 1, (for a more systematic definition, see Section A.1 More formally,
Theorem 1:
The class of GMM-RKS Transformers with positional embeddings is Turing complete.
Proof is in the Appendix A We also show that:
Theorem 2:
The class of GMM-PRF Transformers with positional embeddings is Turing complete.
For a detailed proof, see Appendix A
Since the sampling of in Equations and is random, have some stochasticity in the attention layers of our model. We now show that the Mean Square Error (MSE) of the estimation can be reduced by reducing the cigenvalues of the learnt covariance matrix. In particular, we show that:
Let μ and Σ = S be the learnt mean an covariance matrices of the sampling distribution. Further let and be the parameters, and p = k - q and o = k + be their vector difference and sum respectively, and m be the number of samples. The MSE of the linear approximations is given by:
Theorem 3:
M ESEGMM-REE RKS
2
cos (n² p) -1-111spp]] p112)²
(13)
1 2(1a Ia +|||| -1811- (E ooll ellsT
MSFGMM-PRR PRE
e
(14)
m
It is interesting note that in this formulation, MSFGMM-RR RKS is bounded above by 2 while such bound can be established for MS MSEGMM-PRE For the detailed proof see Supplementary Section B
m
Complexity Analysis: While all of our models have linear complexity with respect to the context length L, differences still exist amongst the various methods. Notable, GMM-RKS and GENERATIVE have quadratic time and space complexity in the query size dg. Both the FASTFOOD methods avoid this approximation, whereas GMM-PRF avoids this by the use diagonal covariance matrix. The complexities are listed in Table
Another factor that controls timing the sampling of Sampling too frequently can lead to significant slowdowns whereas sampling too few times can lead to biased learning. For our experiments, we resample every 100 training iterations, although this can be changed. A detailed list of all hyperparameters along with implementation details are provided in Appendix
3 Experiments
3.1 Does kernel learning improve performance of fixed kernel methods on longer sequences?
Long Range Arena (LRA: Tay et al. 20216) a diverse benchmark for the purpose evaluating the ability of sequence models reason under long-context scenarios. includes tasks that vary both in terms of the context length (ranging from 1K to 4K tokens) well as the data modalities (including text and mathematical expressions). We evaluate the KL-TRANSFORMER architectures introduced in Section 2.1 on the Text, Retrieval and ListOps tasks from LRA which deal with sentiment prediction from IMDB Reviews, document similarity classification and pre-order arithmetic calculations respectively. Both Text and Retrieval use character level encodings, bringing their maximum length to 4000 tokens. The ListOps dataset is synthetically generated and has a fixed length of 2000 tokens. For more details on the tasks, we refer the interested reader to the original paper by Tay al. 2022bb
Setup: ensure fair comparison, we closely follow the same data preprocessing, data split, model size and training procedure as in (Tay et al. 2021b). Within each task, common configuration is used across all KL-TRANSFORMER
For the proof of Theorem 3, the secific and 2 however, hold general
2, Ma μu
Fg which used experiments, Since Theorem

---

Published in Transactions on Machine Learning Research (07/2022)
Model
Complexity
ListOps
Retrieval
Avg.
2K
4K
4K
Random Predictor
NA
10.00
50.00
50.00
36.67
Bascline Models
Softmax Trans. VVaswani
O(L
36.38
64.27
57.46
$2.70
Synthesizer TTay ct al.
O
36.50
61.68
54.67
50.95
Sinkhorn TTy al.
((LLB)²)
34.20
61.20
53.83
49.74
Sparse Trans. (Child et al.)
O(LvL)
35.78
63.58
59.59
52.98
Reformer (Kitáev el al.
OLL L)
36.30
56.10
53.40
48.60
Local Attention (Parmar
O(LK)
15.95
52.98
53.39
40.77
Longformer (Beltagy et al.
O(LK)
36.03
62.85
56.89
51.92
Linformer (Wang ct al.
O(L)
35.49
53.49
52.27
52.56
Big Bird (Zaheer ct al.
O(LK)
37.08
64.02
59.29
53.46
LinearElu (Kthhroopoulos et al.
O(L)
17.15
65.90
53.09
45.38
Performer (Choromanski al.)
O(L)
36.00
65.40
53.82
51.74
Kernelized Transformers
GMM-RKS (Eq.)
(())
18.55
63.95
58.64
47.05
FASTFOOD-RKS (Eq.
O(L)
18.65
65.67
61.92
48.75
GENERATIVE-RKS (Eq
O(L)
18.50
66.50
64.76
49.92
GMM-PRF (Eqs.
10
O(L)
36.96
62.64
65.27
54.96
FASTFOOD-PRF (Fqs.
O(L)
37.05
64.66
71.13
57.61
GENERATIVE-PRF (F4s.
((L)
37.42
62.90
69.81
56.71
Table 2: Experimental results on the LRA benchmark. We report accuracy on the test set, except for Text where validation set used. The best model is boldface and the second best is underlined if within 1% the best. Accuracy scores for all baseline models are from Tay et al. (202Ib). Here, refers the sequence length, K refers to the size of local window and B < L is model specific parameter. For our models, accuracy is averaged over 100 runs.
models based on the configuration specified in the LRA code repository We outline the hyperparameters for all tasks in Table the Appendix.
The results across all LRA tasks are summarized in Table 2 KL-TRANSFORMER variants that learn the kerel function directly from the data in an end-to-end manner outperform the baseline models by occupying both best and second-best performances. We find that AL-TRANSFORMMRS based on PRFs tend to outperform their RKS counterparts which is also reflected on the average LRA score, with FASTFOOD -RF being the best pessfoerfonnin model.
Results:
3.2 Trade-off between Accuracy and Efficiency
We benchmark the efficiency of each KL-TRANSFORMER in terms of peak memory usage and training speed and compare it against three baseline models from the LRA benchmark. Specifically, we compare against other efficient Transformer architectures that employ fixed kernel feature maps (e.g. LinearElu and Performer) as well as the Softmax Transformer which is one the strongest baseline models (see Table 2}. We conduct efficiency benchmarks on the two LRA tasks with sequence length equal 4K in order assess the efficiency these methods in modelling tasks that require long context (results for the other two datasets are included in the Appendix). Speed measurements (steps per second) refer wall-clock training time (including overheads). In both cases experiments are conducted on 8 NVIDIA TITAN RTX GPUs. The comparison is illustrated in Figure On the Text task, GENERATIVE- RRS is the best performing model, although it consumes more memory than the remaining -IL.TTRAN ASFORREE architectures (it is still more efficient than the Softmax Transformer). LincarElu consumes the least amount of memory, while GMM-RKS provides a trade-off between the two. In Retrieval the situation is much clearer, with FASTFOOD-PRE and GENERATIVE- outperforming significantly other models terms of accuracy while having very low memory
hht /tpp:///ittiii www.ggtttbb..... cwm/goog 8- www.arrheheec. wwwparrcchhoog--ooe wwwp://onngg------eeeeeee arens

---

Published in Transactions on Machine Learning Research (07/2022)
8
Transferner
Tranfoom
Memory Consumption (GB)
Memory Consumption (GB)
GenertiveRRSS
Genentive RKS
Genentve-PRF
Fastfoed
GMM-PRE PRF
Footood-PRR
Porforen
FIMM-KKS
0..5
Perforncc
RMM PRF
Lincarlla
0.25
62
64 Text Accuracy
66
52
56
60
64
68
Retrieval Actrievyllcccuaayy
Figure 3: We compare the peak memory consumption (y-axis), performance (x-axis) and speed (denoted by of circle) for the various KL- AL-TRANSFORMER architectures the two LRA tasks with sequence length equal to 4K. Memory usage refers to average memory usage across GPUs and speed (steps per second) for each model is reported relative to the speed of Softmax Transformer (larger circles denote faster models). For similar graph on the ListOps Task see Fig in the Supplementary
consumption. The training speed of KL- EL-TRANSFORMEES is of the same order of magnitude as Performers (as indicated by the area of each circle in Figure 3).
Lastly, in Figure 4 we report the peak memory consumption as the sequence length changes from 1K to 4K on the Text dataset. As expected, all our models have a linear increase in memory consump-
GMM-RKS
Fantfoot-RKS RKSS
Memory Consamption (GB)
Generutive-RSS
GMM PRP
tion with increasing sequence length, as opposed the Softmax Transformer which has dramatic increase in memory consumption. Furthermore, Figure 7 in the Appendix reports the memory us- age of each FASTFOOD-PRE and GENERATIVE-PRR are not only our best performing models on average, but they also consume the least memory among various Thus, among the models proposed in this paper, we can recom- mend FASTFOOD-PRF and GENERATIVE-PRE the model that achieves the best accuracy with the least memory consumption.
Gem
8
across all datasets. We find that
4
across all datasets.
2
Sequence Length (c1000)
Figure 4: Memory vs sequence length
3.3 How do KL-TRANSFORMERS perform short sequence tasks?
We compare the SISFORANS and Softmax Transformer in common transfer learning setting. We adopt the setting of BERT-like models (Devlin et al. 2019: Liu et al.) 2019), except that we have fewer layers and heads (see Table 8 for details) and pre-train all models (including Softmax Transformer) on the WikiText-103 dataset (Merity et al.] 2016) using non-contextual WordPiece embeddings (Wu et al.] (016). Pre-trained models are then fine-tuned on the General Language Understanding Evaluation (GLUE) benchmark (Wang et al. 2019). a collection of resources for training and evaluating language understanding systems. All tasks in GLUE consist of either one two sentences input and therefore be casily captured with context of size 512. Since the context length is rather short, the difference between training and inference time across the various models is minimal. Instead, the main goal of this task is to assess how do KL-TRANSFORMER compare against Softmax Transformers on a set of tasks where the later have been established as the de-facto architecture.
The results on all downstream GLUE tasks are shown Table Crucially, we demonstrate that there is no significant loss in performance compared to Softmax Transformers when kernelized variants used on short language modelling tasks. As illustrated in Table 3 KI-TTRANSFORMERS perform par with the Softmax Transformer.
8

---

Published in Transactions on Machine Learning Research (07/2022)
Model
SST2 MRPC (acc/f1)
QQP
MNLII
QNLI
WNLI
RTE (acc)
CoLA (MCor)
(acc)
(aco/f1)
m/mm (acc/acc)
(acc)
(acc)
Softmax Trans.
0.81
0.70/0.82
0.83/0.76
0.64/0.64
0.68
0.56
0.6
0.18
FASTFOOD-RKS
0.83
0.71/0.82
0.81/0.74
0.57/0.57
0.64
0.59
0.56
0.13
GMM-RKS
0.80
0.70/0.82
0.77/0.69
0.47/0.48
0.60
0.61
0.57
0.07
GENERATIVE-RKS
0.81
0.70/0.82
0.81/0.73
0.59/0.58
0.63
0.62
0.58
0.16
FASTFOOD-PRE
0.81
0.71/0.82
0.81/0.74
0.56/0.57
0.64
0.59
0.58
0.12
GENERATIVE-PRE
0.80
0.71/0.82
78800.77
0.56/0.56
0.61
0.60
0.55
0.10
GMM-PRF
0.82
0.71/0.82
0.81/0.74
0.56/0.56
0.64
0.59
0.59
0.21
Table 3: Results on the GLUE benchmark after fine-tuning on respective tasks. KL-TRANSFORRERS continue to
competitive to Transformers even in short context problems.
Task
Model
Max Egy.
Mean Egy.
Task
Model
Max Egy.
Mean Egv.
GMM-RKS
0.486
0.031
GMM-RKS
0.499
0.053
Text
Retrieval
GMM-PRF
0.096
0.025
GMM-PRF
0.186
0.042
Table 4: Distribution of Eigenvalues for GMM-RKS and GMM-PRF models for the Text and Retrieval tasks. For a head by head distribution see Tables
9
TO
and
4 Empirical Analysis
In section 3 we observed that our models compare favourably with other linear models, with Positive Random Features (PRF) based models usually doing better that the ones based on Random Kitchen Sinks (RKS) (although RKS does better on the Text task). In this section, we want to compare these linear kernels in terms of their empirical variance and how well they deal with sparsity. In Theorem 3, we already noted that the MSE for both GMM-RKS and GMM-PRF models decreases with decrease eigenvalues of the covariance matrices, making it useful learn such a matrix. We also noted that the variance in the output of GMM-RKS is bounded, and thus it should not face issues with sparse datasets when approximating the softmax kernel as show in Choromanski al. (2021b)). We test for these results empirically in subsections 4.1 and 4.2 respectively. We then conclude in 4.3 with discussion on which of our proposed models is best to use in a given scenario.
4.1 Comparison of variance
Looking the eigenvalues of the covariance matrices of final trained models in Table we find that GMM-RKS and GMM-PRF models trained the Text and Retrieval tasks indeed learn to have eigenvalues significantly smaller than 1 (which would be the case for non-learnable co-variances).
While lower eigenvalues reduce the variance in our models significantly, it does not get completely nullified. To understand how much stochastic remains in our models, and also to compare RKS with PRF in regards to their variance, we look the final output&f] the models produce when repeatedly run on the same example. We record the output produced by the model for 100 runs on each datapoint, and calculate the following metrics to quantify the variance:
1. Relative Standard Deviation (RSD): RSD is a standardised measure of dispersion (see Currie & Svehla (1994), def 3.9), defined as the ratio of the standard deviation of a set of observations to the absolute value of their mean. The ratio ensures that the measure is invariant to scaling (eg. multiplying the penultimate layer of the model by a constant factor does not affect RSD) while the absolute of the mean ensures that opposing classes don't cancel each other out.
"We only report cigenvalues for the GMM models since covariance matrices not well defined the Generator and FastFood. Also, the ListOps task left of the entire analysis because its multi-class nature makes the analysis complex Whiile the theory makes claim about the output each attention head, evaluating every head at every layer would give us large number of values to analyse. The output considered before aplying the final sigmnid

---

Published in Transactions on Machine Learning Research (07/2022)
2. Prediction Inconsistency (PI): While RSD quantifies the variance in the model's continuous output, in practice, only the sign of the output is of interest as that decides the predicted class. As way to quantify stochasticity in the discrete output, we count the number of times the output does not have the majority label, i.e., if the output is positive æ times, then we have PI = min (x, 100 æ). Alternately, it can seen as the failure rate if we treat the majority prediction as the true class.
3. Accuracy Gain with Voting (AGV): Since the final output has propensity to change its sign, we can get a more robust prediction at inference time by running the model multiple times and considering the class it predicts more times as its output, instead of taking the prediction from a single inference step. Doing so, we are likely to get closer to the mean prediction (by the Law of Large Numbers, see Révész et al. (2014)), and we get accuracy scores which are slightly higher than those reported Table and we call this value the Voting Accuracy (VA). AGV is then defined as the ratio between the voting accuracy and the regular accuracy.
Text
Generative-RLe RKKS
FastFood-RKS
Retrieval
GMM-PRF
Generative PRF
FastFood
10-6
10-4
10
10
10
60
70
100
101
1.02
RSD
AGV
Figure 5: Values of variance metrics Relative Standard Deviation (RSD), Prediction Inconsistence(PI), Average Gain with Voting (AGV) well Voting Accuracy (VA) on the Text and Retrieval tasks.
The above metrics, in addition the Voring Accuracy, are calculated over all instances in the validation and test set for Text and Retrieval tasks respectively, and plotted in Fig. We note that our metrics are all highly correlated with each other (Rt values 0.95, 0.98 and 0.89 for RSD-AGV, RSD-PI and AGV-PI respectively, see supplementary Fig
10
We notice that our RKS based models, in general, have lower variance compared to our PRF models. We also note that Generator and FastFood are able to further reduce the variance the output, possibly due to additional learnable parameters in Generator and fewer sources randomness FastFood. We also notice that all models have greater variance in the Text task, and is possible that this is related to the better performance of RKS based models in this task, and their inherent lower variance helps them outperform their PRF counterparts.
4.2 Effectiveness on Sparse Datasets
Choromanski et al. (2021)) demonstrated that Random Kitchen Sinks are used to estimate the softmax-kernel, the MSE tends to oo as the raw softmax scores tend to 0. This makes particularly hard to deal with sparse datasets where multiple positions need near-zero attention.
Thus, in order to test how sparsity in the dataset affects models, design synthetic experiment where the model is fed sequence of ordered pairs, where the first element, which we call score, takes a value of -1 1 or 1, and the second element, which we call relevance, takes a value of 0 or 1. Looking at the sequence, the model must predict the sum of the scores at all positions with relevance Note that any position with rele-
0.8
Accuracy
0.6
0.4
GMM-RKS 0.1
GMM-PRF 0.1
vance 0 does not contribute the final answer, and therefore, does not need to be attended to. We construct the inputs by sam- pling the relevance of each position from Bernoulli(p). Thus, the sparsity can be controlled by changing (for detailed description of the task, see Supplementary).
0.2
GMM-RKS 0.5
GMM-PRF 0.5
GMM-RKS 0.9
GMM-PRF 0.9
26
51
76
101
126
151
176
201
226
251
276
301
326
351
376
Iterations
Figure the learning curves for various sparsity levels our synthetic experiment. Due to the simple nature of the task, all models eventually reach 1 100% accuracy. We observe that
Figure 6: Learning curves for our synthetic experi-
ment. The number after the model name is inversely proportional to sparsity in the dataset
10

---

Published in Transactions on Machine Learning Research (07/2022)
the convergence time of GMM-RKS remains more less unchanged with increasing sparsity, while for GMM-PRF the model converges slower as sparsity decreases. We believe that the slower convergence of GMM-PRF is correlated to the variance in its output, which leads to variance in gradients. Observing the gradients propagated to the classifier layer (See Table 14 in Supplementary), we find that this is indeed the case, and not only is the variance in the gradients higher for GMM-PRF, but also the mean is higher, making it take bigger steps in an uncertain direction. This result provides us an insight into which models to use under which circumstances.
4.3 Which Model to Use?
PRF outperforms RKS is terms of accuracy in most of our experiments, especially if we are able to run multiple inference steps. Therefore, in general, we recommend using PRF over RKS. However, if consistency of predictions is important, or rapid training is required on a not-so-sparse task, one may consider RKS as well.
Finally, we observe that Generator and Fastfood methods always outperform the vanilla linearisations. Between them, Fastfood can be significantly faster if properly implemented on GPU. However, generators may perform better if large amount of data is available since they provide greater amount of flexibility.
5 Related Work
5.1 Efficient Transformers
A wide variety of approaches belong to the class of efficient Transformers TTa et al. 2020b). We survey them below:
Sparse models: Memory Compressed Transformer (Liu et al. 2018) uses a convolution kernel (of size K') to sub- sample keys and values reducing the complexity to O K -) Inspired by the notion of sparsity, Child et al. (2019) introduced sparse factorizations of the attention matrix to reduce the overall complexity to O T.). Subsequently, Roy et al. (2020) proposed the Routing Transformer which employs K-means clustering to learn dynamic sparse attention regions, achieving an overall complexity of OLV Recently, Secently.SSummae et al. (2022) proposed yet another sparse (LL approach that learns which bucket each query/key is to be placed in based on final attention values. Sparsity has also been achieved by efficiently subsampling queries and/or keys (Chen et al.) 2021). Kitaev et al. (2020) proposed the Reformer, which reduces complexity to O(L L) by using locality-sensitive hashing to group together similar
symbols. Ye et al. (2019) also proposed OL log L) algorithm using binary partitions of data. There also exist other
works that mainly focus on memory reduction (Liu et al. 20188] Tay et al.) [2020a] Gupta et al.] 2021). While these methods are faster than Softmax Transformers, their asymptotic time complexity remains quadratic. Further approaches
attempt to minimize the constants involved in the quadratic attention, but keep the same assymptotic complexity (Dutta) ct al. | Chen et al. 2022).
Local or Global Attention: Parmar et al.
(2088) was one of the first local attention models that achieved O(L) complexity in both time and space, by using local attention over a constant length-context. similar local attention based method was also utilized by Zhang al. (02T and Liu et al. (202). , the latter of which added special tokens at the start of each local attention block, which attend globally. Another set of methods proposed to approximate the global attention by replacing the softmax to allow changing the order of matrix multiplication, making the calculation of the QKV product linear in the context length (Yorsh et al. 2021 Hua et al. [022]] Qin et al. 2022). Linearised global attention is also used by Performer (Choromanski et al. 2021b), which has been built upon by other works. Chen (202T) and Luo et al. (202T) both try to improve the performance of Performer by incorporating relative positional embeddings which they show makes it strictly stronger in terms of representability, [Choromanski et al. (2011)) tries to combine RKS and PRF attentions by learnable weight parameter to get the best of both worlds. Both these methods can be applied is to our work. [Schlag al. (2021a) a) take a different view on this, claiming that having stochasticity hampers performance, and propose deterministic feature map to avoid this.
Multiple Types of Attention: Beltagy et al. (020 proposed ((L) method that combines the above two approaches by using local sliding windows as well as global attention components. Zaheerr et al. (2020) proposed Big Bird, another sparse attention mechanism which combines global attention, random attention and local attention to reduce the
complexity to (L). A similar construction was previously used by /Ainslie et al. (2020). More recently the combination of global and local attention has been utilized by Zhu al. (202T) and Nguyen et al. (202T). Further, Xiong et al. (202T)
adapted the Nyström method approximate standard self-attention with O(I) complexity. Lu et al. (2021) also makes
11

---

Published in Transactions on Machine Learning Research (07/2022)
use of a similar near-field and far-field attention mechanism, where the far field attention is calculated via low rank approximation using lonh and ReTU non-linearities. A different line of work has attempted to limit the number of key-value pairs that are to be attended to, by attempting to summarise the variable length context with fixed number of memory-cells (Zhang & Cai 2022 Ma et al. (021). While some these more nuanced approaches outperform most their innovations are orthogonal to ours, and none of these approaches explore kernel learning within the attention mechanism of Transformers.
5.2 Kernel Learning
While kernel methods have long been used to solve non-linear statistical problems, they traditionally scaled poorly with the number of data points thereby limiting their applicability on large datasets (Vapnik et al. 1997; Cortes &
Vapnik 1995 Schölkopf et al. 1998; Schölkopf & Smola 2001 Hofmann et al. 2008). Prior to RKS, several kernel approximation techniques have been proposed to improve the scalability of kernel methods, including greedy basis selection techniques (Smola & Schökopf 2000), hivide-and-conquer approaches (Hsieh et al., (214)) /Zhang al.) 2013
Liu al. ). non-stationary spectral kernels (Remes et al. 2017). generalized spectral kernels (Samo Roberts
2015) well as Nyström methods (Williams & Seeger 2001
The method of Random Kitchen Sinks has been revisited several times, either to improve the approximation quality Yu
al. 60 Choromanski et al. 2017 Li 2017 Avron et al. [0166] Lyu 2017). reduce the time and memory complexity
[Le et al. 2013; Choromanski & Sindhwani 2016 Feng et al. 2015 Dao et al. 2017) or analyze theoretically the risk
and generalization properties of the algorithm (Sutherland & Schneider 2015 Sun et al. 2018 Li ct al. 2019b). A
systematic survey of random feature methods for approximating kernel functions can be found in (Liu et al. 2021
Lastly, there exist class of methods that extend the RKS framework to enable kernel learning. Representative approaches involve either one-stage (Yang et al.| 2015 Fang et al.) 2020) two-stage procedure (Sinha & Duchi 2016 Li et al. 2019a; Bullins et al. [2118] Shen et al. 2019). Two-stage approaches involve an intermediate step in which the problem of sernel-alignmen is solved (Cristianini et al. 2002). However, solving the semeel-algnmmnt problem requires accessed to labeled data which is not available in this case, as inputs to the kernel learning algorithm are the intermediate representations of the input sequence.
6 Conclusion
In this paper, we bridged the gap between advances in kernel learning and efficient Transformers by proposing several kernel learning methods for Transformers that increase the expressiveness of Transformers while keeping the computational complexity linear in sequence length. We showed that our proposed KL-TRANSFORMER are Turing-complete and can control their variance. Experimentally our proposed models perform on par with, and possibly exceed the performance of existing efficient transformer architectures on long context tasks without falling behind on short context tasks. We also found that for some datasets such as LislOps, RKS based models tend to fall short of their PRF counterparts. Our experiments further demonstrate that the memory consumption of our models scales linearly with the sequence length.
Ethical Considerations
Our work is on making Transformers computationally efficient without losing expressiveness. Our models were evaluated on publicly available benchmark datasets. The datasets used in our work do not contain sensitive information to the best of our knowledge.
Reproducibility: We plan to open source the entire code of the KL-TRANSFORRER framework (including the implementation of all models as well as the code for replicating all of our experiments) before the camera ready version of the paper. As part of this submission, we include code for all the methods proposed by us along with instructions on how to reproduce results. detailed description of all hyperparameters (for both LRA as well as GLUE benchmarks) has been included in Appendix C Finally, regarding our theoretical contributions, we present detailed theoretical analysis in Appendix A
12

---

Published in Transactions on Machine Learning Research (07/2022)
References
Milton Abramowitz and Irene A Stegun. Handbook of mathematical functions with formulas, graphs, and mathematical tables. national bureau of standards applied mathematics series 55, tenth printing. 1972.
Joshua Ainslie, Santiago Ontañón, Chris Alberti, Philip Pham, Anirudh Ravula, and Sumit Sanghai. ETC: encoding long and structured data in transformers. CoRR, abs/2004.08483, 2020.
Haim Avron, Vikas Sindhwani, Jiyan Yang, and Michael W. Mahoney. Quasi- carlo feature maps for shift-invarian kemels.. Journal of Machine Learning Research, 17(120):1 2016.
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer, 2020.
Tom Brown, Benjamin Mann, Nick Ryder, Mclanic Subbiah, Jared 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, Benjamin 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.
Brian Bullins, Cyril Zhang, and Yi Zhang. Not-so-random features. In International Conference on Learning Representations, 2018.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers, 2020.
Peng Chen. Permuteformer: Efficient relative position encoding for long sequences, 2021. URL htttt::////ax
brggabs//2199.02 02377
Yifan Chen, Qi Zeng, Dilek Hakkani-Tur, Di Jin, Heng Ji, and Yun Yang. Sketching as a tool for understanding and accelerating self-attention for long sequences, 2021.
Zhaodong Chen, Yuying Quan, Zheng Qu, Liu Liu, Yufei Ding, and Yuan Xie. Dynamic nam fine-grained structured sparse attention mechanism, 2022. URL htttss://aaxiiii http:////wxw..gg/aa///////// wotg/bos/////e///eoo /203. 00091
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers, 2019.
Krzysztof Choromanski and Vikas Sindhwani. Recycling randomness with structure for sublinear time kernel expansions. In Proceedings of the 33rd International Conference on International Conference on Machine Learning Volume 48, ICML'16, pp. 2510. JMLR. 2016.
Krzysztof Choromanski, Haoxian Chen, Han Lin, Yuanzhe Ma, Arijit Schanobish, Deepali Jain, Michael S. Ryoo, Jake Varley, Andy Zeng, Valerii Likhosherstov, Dmitry Kalashnikov, Vikas Sindhwani, and Adrian Weller. Hybrid random features. CoRR, abs/21 10.04367, 2021a. URL
04367
Krzysztof M Choromanski. Mark Rowland, and Adrian Weller. The unreasonable effectiveness of structured random orthogonal embeddings. In Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.). Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quiney Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J Colwell, and Adrian Weller. Rethinking attention with performers. In International Conference on Learning Representations, 2021b.
Corinna Cortes and Vladimir Vapnik. Support-vector networks. Mach. Learn., 20(3):273-29) September 1995. ISSN 0885-6125. doi: 10.1023/A:1022222441144
Nello Cristianini, John Shawe- Taylor, André Elisseeff, and Jaz Kandola. On kernel-target alignment. In Dietterich, S. Becker, and Z. Ghahramani (eds.), Advances Neural Information Processing Systems, volume 14. MIT Press, 2002.
13

---

Published in Transactions on Machine Learning Research (07/2022)
L. A. Currie and G. Svehla. Nomenclature for the presentation of results of chemical analysis (iupac recommendations 1994). Pure and Applied Chemistry, 66(3):695 608, 1994. doi:
URI. https
Tri Dao, Christopher M De Sa, and Christopher Ré. Gaussian quadrature for kernel features. In Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Jacob Devlin, Ming-Woi Wei Chang, Kenton Lee, and Kristina Toutanova. BERT:
Pre- Pre-training of deep bidirectional transformers for language understanding. Proceedings of the 2019 Conference of the North American Chapter the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-441111 1171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguisties. doi: 10.18653/v L/N19 1423.
Subhabrata Dutta, Tanya Gautam, Soumen Chakrabarti, and Tanmoy Chakraborty. Redesigning the transformer architecture with insights from multi-particle dynamical systems. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 5531-5544, Curran Associates, Inc., 2021. URL
FLlc/22d388/3122331233 www.:afffcee0000000000000 Paper .pf
Kun Fang. Xiaolin Huang, Fanghui Liu, and Jie Yang. End-to-end kernel learning via generative random fourier features, 2020.
Chang Feng, Qinghua Hu, and Shizhong Liao. Random feature mapping with signed circulant matrix projection. In Proceedings of the 24th International Conference on Artificial Intelligence, IJCAI'15 pp. 3490-3496. AAAI Press, 2015. ISBN 97815 577357377.
lan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K. Q. Weinberger (eds.), Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014.
Ankit Gupta, Guy Dar, Shaya Goodman, David Ciprut, and Jonathan Berant. Memory-efficient transformers via top-k attention, 2021. URL htttps://arxiv..om abs / 2106. 06899
Thomas Hofmann, Bernhard Schölkopf, and Alexander J. Smola. Kerel methods in machine learning, The Annals of Statistics, 36(3): 171 - 1220, 2008. doi: 10. 10.1244/000334400000007777
Cho-Jui Hsieh, Si Si, and Inderjit S. Dhillon. A divide-and-conquer solver for kernel support vector machines. In Proceedings of the 31st International Conference International Conference Machine Learning Volume 32, ICML'14, pp. 1-566-1----111 1-574. JMLR.org. 2014.
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoe V. Le. Transformer quality in linear time, 2022. URL https: //axxiv.org/abs//// 10447
John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph-based protein design. In H. Wallach, H. Larochelle, A. Beygelzimer, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32, pp. 15820-15831. Curran Associates, Inc., 2019.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers RNNs: Fast toregressive transformers with linear attention. In Hal Daumé IIII and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 1 119 of Proceedings of Machine Learning Research, pp. 5156-5165. PMLR, 13-18 Jul 2020.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes, 2014.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
14

---

Published in Transactions on Machine Learning Research (07/2022)
Quoc Le, Tamas Sarlos, and Alexander Smola. Fastfood computing hilbert space expansions in loglinear time. In Sanjoy Dasgupta and David McAllester (eds.), Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 244-252, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR.
Chun-Liang Li, Wei-Cheng Chang, Youssef Mroueh, Yiming Yang, and Barnabas Poczos. Implicit kernel learning. In Kamalika Chaudhuri and Masashi Sugiyama (eds.), Proceedings of Machine Learning Research, volume 89 of Proceedings of Machine Learning Research, pp. 2007-2016. PMLR, 16-18 Apr 2019a.
Ping Li. Linearized gmm kernels and normalized random fourier features. In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining. KDD '17, pp. 315-324, New York, NY, USA, 2017. Association for Computing Machinery. ISBN 9781450348874, doi: 011444339998830998800
7hu Li, Francorss Ton, Dino Oglic, and Dino Sejdinovie. Towards unified analysis of random Fourier features. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 3905-3914, PMLR, 09-15 Jun 2019b.
Fanghui Liu, Xiaolin Huang, Chen Gong, Jie Yang. and Li Li. Learning data-adaptive non-parametric kernels. Journal of Machine Learning Research, 21(208):1-39, 2020.
Fanghui Liu, Xiaolin Huang, Yudong Chen, and Johan A. K. Suykens. Random features for kernel approximation: A survey on algorithms, theory, and beyond, 2021.
Peter J. Liu, Mohammad Saleh, Etienne Pot, Ben Goodrich, Ryan Sepassi, Lukasz Kaiser, and Noam Shazeer. Generating wikipedia by summarizing long sequences. In International Conference on Learning Representations, 2018.
Yang Liu, Jiaxiang Liu, Li Chen, Yuxiang Lu, Shikun Feng, Zhida Feng, Yu Sun, Hao Tian, Hua Wu, and Haifeng Wang. Ernic-sparse: Learning hierarchical efficient transformer through regularized self-attention, 2022. URL
https::/ httpa/aaww.oogrbba///aaaa 12276
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, 2019.
Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, and Li Zhang. Soft: Softmax-free transformer with linear complexity, 2021.
Jiasen Lu, Dhruv Batra, Devi Parikh, and Stefan Lee. Vilbert: Pretraining task-agnostic visiolinguistic representations for aasion-and-language tasks. In H. Wallach, H. Larochelle, A. Beygelzimer, E. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32, pp. 13-23. Curran Associates, Inc., 2019.
Shengjie Luo, Shanda Li, Tianle Cai, Di He, Dinglan Peng. Shuxin Zheng, Guolin Ke, Liwei Wang, and Tie- Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding, 2021. URL https //
12566
Yueming Lyu. Spherical structured feature maps for kernel approximation. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 2266-2264. PMLR, 06-11 Aug 2017.
Xuezhe Ma. Xiang Kong, Sinong Wang, Chunting Zhou, Jonathan May, Hao Ma, and Luke Zettlemoyer. Luna: Linear unified nested attention. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34. pp. 2441- 2453. Curran Associates, Inc., 2021. URL htttps://ppp.. htt.:p:://proceeedngg. www.senurippi pwtppaper/222///1eee (tt.///00000044baaaaa Paa---apeeeeee pdf
1431
Ali Madani, Bryan McCann, Nikhil Naik, Nitish Shirish Keskar, Namrata Anand, Raphael R. Eguchi, Po-Ssu Huang, and Richard Socher. Progen: Language modeling for protein generation, 2020.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016.
15

---

Published in Transactions on Machine Learning Research (07/2022)
Tan Nguyen, Vai Suliafu, Stanley Osher, Long Chen, and Bao Wang. Fmmformer: Efficient and flexible trans- former via decomposed near-field and far-field attention. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 29449- 29463. Curran Associates, Inc., 2021. URL http:ttsp//////ee.nnnn////////ee neurips F621 05215566f2224555577733333333 ptap:p/ppppppppppppaaaeerrr.ppp
Paper.
Junier B. Oliva, Avinava Dubey, Andrew G. Wilson, Barnabas Poczos, Jeff Schneider, and Eric P. Xing. Bayesian nonparametric kermel-lcarning In Arthur Gretton and Christian C. Robert (eds.), Proceedings of 19th International Conference on Artificial Intelligence and Statistics, volume of Proceedings of Machine Learning Research, pp. 1078-1086, Cadiz, Spain, 09-11 May 2016. PMLR.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz. Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 4055-4064, Stockholmsmässan, Stockholm Sweden, 10-15 Jul 2018, PMLR.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vito,, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Buc, F. Fox, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention. CoRR, abs/2103.02143, 2021.
Jorge Pérez, Javier Marinkovic, and Pablo Barceló. On the turing completeness of modern neural network architectures. CoRR, abs/1901.03429, 2019.
Jorge Pérez, Javier Marinković, and Pablo Barceló. On the turing completeness of modern neural network architectures. In International Conference on Learning Representations, 2019.
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention, 2022.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language understanding by generative pre-training, 2018
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yangi 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(140):1-67, 2020.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In Proceedings of the 20th International Conference on Neural Information Processing Systems, NIPS 07, pp. 1177-1184, Red Hook, NY, USA, 2007. Curran Associates Inc. ISBN 9781605603520.
Sami Remes, Markus Heinonen, and Samuel Kaski. Non-stationary spectral kernels. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
P. Révész, Z. W. Birnbaum, and E. Lukacs. The Laws of Large Numbers. Probability and mathematical statisiics.. Elsevier Science, 2014. ISBN 9781 483269023. UtLpttps: htttttps://booo.... http://boooooo.ggll.oo google ch.books (tt.:oooko....oooo.occc.........
Eitan Richardson and Yair Weiss. On gans and gmms. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa- Bianchi, and Garnett (eds.), Advances Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
Alexander Rives, Joshua Meier, Tom Sercu, Siddharth Goyal, Zeming Lin, Demi Guo, Myle Ott, C. Lawrence Zitnick, Jerry Ma, and Rob Fergus. Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences. bioRxiv, 2020. doi: 10.1101/622803.
16

---

Published in Transactions on Machine Learning Research (07/2022)
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-basee sparse attention with routing transformers, 2020.
Walter Rudin. Fourier Analysis on Groups. John Wiley & Sons, Ltd, 1990.
Lars Ruthotto and Eldad Haber. An introduction to deep generative modeling, 2021.
Yves-Laurent Kom Samo and Stephen Roberts. Generalized spectral kernels, 2015.
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 2021a. URL
ngs..... w www.s/ressssvv... schh aaahllggla.ht..
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 on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 9355-9366. PMLR, 18-24 Jul 2021b.
Bernhard Schölkopf and Alexander J. Smola. Learning with Kernels: Support Vector Machines, Regularization, Optimization, and Beyond. MIT Press, Cambridge, MA, USA, 2001. ISBN 0262194759.
Berharrd Schölkopf, Alexander Smola, and Klaus-Roher Müller. Nonlinear component analysis kernel eigenvalue problem. Neural Computation, (K5):1299-133 1998. doi: 10.1162/089966698300144444
Zheyang Shen, Markus Heinonen, and Samuel Kaski. Harmonizable mixture kernels with variational fourier features. In Kamalika Chaudhuri and Masashi Sugiyama (eds.), Proceedings of the Twenty-Second International Conference Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pp. 3273-3282. PMLR, 16-18 Apr 2019.
Bernard W Silverman. Density estimation for statistics and data analysis, volume 26. CRC press, 1986.
Aman Sinha and John C Duchi. Learning kernels with random features. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
Alex J. Smola and Bemhard Schökopf. Sparse greedy matrix approximation for machine learning. In Proceedings of the Seventeenth International Conference on Machine Learning, ICML '00, pp. 911-918, San Francisco, CA, USA, 2000. Morgan Kaufmann Publishers Inc. ISBN 1538607072.
Yitong Sun, Anna Gilbert, and Ambuj Tewari. But how does it work in theory? linear svm with random features. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Desa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018
Zhiqing Sun, Yiming Yang, and Shinjae Yoo. Sparse attention with learning to hash. In International Conference on Learning Representations, 2022. URL http:st///peeeeeeeieeeee http://wwwwweeeffoooooeoe wtwp://ooreiddeeeeeeeee w
Dougal J. Sutherland and Jeff Schneider. On the enror of random fourier features. In Proceedings of the Thirty-Firy Conference on Uncertainty in Artificial Intelligence, UAI'1,, pp. 862-871, Arlington, Virginia, USA, 2015, AUAI Press. ISBN 9780996643108.
Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da- Da-Cheng Juan. Sparse sinkhorn attention. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 9438-9447. PMLR, 13-18 Jul 2020a.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, 2020b.
Yi Tay, Dara Bahri, Donald Metzler, Da-Cheng Juan, Zhe Zhao, and Che Zheng. Synthesizer: Rethinking self-attention in transformer models, 2021a.
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 Learning Representations, 2021b.
17

---

Published in Transactions on Machine Learning Research (07/2022)
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philipge Morency, and Ruslan Salakhutdinov. Transformer dissection: An unified understanding for transformer's s attention via the lens of kernel. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP),, pp. 3344-4353, Hong Kong, China, November 2019. Association for Computational Linguisties. doi: 10. 1/111//////D1-1 1443.
Vladimir Vapnik, Steven Golowich, and Alex Smola. Support vector method for function approximation, regression estimation and signal processing. In M. C. Mozer, M. Jordan, and T. Petsche (eds.), Advances in Neural Information Processing Systems, volume 9. MIT Press, 1997.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, F. ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Guyon, U. V. Luxburg. S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30, pp. Associates, Inc., 2017.
Curran
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill. Omer Levy, and Samuel R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding. In International Conference on Learning Representations, 2019.
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020.
Christopher Williams and Matthias Seeger. Using the nyström method to speed up kemne machines. In T. Leen, T. Dietterich, and V. Tresp (eds.), Advances in Neural Information Processing Systems, volume 13. MIT Press, 2001.
Andrew Wilson and Ryan Adams. Gaussian process kernels for pattern discovery and extrapolation. In Sanjoy Dasgupta and David McAllester (eds.), Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 1067-1075, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR.
Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V. Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, Jeff Klingner, Apurva Shah, Melvin Johnson, Xiaobing Liu, Lukasz Kaiser, Stephan Gouws, Yoshikiyo Kato, Taku Kudo, Hideto Kazawa, Keith Stevens, George Kurian, Nishant Patil, Wei Wang, Cliff Young, Jason Smith, Jason Riesa, Alex Rudnick, Oriol Vinyals, Greg Corrado, Macduff Hughes, and Jeffrey Dean. Google's 's neural machine translation system: Bridging the gap between human and machine translation, 2016.
Yunyang Xiong, Zhanpeng Zeng, Rudrasis Chakraborty, Mingxing Tan, Glenn Fung, Yin Li, and Vikas Singh. Nyströmformer: A nyström-based algorithm for approximating self-attention. Proceedings of the AAAI Conference on Artificial Intelligence, 14138-14418-144848 May 2021.
Zichao Yang, Andrew Wilson, Alex Smola, and Le Song. A la Carte - Learning Fast Kernels. In Guy Lebanon and S. N. Vishwanathan (eds.), Proceedings of the Eighteenth International Conference Artificial Intelligence and Statistics, volume 38 of Proceedings of Machine Learning Research, pp. 1098-1106, San Diego, California, USA, 09-12 May 2015. PMLR.
Zihao Ye, Qipeng Guo, Quan Gan, Xipeng Qiu, and Zheng Zhang. Bp-transformer: Modelling long-range context via binary partitioning, 2019.
Uladzislau Yorsh, Pavel Kordík, and Alexander Kovalenko. Simpletron: Eliminating softmax from attention computa- tion, 2021. hwR.:htttss//aarri.oorrr///////p
15588
Felix Xinnan Yu, Ananda Theertha Suresh, Krzysztof M Choromanski, Daniel N Holtmann-Rice and Sanjiv Kumar. Orthogonal random features. In D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett (eds.) Advances Neu- ral Information Processing Systems, volume 29, Curran Associates, Inc., 2016a. URL httl:tttts///ooeeiiiiiii heurips. http:ipps..cpppppe//1e1ee/ http:///cppppeee.eoeee//aaaaaaaa
Felix Xinnan Yu, Ananda Theertha Suresh, Krzysztof M Choromanski, Daniel N Holtmann- and Sanjiv Kumar. Orthogonal random features. In D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016b.
18

---

Published in Transactions on Machine Learning Research (07/2022)
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: Transformers 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-1722 17297. Curran Associates, Inc., 2020.
Hang Zhang, Yeyun Gong, Yelong Shen, Weisheng Li, Jiancheng Lv, Nan Duan, and Weizhu Chen. Poolingformer: Long document modeling with pooling 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. 12437-12446. PMLR, 18-22 Jul 2021. URL UtL.tt.ps://prroeeeiiii mlr, wttpppeeeee/e/eeeeeeeee zha.gg222
html
Yizhe Zhang and Deng Cai. Linearizing transformer with key-value memory, 2022.
Yuchen Zhang, John Duchi, and Martin Wainwright. Divide and conquer kernel ridge regression. In Shai Shalev- Shwartz and Ingo Steinwart (eds.), Proceedings of the 26th Annual Conference on Learning Theory, volume 30 of Proceedings of Machine Learning Research, pp. 592 617, Princeton, NJ, USA, 12-14 Jun 2013. PMLR.
Chen Zhu, Wei Ping, Chaowei Xiao, Mohammad Shoeybi, Tom Goldstein, Anima Anandkumar, and Bryan Catanzaro. Long-short transformer: Efficient transformers for language and vision. In M. Ranzato, A. Beygelzimer, Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 17736. Curran Associates, Inc., 2021. URL File/9tm
Pwp.saapperppe pdf
19

---

Published in Transactions on Machine Learning Research (07/2022)
Learning The Transformer Kernel - Appendix
A Detailed proof of Theorems
A.1 Definitions
Transformer:A transformer consists of an Encoder and Decoder, which in turn consist of several encoder and decoder layers respectively. A single encoder layer consists of an attention layer(Att) and 2 layer feed-forward neural network(O):
Att x,, W&X, W,X) +x, = (a,)) +
(15)
(16)
In our case, the feed-forward neural network uses perceptron activations(ie fpere(#)) = iff x: >> O and 0 otherwise)
and the attention is (discussed in detail later). The final layer of the encoder is followed by couple of two layer output neural networks, which produce the Encoder Key(Ke) and Encoder Value(Ve) to be used by the decoder. In our proof, we assume these have ReI.U activation neLu (ar) = max((, x))
The decoder layers are similar to the encoder layer except for an additional cross attention layer which attends to the encoder output:
p; Au,,,,; W,Y, W, Y) +y: Au(W", Ke V.) +y:
(17)
(18)
2: O(a;) a;;
(19)
Unlike the encoder, the decoder self attention if Eq. 17 can only attend previous position. After the final layer we have two layer feed-forward neural network with ReLU activation to produce the output. The decoder is initialised with a special seed vector and is repeatedly applied with the right shifted output of the last application as the input of the current application, until some termination condition is fulfilled.
Both the encoder and decoder can further use position embeddings, which have the same dimension as the output of each layer, and are added to the input prior to the first layer. These help in establishing the order of the input
Since the output of any unit of a layer is independent of values to its right, these do not change with time and can be cached. The output of the final layer of the rightmost cell can therefore be regarded as the model state encoding(v)
Turing Machine:A Turing Machine is an abstract construct which consists of a right infinite tape and a read-write head. Each cell of the tape can hold one of many symbols from a predefined alphabet Σ which includes a special blank symbol Additionally, the writ-write head can be in one of many possible states within the state-space Q which includes a special initial state ginit and subset of final states F.
Initially, the tape contains the input followed by infinite number of blank symbols, while the head starts off in the last non-blank cell. In each step, the head executes in accordance with transition function ((s()) ())) (u(1,, (i++)), ()) ie, based on the symbol currently under the head and the current state, it decides the symbol it wants to overwrite the current symbol with, the state it will be in the next step and the direction it wants to move, which can be either left(-1) or right(1). 1). We assume that the transition function already makes sure that the head never moves left from the leftmost cell.
For the purpose of our proof, we additionally define (0) as the index of the cell to which the head currently points, ((),, which represents the step number when the head last pointed to the current cell, ie {(t) max(j]e(d) clA). In the special case where the current cell is being visited for the first time, we have {(i) - 1
A.2
The Proof
In this section, we provide general proof which works for both Theorem and Theorem in the paper. This is possible since the construction only makes of the dual of the kernel functions used, i.e. the gaussian. The fact that
20

---

Published in Transactions on Machine Learning Research (07/2022)
both kernel functions map the gaussian is shown in lemma S.2 (Sec...6)
Our proof is based on the similar proof in Pérez et al. 2019). Any symbols not explicitly defined have same meanings from that paper. We begin the proof by defining our model encoding(v):
V q S1, 21, 22
92 52, 43, 24, 25, 26.
(20)
37, 84: 388
Ig, 210, 111, 2112
where q, E QI@), 8, Q!7,, and 2, Q. giving a total model size of 2/Q) 3|Σ] + 14. Hereafter, [æ]] represents the one-hot encoding for the state symbol depending the position it is being used in. Oq represents all 0's in state field, and represents the Icopy state discussed later, while 0g represents all 0's in symbol field, and represents the blank symbol. Further, ( min((, n) where is the size of the encoder and ( represents the symbol at position g in the encoder. We assume that atleast the last cell of the encoder contains blank symbol.
This differs from Pérez et al. (2019) in the addition of a fourth scalar in the first group, in which we intend to store the current position ( of the head.
Our invariant is that y, the output from the decoder at timestep stores:
1. The current state of the Turing Machine(q())
2. The symbol under the (())
3. The direction of movement of the head in the previous timestep (m
4. The current position of the (())
In all, we get y, (1 [[ 1. (+-1) ........0
Positional Embeddings: The last group (Ig, I10, I11, 212) is dedicated to the positional embeddings, which are given as(I, To (t) These same embeddings are added on both the Encoder and Decoder side.
Encoder: The encoder consists of single layer. It gets as input the symbol at position and the positional embeddings, ie input; [0g. Oz, O, O, 0g. [1 0, O, O, 0, i, O,, 1, i, # which has trivial attention layer(ie, one that outputs all zeroes) and feed forward layer which separates the positional embeddings from the symbols, giving kç [0. ... 0, 1,, 0, 0] and vf 10 O., 0, 0, O [s O, O, 0, 0, i, O,, 0,0,0,0,0).
Decoder Layer 1: The first layer of the decoder calculates the next state, the symbol to be written and the direction movement of the head. This includes 2 cases:
1. Initially, the Decoder starts off with in the state 9copy- While the state is still Icopy- the head writes the symbol at the tth position in the encoder and moves right, until blank symbol is seen. Once a blank symbol is reached, the tape rewrites the blank symbol, moves left and the state changes to ginit-
2. Once we move into qinst. the output is fully defined by the current state and symbol under the head.
To facilitate the first case, make use of the cross attention layer, to get
Att(q. K2 V*) jO,...,.. O,
0g. [a (1) 0,0,0,0
8(9) 0g, 0,
(21)
0,0,0.0
V 800
21

---

Published in Transactions on Machine Learning Research (07/2022)
The details of this process are explained in lemma S.l(see sec. A.4) Adding in the residual connection, we have:
[[ Is I.m (111) (0))
a
Og [a 1 1,,000,,0 0,000,0,0
8(0) 0,,0,
(22)
1 (i 1) (i 1)²
Hereafter, we make use of the feed-forward layer to get:
O(af) = - I. 1 (--)) mm (0)
11
01 [ (t) m (1-1) 0,0,
(23)
0,..., ,
0, 0,
If the state is qmitt then set [ O,, else have [ [ul(1]. Note that this gives [ [u[
[a
To get all the required values, we first ( and [s(*]| to one-hot encoding of Q Σ. from there, we can calculate all the required values in a look-up table fashion. if the state is qinit then we set [ 02
The final output of this layer is then:
2 = c(+1)
(12),,,, m (6())) 0,00,
3 0,
(24)
1,i+1 1, (i+1)' (i+ 1)²
Decoder Layer 2: In this layer we calculate the symbol under the head in the next timestep. In order to do so, we first
use the self attention layer calculate [v((((+)))] and (f(i + (For details, see see A.7
Att(W322 WZz wizz)
0.....,0.
O,......
(25)
0, (1( (E(i+ 1)),
0,0,0,0
Adding the residual layer, have
a² 0,, ((+1))
(10+ 10 m() m (i-1), 1,,0,,0000 0,0,
800, [1l (Eli + 1)),
(26)
1 1, (1=1) ( 1)²
The feed-forward layer then gives (2(a?) = Ilq (++1),, [ Spere ((rr)))))))(((()))---_)))) (1-+))+--(())))) (f-1), O,... M] where M large negative value. The perceptron function the sı is added positionwise. and is 0 unless f(i+11 In this special case, it makes s contain only or which is converted into O, by the ReLU activation in the output MLP. The same is also true for every field after the first 4, where add a large negative value make the ReL.U output O.
22

---

Published in Transactions on Machine Learning Research (07/2022)
A.3 The Attention Mechanism
The attention mechanism in the Gaussian kernel is defined as follows:
Attm(Q, K, V) = V(CooNorm(((()
(27)
where @ is PAKS (Eq 34 for Theorem 1 and PRRF (Eq 33 for Theorem 2, and w, is sampled from gaussian with zero mean and diagonal covariance. However, for the proof construction, we use hard version of this attention mechanism, and limit ourselves to the standard gaussian for wsincee the mean and sigma is learnable, this can always be achieved). To begin with, we replace the kernels with their common dual, using lemma S.2 (sec [A.6p. In our construction, we do not require learnable means and variances, so we fix them to be Od, and hereafter:
h-1 Σ wov (ColNoom
[[e
d 1,d-1
( K.
(28)
1=0m=0///
where [f m m))li denotes an a x 3 matrix whose (l, m))* entry is f(l, m). ColNorm(() indicates the matrix X with its columns normalised to and d is the dimension of the query/key vector.
While this definition does not seem to allow multiplying the exponent, one must remember that the query and key matrices are calculated using projection matrices, and any required scalar factor can be incorporated into them. Therefore, we define hard gaussian attention
score (u, v) v)= =-Iuu 2
(29)
Hard attention is them computed as
[[score(qi kj) (mar, ky v) n
Au((g K, V)
(30)
== I|score (qi, kj) maary(scoreqq. ky.)]
Here I in the indicator function.
A.4
Lemma S.1
A.5 Statement
Given
q
1,4,
kj = 00 0,
0, 0,
O,... 0,
j. 1,,0, 0,
(31)
vî = [0. .0....,,,
Opp [30
j. 0,0g0, O.
0,0,0,0
For j (O,..... n., we need construction that gives
All(q. K", Vo) = 0,...,.. 0,
0g [a 10],0,0,0000
(32)
0,0,0,0
23

---

Published in Transactions on Machine Learning Research (07/2022)
A.5.1 Proof
Note that while the key and value comes from the encoder, and is therefore fixed, the query from the decoder and thus be projected as please. It is easy to construct projection matrix that gives Wqq [0,...,š,-1] -,,-1,100 0). Then we have score(q, ky) -Ili = j|| -(t - j)3, whose maxima on j' is unique and occurs at i = B((), Thus, we have Att (q. K Vo) = Vaoi which is exactly what we wanted.
A.6
Lemma S.2
A.6.1
Statement
Let
(x) = (oxp(-11xx + wh exp(= + wx)
(33)
2m
① (RKS(x)= =
[cos(w] "os æ). sin(ws x), sin x)]
(34)
We want to show that if w (O, then the kernels as defined above corresponds to the GMM kernel, ic.
(RKs(()@RKkk) = @PRF PRF (Y) me -
(35)
A.6.2 Proof
Positive Random Features:The proof actually follows from a trivial extension of Lemma in (Choromanski al. 2021b), but present it here end to end for the convenience of the reader.
First we observe that
|
1|x|1 2-x 22x Ty++ y112
2+||x|12 -1 2--111111111100
(36)
||x|| |x| IIxty1e 11y1
||ا=س|| dw in to evaluate the second factor above:
Next we leverage the fact that (2x) f
1+
=(2m) (2/7)d d,/2 ||x+y|
||w x
dw
dg/12 /
|| | 2- | 2z=" 2n To
-2/))_1/22
dw
(37)
Mull xxTxxuu" yow
(2x) d,/2
e
BB. xew"y)
The terms e -|x||² and elx|l in the last line of Eq 36 are independent of w and can thus be pushed into the expectation. Finally, we approximate the expectation by sampling in order to get the required result. Random Kitchen Finks:For the second kernel, we start with Eq. 7.4.6 in Abramowitz Stegun (972 and extend it to vectors. We have,
24

---

Published in Transactions on Machine Learning Research (07/2022)
-1 c2s(2t? x)dt
e
mm
140 m-1 i-1
cos
2nntoo +
Σ
1:00g
dtodts ... ..
1m
i=1
m
11/+ i=1
cos (2xoto) cos
2
t dtydt, ... dtm-1
mmm
m-1
e
(18
E. f
sin(2 sin (2
Σ
ditdtt ... dtm-1
Rm
The second integral involving sin odd and therefore evaluates to O. That leaves us with:
1
(t
cos(22n))) cos 2
Σ tytts
digdts ... .f.dfm-)
Rm
{
(fe
e
(os(2ralo)dtt je
Em 2 cos 2 Σ 4:0 hr0s) dll
d1m-1
Rm-1
2 VVTe
c
m- Emm
t cos
> te
dt ... dd
pm--
This process be repeated for every dimension of t and x to finally give:
f
mm/ ||4 cos(2t" x)dt
2m e |x|
(38)
Using Eq.
38
we can now get
2m k-1
(aks(X)@nKsll
cos(w x) cos(w] y) + sin(w x) sin(w( y))
2m
cos (wi(x y))
k
1=0
RJ 2
E
cos (w(x y))
(39)
11y112
cOs cos (2w T. -y y
2
2 )dw
(2 (2m)mm/² m/2
Rm
2
-1 1/2
COS(2
w T'x -
V2 )d w
m/2
Rm
/2
/2
v1??
In either case, the error stems from the approximation of the expectation by sampling, which can be made arbitrarily small by increasing
25

---

Published in Transactions on Machine Learning Research (07/2022)
A.7
Lemma S.3
A.7.1 Statement
Given,
2) 10,...,,,,ch cl42))
[1 [o ( mli-1)10,00
g(it1) 0,
(40)
1
1, (i+
(i (i+1) (i+ (1
we need a construction that gives
Au(W3 WZz? W222) [0,. (0,...,00 O, O,...,0,
(41)
0, [. (E(i 1)),
0,0,0,0)
A.7.2 Proof
We set the weight matrices to get q W?zz [0.... 0, 0lc1++1) 00,
0, 0). kj W223 0,....,0,c!1)
c(++1) (9) 0, T- and vj 2 = [o....,0, [u 1, 0.0.0.0.000 0J. All these are partial permutations and
0,
therefore can be done using appropriate binary matrices.
Note that the required output is exactly the value at j = E(i 1), so it sufficient to show that the score(qi, kj) is maximised in j for E(i i..
max (2'(c00) = c(+1)) if 3j' s.t. cl'') = c(f+1)
(42)
i,
otherwise
Now we have score(q,, kj) (cl(i+1) c(1)2= (j+1)2 For j such that c(++ + cli)), the score is almost -1 since e an integer. there 3f s.t. cli) (i+1) then the corresponding score is greater that 1, and the maxima is achieved the highest such value of If such j does not exist however, then Vjj < i, score(q, k,) -1 (.1 and therefore, the maxima achieved j i.
B Mean Square Error of Linear Approximations
In this section, we calculate the Variance/Mean Square Error(MSE) in the Linear approximation of the Gaussian Kernel. Our proof is based on the similar proof in Choromanski et al. (2021b).
B.1
Random Kitchen Sinks
B.1.1
Statement
For a GMM-RKS estimator with m samples for Normal distribution with mean vector y and Covariance Matrix Σ ST'S, the variance of the estimate around its mean is given by:
MSE(PcMM-RKaS "K dUMM ((MM))RK((K))) cos 2 (44 T (k - 91) exp(-I|sT -0- ))
(43)
26

---

Published in Transactions on Machine Learning Research (07/2022)
B.1.2 Proof
MSE(OGMM-RKS(H) (GMM-RRS(K))
1 Vorus-Wu.D
(cos(w) g) cos(wW k) + sin(w; q) sin (w, k)
m
cos 9) cos( (x k) sin q) sin k)))
1 Var Naru-NNiP m²
cos(w] (q + (q - k)))
1 Var ~00,1) m²
m
(cos((n) ST + +² (q - k)) + cos((1)) St - )g k))))
(44)
4
m
Var Nu..)
(os(1) (g cos(H ( (o)
m²
4
008 2n k:) No..
cos(1] T 5T (g - k)))
m²
cos2 (1¹ (q -
(5 (q - k)))
2 cos (1² (q - k)) (x 1)))
m
Note that there is slight abuse of notation in that and X are not independently sampled, but transforms of the same sample, making the change of variable valid. Since all are iid, we can pull the summation out of the variance. Thereafter we apply Lemma from Yu et al. (2016a) to calculate the final variance.
B.2 Positive Random Features
For a GMM-PRF estimator with m samples for Normal distribution with mean vector y and Covariance Matrix Σ = ST the variance of the estimate around its mean is given by:
MSF (AGMM- &GMM
(45)
1)p(-2(llal)-) || - ( ( (x))))))) (q+ A)II) - exp(115 k)DD) m
B.2.1 Proof
MSF(@GMM (aMM))@@MMPPPtPre
Vara.N((u.)) (exp(wf (q (q+ - -100|1-1111113)))
m²
exp(
11112//lll+++111)//rrr>>>...... (1 (g (g+ k))
exp( -2(121 + + S 5 + k) + + (q k)))
(46)
m
1 exp(-2([lal1 || -² -1
m
2nT ST (q + b))) - 2 -o ST k)))))
1 exp(-2([lal1) + || - (e)xxl) (q+ k)ll) - exp k))))
m
Where the last step follows from Eq. 16 in Choromanski et al. (021b) which in turn follows from the fact that GMM-PRF is an unbiased estimator for gaussians.
27

---

Published in Transactions on Machine Learning Research (07/2022)
Model
Image
Random Predictor
10.00
Baseline Models
Softmax Trans.
Crans.Vasswan) (Vaswani et al.)
42.44
Synthesizer Tay et al.)
41.61
Sinkhorn Tay et al.)
41.23
Sparse Trans. Trans.(Chil) (Child et al.)
44.24
Reformer Kitaev et al.)
38.07
Local Attention (Parmar al.)
41.46
Longformer(Beltagy et al.)
42.22
Linformer(Wann WWag et al.)
38.56
Big Bird Zaheer al.)
40.83
LinearElu Katharopoulos et al.)
42.34
Performer Choromanski et al.)
Performer(Choromannee
42.77
Kernelized Transformers
GMM-RKS (Eq.
5)
42.33
FASTFOOD-RKS (Eq.
36.74
GENERATIVE- RKS
(EC
39.84
GMM-PRF (Eqs.
39.94
FASTFOOD-PRE (Eqs.
12
38.31
GENERATIVE-PRT (Eqs.
40.01
Table 5: Experimental results the Image dataset with 1024 tokens from LRA benchmark.
C
Experimental Details
C.1
Source Code
We implemented KI.- TR-NSORRRORMR in Python and PyTorch (Paszke al. 2019) and plan to open-source the code for reproducing all experiments upon acceptance.
C.2 LRA image dataset results
For completeness we also give results on Image datasets from LRA task where x N image is flattened into a sequence of N2 pixels which is then provided as input to the model. The gray-scaled CIFARIO image classification dataset (Krizhevsky 2009) used, resulting in sequence length of 1024.
C.3
Hyperparameters for LRA Tasks
Further Notes:
To benchmark memory in Figure used batch size of 32 for Text and batch size of for Retrieval.
For GMM-RKS and GMM-PRF the number of components the mixture was set to C'
28

---

Published in Transactions on Machine Learning Research (07/2022)
Parameter
ListOps
Text
Retrieval
Image
Batch Size
32
32
256
Learning Rate
x 10-3
5 x 10-2
X 10-2
X 10
Training Steps/Eppochs
10K/NA
20K/NA
5K/NA
NA/200
Optimizer
Adam with Weight Decay (81
0.9, 82
0.98)
Weight Decay
0.1
0.1
0.1
0.0
Warmup Steps
1000
8000
8000
175
Scheduler
Sqrt Decay
Sqrt Decay
Sqrt Decay
Cosine Decay
Loss
Cross Entropy
Sequence Length
2000
4000
4000
1024
Num. Layers
6
4
Num. Heads
8
4
Embedding Dim.
512
256
128
128
Key/Query/Value Dim.
64
64
32
Feedforward Dim.
2048
1024
512
128
Dropout Rate
0.1
0.1
0.1
0.3
Activation Function
Gelu
Gelu Sinusoidal CLS
Gelu
Gelu
Positional Encoding
Sinusoidal CLS
Sinusoidal CLS
Learnable
Pooling Mode
CLS
Table 6: Hyperparameters for LRA tasks.
Model
ListOps
Text
Retrieval
Image
GMM-RKS
256
128
64
128
FASTFOOD-RKS
64
64
GENERATIVE-RKS
256
256
128
128
GMM-PRF
256
256
128
128
FASTFOOD-PRE
64
64
32
8
GENERATIVE-PRF
128
128
64
8
Table 7: Number of random samples M used within each KL.RRANSFORMER.
29

---

Published in Transactions on Machine Learning Research (07/2022)
C.4 Hyperparameters for GLUE Tasks
Parameter
Value(s)
Pre-Training Batch Size
64
Batch Size
Pre-Training Learning Rate (Apec)
10 10
Pre-Training Learning Rate at Step
min 10000 Igm 10000
* Tpe
Training Learning Rate (Main))
{2 10-3
10 10-4__ 5 x 10
2 10
5,5 5 10-0}
Training Learning Rate Step (Yhmin)
min(
June 0.9
Jure
1/rain
Pre- Pre-Training Epochs
Training Epochs
10
Optimizer
Adam with Weight Decay (81 0.9, B2 0.999)
Weight Decay
0.01
Loss
Cross Entropy
Sequence Length
512
Num. Layers
Num. Heads
10
Embedding Dimension
300
Key/Query/Valuu Dimension
64
Transformer Feedforward Dimension
512
Classifier Feedforward Dimension
128
Dropout Rate
0.1
Transformer Activation Function
Gelu
Classifier Activation Function
Tanh
Positional Encoding
Sinusoidal CLS
Pooling Mode
Num. of Samples from Distribution
{64, 128)
Table 8: Hyperparameters for GLUE tasks. Where multiple parameters were tried, they are listed in curly brackets. Ipre denotes the total number of pre-training steps, whereas Iune denotes the total number of fine-tuning steps on each GLUE task.
Our model has significantly fewer number of parameters as compared to Devlin al. (2019) and therefore we perform poorer than them all datasets. They use 24 layers with 16 heads each. If reported in the same order the columns of Table 3
their numbers would look like: 97.5, 89.3/ 85.4, 72.1/89.3, 87.2 87.2/86.4 92.7, 65.1, 70.1.
While we had to limit model sizes due to resource limitations, this handicaps all models equally, and therefore should not prevent comparison across various models reported in our paper.
C.5 Further Results on Efficiency Benchmarks
Figure shows how much memory each model uses. Piouse this memory usage against the model performance
C.6 Correlation of Variance Metrics
Figures 9 and 10 show the correlation of Average Gain by Voting with the other twp variance metrics for the text and retrieval tasks respectively
30

---

Published in Transactions on Machine Learning Research (07/2022)
Layer
Head
Minimum
Maximum
Mean
7.601e-7
4.584e-1
3.531e-2
2.380c-8
4.863e-1
3,576e-2
1.400e-5
4.469e-1
3.319e-2
4.214e-6
4.197e-1
3.208c-2
4.242c-6
2.790c-1
3.070e-2
2.971e-6
3.238e-1
3.161e-2
2
3
5.081e-6
2.842e-1
3.078e-2
4
5.212e-7
3.006e-1
3.086e-2
2.297e-6
2.219e-1
2.872e-2
7.094e-6
2.036e-1
2.594e-2
3
1.074c-7
1.953e-1
2.772c-2
6,330e-9
1.836e-1
2.442e-2
Table 9: Distribution of Eigenvalues of covariances of GMM-RKS model for the Text task
Layer
Head
Minimum
Maximum
Mean
1.363e-5
9.660e-2
2.440e-2
6.659e-5
7.900e-2
2.200e-2
7.241c-6
9.050e-2
2.810c-2
4
2.760e-5
7.790e-2
2.630e-2
2.068e-5
7.520e-2
2.430e-2
2
2.386e-5
7.580e-2
2.400e-2
2
7.256e-6
7.400e-2
2.370e-2
1.911e-4
8.080e-2
2.680e-2
2.871e-5
6.710e-2
2.570e-2
6.117e-6
7.370e-2
2.190e-2
7.830c-6
6.600e-2
2.010e-2
2.186e-5
7.620e-2
2.810e-2
Table 10: Distribution of Eigenvalues of covariances of GMM-PRE model for the Text task
Layer
Head
Minimum
Maximum
Mean
4.011e-05
4.994c-01
5,998e-02
2
4.250e-05
2.294c-01
5.390e-02
3
1.918e-04
3.427e-01
5.842e-02
4
5.615e-06
3.943e-01
5.866e-02
1
3.232e-06
2.605e-01
5,469e-02
2
2.746e-06
2.581e-01
5.333e-02
3
6.364e-06
2.329c-01
5.091c-02
4
2.071e-05
1.743e-01
5.020e-02
1
4.358e-05
1.919e-01
5.157e-02
2
1.143e-04
1.738c-01
4.819c-02
3
3
2.546e-05
1.735e-01
4.705e-02
4
9.697e-07
2.075e-01
4.982e-02
Table 11: Distribution of Eigenvalues of covariances of GMM-RKS model for the Retrieval task
31

---

Published in Transactions on Machine Learning Research (07/2022)
GMM RKS
Prstfood-RKS
Memory Consumption (GB)
Generative-RKS
GMM-PRI
Fastfood PRF
Generative PRF
Image
ListOps
Text
Retrieval
Figure Peak memory used by KL-TRANSFORMERS across different datasets.
Transformer
Memory Consumption (GB)
GMM RKS
GMM-PR PRE
Perlomaer
Fantfond-PRI
20
24
28
32
36
ListOps Accuracy
Figure 8: We demonstrate the peak memory consumption (y-axis) and performance (x-axis) of the various Kernelized Transformer architectures on the ListOps dataset from LRA. Memory usage refers to device memory usage across each GPU.
C.7 Eigenvalues of Trained Models
D Ablation Studies
D.1 FastFood Attention
In the main paper, use FastFood-SGB, which has all the diagonal matrices learnable. However, B and G' matrices have very special structure (their elements being drawn from Bernoulli(-I.) (0.5) and / 1) respectively), which is lost if we make them learnable. Therefore, it makes sense to have FastFood-S, which only has S learnable. Finally, we can also have everything fixed, giving us the basic FastFood version. The results of these two versions, along with the original FastFood-SCCe kernel on the GLUE benchmaark are summarised in Table [3] As one can see, FastFood- SGB is either the best or close except for WNLI and CoLA, therefore we choose use this version for our main analysis.
32

---

Published in Transactions on Machine Learning Research (07/2022)
PI vs AGV
RSD vs AGV
GMM-RKS
1.015
Generative-RKS
AGV
1.010
FastFood-RKS
GMM-PRF
1.005
Generative-PRF
1.000
FastFood-PRF
10
0.0
0.2
0.4
PI
RSD
Figure 9: Correlation between AGV the other two variance metries on Text task.
Pl vs AGV
RSD vs AGV
GMM-RKS
1.006
Generative-RKS
1.004
FastFood-RKS
AGV
GMM-PRF
1.002
Generative-PRF
1.000
FastFood-PRF
5
0.0
0.2
Pl
RSD
Figure 10: Correlation between AGV and the other two variance metrics on the Retrieval task.
E Sparsity Synthetic Experiment
E.1 Task Description
Given a sequence ordered pairs (vi, ax), where v {-1,1} and a; E {0, 1}, the task is to output Σ 0;4; Here, V2 can be seen as the value of given position, while a; indicates whether or not we need to attend to that position. The dataset generated by (pseudo-randomly flipping coin independently for a, and u. The bias in the flip for defines the sparsity of the dataset. The flip for is appropriately biased to ensure that no prefix has absolute of more that 4. The final prediction is outputted 9-way classification over integer values between -4 and (both
Layer
Head
Minimum
Maximum
Mean
1
8,225e-05
2.114c-01
4.818e-02
2
5.386e-05
1.525e-01
4.066e-02
3
9.659e-05
1.707c-01
5.638e-02
4
1.498e-07
1. 1.467e-01
4.166e-02
1.294e-06
9.798c-02
3.024c-02
2
9.306e-05
1.289c-01
4.334c-02
3
4.300e-06
1.865e-01
4.846e-02
4
2.418e-05
1.047e-01
4.458e-02
1
2.466e-08
1.133e-01
3.987e-02
2
2.291e-04
1.429e-01
5.099e-02
3
3
3.997e-05
1.134e-01
3.380e-02
4
6.653e-05
1.129c-01
2.827c-02
Table 12: Distribution of Eigenvalues of covariances of GMM-PRF model for the Retrieval task
33

---

Published in Transactions on Machine Learning Research (07/2022)
SST2
MRPC
MRPC
QQP
QQP
MNLI
MNLI
QNLI
WNL.
RTE
CoL.A
Dataset
(acc)
(acc)
(fl)
(acc)
(f1)
(mat)
(mis)
(acc)
(acc)
(acc)
(MCorr)
FastFood
0.814
0.713
0.820
0.811
0.738
0.571
0.568
0.629
0.634
0.563
0.152
FastFood-S
0.807
0.706
0.822
0.810
0.741
0.571
0,571
0.642
0.606
0.570
0.101
FastFood
0.828
0.707
0.820
0.810
0.739
0.572
0.638
0.592
0.129
SGB
Table 13: Ablation studies using FastFood variants on the GLUE benchmark.
Dataset Sparsity
Checkpoint
GMM-RKS
GMM-PRF
GMM-RKS/GMM-PH
Acc.(%)
Std. Dev.
Abs. Mean
Std. Dev.
Abs. Mean
Std. Dev.
Abs. Mean
20
0.083064
0.144132
0.079822
0.139948
0.960974
0.97097
40
0.094528
0.16609
0.11077
0.193444
1,171828
1.164689
0.1
60
0.099791
0.176456
0.116044
0.201462
1.162868
1.141712
80
0.131252
0.228456
0.154802
0.275594
1.179433
1.206332
20
0.069646
0.121306
0.434024
0.53019
6,231835
4.370671
40
0.11552
0.201468
6.432806
12.17245
55,68558
60.41891
0.5
60
0.142631
0.249917
3.752173
5,886064
26.30681
23.55206
80
0.166899
0,295308
22.46943
42.14157
134.6289
142.7036
20
0.059153
0.102415
0.065438
0.115866
1.106237
1.131343
40
0.108255
0.186704
0.113171
0.211465
1.137784
1.132622
0.9
60
0.119243
0.206855
0.126247
0.217067
1.058737
1.049365
80
0.136816
0.240525
0.168633
0.295976
1.232546
1.230542
Table 14: Means and Variances of gradients received at the classifier layer for the synthetic experiment
inclusive). The bias against higher absolute values causes -4 and classes appear less often. This is balanced out by overgeneration and sampling.
We generate 3 datasets corresponding to sparsities 0.1, 0.5 and 0.9. Each dataset has 200K instances, of sequence length 200. Of these, we use 80% as the training set and the rest for validation.
E.2 Model Description
We 3 layer transformer with dmoode! d feedf orward = 64 and dquery == dvalue 16. The input is encoded 3-d many-hot vector (v; 1, 1; 1, a:). this is then passed through an embedding layer and added to learnable position embeddings and passed through the transformer. The final embedding of the t position is then passed through hidden layer with 64 units and then passed on to the final layer for 9-way softmax. For both linear attention models, we use 64 samples.
All models are trained using AdamW optimizer with 32 0.98, 10-9, weight decay=0.11 1, Learning Rate=5 10-6 and all other parameters default. We use SGD with batch size of 400 and entropy loss.
E.3 Gradients
We create checkpoints for the GMM-RKS and GMM-PRF models when they first achieve validation accuracies of 20%, 40%, 60% and 80%. For each of these checkpoints, we pass the first 50 validation datapoints and record the gradients the classifier layer. This process repeated 50 times. Thereafter, we calculate the mean and standard deviation of gradients each neuron. To avoid cancellation of opposite signs, the mean is calculated over absolute values. The final reported numbers averages over the 64 neurons. The results are shown in Table [4
34


Reviews:
[{"review_id": "x6WWJbiQae", "reviewer": "Reviewer_tXbd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\nThe paper is well written and structured. The problem is clearly motivated and abstracted. Related work is discussed sufficiently. Although the ingredient are not new, the overall approach gives sufficient noveltely and is shown that it can work more accurately than strong baselines. The computational complexity is linear. A sound and non-trivial theoretical analysis is presented. Although the novelty of the approach could be described a bit better and delineated from the related work, all in all I find this a sound and relevant contribution to the field, worth publishing in TMLR.\n\n\n", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors present a new transformer. As proposed in Tsai et al (2019), the self-attention mechanism is here defined using a kernel function to evaluate the similarities of keys and queries. The approach approximates the transformer kernel as a dot product of spectral feature maps, similar to Choromanski et al. (2021). The approach's core novelty (given in Section 2.1) is a method to learn the transformer's kernel. To due so, the authors invoke the framework for kernel learning using GMMs in regression and classification tasks by Oliva et al. (2016) and adapt it to transformers. The proposed transformer is shown to achieve higher accuracy on the LRA benchmark than several baseline methods (including Performer), while having on par (or better) computational complexity. Further experiments reveal a tradeoff in accuracy and efficiency and analyze the effectiveness of sparse data. The appendix shows that the method's accuracy can also be (slightly) inferior to some baseline models (e.g., Table 5). A theoretical analysis is presented showing that the proposed method(s) are Turing complete and giving an analytical characterization of the MSE of the approximation. Proofs are shown in the appendix.\n", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "Minor comments:\n- Describe novelty better and delineate from related work. (Your core contribution seems the aspect of learning the kernel. However, the name \"We propose the kernelized transformer\" indicates as if you were the first to use kernel methods in transformers.)\n- PRF is defined twice on page 2\n- space missing before parentheses in Table 2", "broader_impact_concerns": "-", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Htcgko9HfP", "reviewer": "Reviewer_fZ1L", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. The paper is very well-written and easy to follow.\n2. This paper is \"complete\" with clear motivations, efficient methods, detailed analyses, and strong empirical results. Overall, this paper attempts to attack an important problem of quadratic time complexity in the self-attention module. Though there are many other existing works already, this paper makes distinct contributions by introducing learned kernels and efficient implementations.\n3. The authors have conducted comprehensive empirical studies with different tasks and situations. Overall, the empirical results seem to be convincing, but I'm not an expert and are not familiar with the datasets used in this paper.\n\nWeaknesses:\n1. Mistakes in time and space complexity for FastFood: if I understand correctly, FastFood parameterization reduces the time complexity to logarithmic time, but this is not reflected in Table 1.\n2. Discussions of the empirical results: in Section 3.1, the authors present the results on long sequences with improved performance. However, I find that the results are quite different for text and retrieval. While the RKS kernel is better for text, PRF is better for retrieval. I feel some discussions are required. The authors do give an explanation based on the variance analysis, but it is more of a guess. Besides,  both GMM PRF and generative PRF underperform compared to Performer and LinearElu, which is concerning. I think the authors need to dig a bit deeper and give a clear recommendation to other users.\n3. Some conclusions are a little flimsy. In terms of the variance analysis, it seems that the PRF kernel has a higher variance, but on the other hand, it has smaller eigenvalues. I don't know how the authors reach the conclusion that it is consistent with Theorem 3. Also, it is unclear about the casual relationship between variance and performance. Often the case, high variance leads to slow convergence in training, but not necessarily bad performance. To make any scientific arguments (or conclusions), the authors need to design control experiments to verify, rather than just loosely connecting them.", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes to learn the kernel of the self-attention layer in Transformer architectures. Inspired by kernel methods and random Fourier feature models, the authors approximate the kernel function with explicit feature maps by sampling from a spectral distribution. For computational reasons, the authors further approximate the spectral distribution with either a mixture of Gaussians (with a clever parameterization) or a flexible generator. As a consequence, the resulting kernelized transformer attains a linear complexity in terms of the sequence length.\n\nIn theory, the authors proved that the resulting architecture with positional encodings is Turing-complete. In practice, the model with learnt kernels improves the performance of long-context tasks, while staying competitive to standard transformer with softmax kernel in short-context tasks.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "Overall, I like this paper. Albeit the idea of learning the kernel is straightforward, this work makes solid contributions to reducing the complexity of the self-attention module while retaining the performance.\n\nMy main request is to remove these flimsy arguments (as mentioned in weaknesses). Many conclusions in the paper are merely scientific guesses or conjectures.\n\nThe authors have discussed the variance of the kernelized attention module, but less is done for bias analysis. I think empirical bias analysis should be easy and it is quite interesting to see how the approximation quality or model performance changes as a function of the number of samples. To further reduce the bias, I recommend the authors check out the Russian Roulette Estimator of approximating an infinite sum. See for example [1].\n\nReference:\n1. Efficient optimization of loops and limits with randomized telescoping sums. ICML 2019.\n", "broader_impact_concerns": "The main contribution of the paper is on making the transformer architecture more computationally efficient, so it is more on the methodology. I do not expect there to be any direct negative societal impact rom this contribution. In addition, all experiments were done in public datasets. That being said, a better transformer model might be used in many applications (e.g., language model) with toxic outputs. However, this is out of the scope of this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "JTpsGq36YZ", "reviewer": "Reviewer_x9TU", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strenghts:\n- The idea is sound and the paper is easy to read, except for one contribution (see below).\n- Experiments are comprehensive, except for the efficiency evaluation (again, see below).\n- Theoretical claims are clear, but I have not checked the full derivations in the appendix.\n\nWeaknesses:\n- The novelty can be considered minor, since the idea of using RFFs inside transformers is known, and the methods they use to learn the kernel approximation are all leveraged from the corresponding literature, with only one exception.\n- The paper seems to be stuck at mid-2021 in terms of citations, comparisons, related works, etc. I am assuming this is a prior submission that was not updated. \n- The use of deep generative techniques is not clear (see below on requested changes).\n- I do not understand why the efficiency comparison is made by considering a different set of methods compared to the accuracy claims.\n\nOverall, I would vote for accepting the paper if these points are addressed.", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper builds on the idea of linear transformers, which reformulated the standard attention mechanism by rewriting it in terms of kernel functions. In this paper, the kernel is approximated with random Fourier features (RFFs), and the underlying spectral distribution is learnt with different mechanisms. A few theoretical claims are made about the method (notably, that it keeps Turing completeness of the original mechanism). Experiments show that the method performs better (in terms of accuracy and/or speed) compared to some baselines.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "The only part of the paper which is not clear to me is the use of generative models (Fig. 2). If the generator is trained together with the rest of the network, sampling should be made differentiable, but this not mentioned. If it not trained together (somewhere in the text it is mentioned that samples are generated every fixed number of iterations), then, how is it trained?\n\nIf I understand correctly, a different kernel is learnt for every head in the network. It would be interesting to check whether it is possible to learn a shared distribution for the heads, or even for different layers.", "broader_impact_concerns": "I do not see significant ethical concerns here.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers were generally appreciative of the article and in favor of acceptance.  Reviewer fZ1L has remaining concerns about whether claims about the variance are supported.  In particular, they disagree with the logical entailment of the claim that the empirical results are consistent w/ Theorem 3 (last sentence of 4.1), and "it is unclear about the casual relationship between variance and performance" (point 3 of "weaknesses" in their original review).  These claims should be modified in the final version.
Decision: Accept (poster)