Hybrid Random Features
Authors: Krzysztof Marcin Choromanski
Han Lin
Haoxian Chen
Arijit Sehanobish
Yuanzhe Ma
Deepali Jain
Jake Varley
Andy Zeng
Michael S Ryoo
Valerii Likhosherstov
Dmitry Kalashnikov
Vikas Sindhwani
Adrian Weller
Venue: ICLR 2022 Poster (2022)
Keywords: random features
softmax kernel
attention mechanism
compositional kernels

Abstract: We propose a new class of random feature methods for linearizing softmax and Gaussian kernels called hybrid random features (HRFs) that automatically adapt the quality of kernel estimation to provide most accurate approximation in the defined regions of interest. Special instantiations of HRFs lead to well-known methods such as trigonometric (Rahimi & Recht, 2007) or (recently introduced in the context of linear-attention Transformers) positive random features (Choromanski et al., 2021). By generalizing Bochner’s Theorem for softmax/Gaussian kernels and leveraging random features for compositional kernels, the HRF-mechanism provides strong theoretical guarantees - unbiased approximation and strictly smaller worst-case relative errors than its counterparts.  We conduct exhaustive empirical evaluation of HRF ranging from pointwise kernel estimation experiments, through tests on data admitting clustering structure to benchmarking implicit-attention Transformers (also for downstream Robotics applications), demonstrating its quality in a wide spectrum of machine learning problems.

Published as a conference paper at ICLR 2022
HYBRID RANDOM FEATURES
Krzysztof Choromanski" 12, Haoxian Chen Han Lin Yuanzhe Ma* Arijit Sehanobish Deepali Jain', Michael S Ryoo', , Jake Varley', Andy Zeng', Valerii Likhosherstov Dmitry Kalashnikov', Vikas Sindhwani', Adrian Weller 'Google Brain Robotics, "Columbia University, 3 University of Cambridge, "The Alan Turing Institute
ABSTRACT
We propose new class of random feature methods for linearizing softmax and Gaussian kernels called hybrid random features (HRFs) that automatically adapt the quality of kernel estimation to provide most accurate approximation in the defined regions of interest. Special instantiations of HRFs lead to well-known methods such trigonometric (Rahimi & Recht 2007) or (recently introduced in the context of linear-attention Transformers) positive random features (Choro- manski al., (021bb. By generalizing Bochner's 's Theorem for softmax/Gaussian kernels and leveraging random features for compositional kernels, the HRF- mechanism provides strong theoretical guarantees unbiased approximation and strictly smaller worst-case relative errors than its counterparts. We conduct ex- haustive empirical evaluation of HRF ranging from pointwise kernel estimation experiments, through tests on data admitting clustering structure to benchmark- ing implicit-attention Transformers (also for downstream Robotics applications), demonstrating its quality in a wide spectrum of machine learning problems.
INTRODUCTION & RELATED WORK
Consider the softmax and Gaussian kernel functions K Rdxd + I defined as follows:
SM(x, y) def exp(x) y), K gaauss (x exp(
|
These two are prominent examples of functions used in the so-called kernel methods (rreton al., 2005 Zhang et al. 2018) and beyond, i.c. in softmax-sampaa Teatures (RFs. Rahimi & Recht 2007 Liu al. 2020 Peng et al. (021) yield powerful mechanism for linearizing and consequently scaling up kernel methods with dot-product kernel decompositions disentangling x from y in the formulae for kernel value K(x, y) ~ "(x)) ((y) via data-agnostic probabilistic (random feature) maps
(Blanc & Rendle 2018). Random
Rd + RM. The tight relationship between softmax and
2 SMM(X, y) exp( 11yl12 2 provides a mapping of any random feature vector OsM (u) for the softmax kernel to the correspond- ing one @gauss (u) = exp(- 2
Gaussian kernels given by the transformation Kgauss(X, yx) = exp( - ||x|
F (u) for the Gaussian kernel, thus we will focus on the former kernel. The classic random feature map mechanism øtrig for the softmax kernel, obtained from Bochner's Theorem applied to the Gaussian kernel (Rahimi & Recht 2007). is of the form:
trig
1
cxp(
||u||
(si sin..,,, u), cos(w] u). .... cos( uo))
(2)
m 2
where m stands for the number of random features and wi, Wm id
(0,
The above (common in most downstream use-cases) method for linearizing softmax/Gaussian ker- nels was recently shown to fail some the new impactful applications of scalable kernel methods such implicit-attention Transformer architectures called Performers (Choromanski et al. 2021bb. We denote by MSE the mean squared error of the estimator (i.e. its variance since all estimators considered in this paper unbiased). The above mechanism struggles to accurately approximate close-to-zero kernel values, as characterized by the particularly large MSE in that region. This is a crucial problem since most of the entries of the attention matrices in Transformers' models are very
"Equal Contribution, Correspondence to kchoro @ google.com
AAthhrsshi in alphabetical order Independent researcher. Work done during postdoc Yale University

---

Published as a conference paper at ICLR 2022
small and the approximators need to be particularly accurate there. Otherwise the renormalizers computed to make attention matrices row-stochastic (standard attention normalization procedure in Transformers) would be estimated imprecisely, potentially even by negative values.
The solution to this problem was presented in FAVOR+ mechanism (Choromanski et al. 2021b), where new positive random feature map for unbiased softmax kernel estimation applied:
(u) 1 exp(- 2
|uul? (exp(w] u). ..... exp(w,n, u), exp(-w]] u). ..., exp(-w, w))
22m
Even though, very accurate in estimating small softmax kernel values (which turns out to be crucial in making RFs work for Transformers training), this mechanism is characterized by larger MSE for large kernel values. In several applications of softmax kernels (in particular Transformers, where attention matrices typically admit sparse combinatorial structure with relatively few but critical large entries and several close-to-zed ones or softmax sampling) the algorithm needs to process simulta- neously very small and large kernel values. The following natural questions arise:
Is it possible to get the best from both the mechanisms to obtain RFs-based estimators particularly accurate for both very small and large softmax kernel values Furthermore, can those estimators be designed to have low variance in more generally pre-defined regions? ?
We give affirmative answers to both of the above questions by constructing a new class of random feature maps techniques called hybrid random features HRFs. Theoretical methods used by us to develop HRFs: (a) provide a unifying perspective, where trigonometric random features from Bochner's Theorem and novel mechanisms proposed in (Choromanski et al.] (021b) are just spe- cial corollaries of the more general result from complex analysis, (b) integrate in the original way DDemmns-WWilimson several other powerful probabilistic techniques such as Goemans- Willimson method (Goemans & Williamson 2004) and random features for the compositional kernels (Daniely et al. 2017).
We provide detailed theoretical analysis of HRFs, showing in particular that they provide strictly more accurate worst-case softmax kernel estimation than previous algorithms and lead to computa- tional gains. We also conduct their thorough empirical evaluation on tasks ranging from pointwise kernel estimation to downstream problems involving training Transformer- or even end-to- end robotic controller-stacks including attention- architectures.
Related Work: The literature on different random feature map mechanisms for Gaussian (and thus also softmax) kernel estimation is voluminous. Most focus has been put on reducing the vari- ance of trigonometric random features from (Rahimi & Recht 2007) via various Quasi Monte Carlo (QMC) methods, where directions and/or lengths of Gaussian vectors used to produce features are correlated, often through geometric conditions such orthogonality (Choromanski
al. 2017 Rowland et al., 2018; Yu et al. 2016 Choromanski et al. 2019 Choromanski & Sindhwani 2016). Our HRFs do not compete with those techniques (and can be fact easily combined with them) since rather than focusing on improving sampling mechanism for a given approximation algorithm, they provide a completely new algorithm. The new application of random features for softmax ker- nel in Transformers proposed in (Choromanski et al.) 20200 2021b) led to fruitful research on the extensions and limitations of these methods. [Schlag et al. (202)) replaced random features by sparse deterministic constructions (no longer approximating softmax kernel). Luo et al. (222I) observed that combining L2 noomaalllzttoonn of queries and keys for variance reduction of softmax kernel es- timation with FFT-based implementations of relative position encoding and FAVOR+ mechanism from Performers helps in training. Trigonometric random features were applied for softmax sam- pling in (Rawat al. (2019). Several other techniques such Nyström method (Yang et al. [012]] Williams & Seeger] 2000) Rudi et al.) 2015) were proposed to construct data-dependent feature representations. Even though,
we show in Sec. [2.4] certain instantiations of the HRF mecha- nism clearly benefit from some data analysis, our central goal remains an unbiased estimation of the
softmax/Gaussian kernels which is no longer the case for those other techniques.
HYBRID RANDOM FEATURES
2.1 PRELIMINARIES
Whenever we do not say explicitly otherwise, all presented lemmas and theorems are new. We start with the following basic definitions and results.
Definition 2.1 (Kernel with a Random Feature Map Representation), We say that a kernel function
K Rd Rd - R admits random feature (RF) map representation it can be written as
2

---

Published as a conference paper at ICLR 2022
kxx,y) = Eu
Σ Ei w) Ei(Y)
(4)
for some & Rd R R, and where wi sampled from some probabilistic distribution Ω P(RR). The corresponding random condmmfeature map. for given m E N. is defined as
@m(u)
vm Ln (u) *... o (u) E g
(5)
where 6'n (E, (u, wi).... ..... (u wm)) stands for vertical concatenation, and wT, ~
Ω2
Random feature maps can be used to unbiasedly approximate corresponding kernels, as follows:
R(x.y/ y) 6m (x) syly).
Using the above notation, a trigonometric random feature øirig from Equation 2 can be encoded as applying 1 = Ç (u, w) sin(w" u), (2, cos(w u). Similarly, positive random features can be encoded as taking / = 2 and E (u, w) = V2
exp(w T u), (2 w) = V2 exp(-w u).
The following result from (Choromanski et al., (02Tb) shows that the mean squared error (MSE) of the trigonometric estimator small for large softmax kernel values and large for small softmax kernel values, whereas an estimator applying positive random features behaves in the opposite way.
Denote by SMm trig, (x,y) an estimator of yM for x,y Rd using trigonometric RFs and W1.....
idd (O, Ia). Denote by SM,, m (x, its analogue using positive RFs. We have:
Lemma 2.2 (positive versus trigonometric RFs). Take A = x y, 2 = x + y, [1 = (2m) exp(u²) (SM-2(xxxx fz(u) = (2m)-1 exp(u?)SMM (x. fa(u) = -1- - cxp((up))
The MSEs of these estimators are:
trig
MSE(SM (x,y)) (1 MSE(SM, (x.y)) = fe(lz(l))ss(lllll (7)
2.2 THE ALGORITHM
We are ready to present the mechanism of Hybrid Random Features (HRFs), Denote by E
(x, list of estimators of SM(x,y) (the so-called base estimators) and by A (SM &
(x*(x, a list of estimators of {X*(x, Po for some functions X .dd x IRd + [0, 1]. constructed independently from &, Take the following estimator of SM)
SM E.A
(x,y)
Σ (xxx)))SN'xxx))) + (1-I (x,
SM (x,y)
(8)
In the next section, we explain in detail how base estimators are chosen. We call SM yx y) hybrid random feature (HRF) estimator of SM(x, y) parameterized by E, The role of the X- coefficients is to dynamically (based on the input (x, prioritize or deprioritize certain estimators to promote those which are characterized by lower variance for given input. Note that if elements SM E.A.
E. (x, y) is also an unbiased estimator
of & are unbiased estimators SM(x, y). then trivially of SM((xy)) y). Assume that each SM (x,y) ) is of the form SM (x,y) = (ot, (x)) @b.m (y) for o (u) -
w jim (u) ...
tx o (u), & and jim 2.m
fe Rd R", where j {1,2).
Assume also that X (x( y) can be written as:
(xxx)) as Erann (x, (f
for some scalars E R, distribution Ω E P(RR)) (where P(Rd stands for the set of probabilis- tic distributions on RRd), mappings E m/. RRd Rdd R and that the corresponding estimator
{*(x))) = (x y) of y) is of the form:
(x,
(x)) Tx py,
(10)
3

---

Published as a conference paper at ICLR 2022
for (u) = w j.n (u) * .... Pjjn
1.k
( (u), (j) (u) (1) (u, Tn Tn)) ",,..., On ?2
and where j E (1. Linearization of the X-coefficients given by Equation is crucial to obtain linearization of the hybrid estimators, and consequently random feature map decomposition.
We denote a hybrid estimator using n random features for its base estimators and m to approxi- hyb mate A-coefficients as SMm,n Furthermore, for two vectors u, v, we denote their vectorized outer- hyb product by u @ Finally, we denote by II" cectors-concatenattoo Estimator SM,, (x,y) y) can be rewritten as a dot-product of two (hybrid) random feature vectors, as the next lemma shows.
kybb Lemma 2.3. The HRF estimator SM m.n Xx y) satisfies SM 5Mm,n (x,,)(x,y)) = V (x) T V2(y),, where V; for j (1,2) is given as vz(z)= = 3 (z) Vz (2((2) (2) and:
hyb
(Z)
II Loop
ak 1 Pim (z) * ... * ta (z)
m
(2(2)) = /
II
Pi. @ sf (z)
l= I. je ( d. (x)
(11)
Vi((2)
k= Ok 11.pt J.m
(z) *... 20) 0 (z)
m
(z)
II
J.m (z)
mm
L (E (1.....a xx (1.....pp11
Bipolar estimators:
A prominent special case the general hybrid estimator defined above is the hyb (x,y)), Thus consider the following estimator SMM,..
trig
one where: (SM
yx, SM
hyb SM, (x,y) (x. (x.y) + (1 - A, (x, y) SSM trig (x. y).
(12)
trig The question arises whether SM defined in such way can outperform both SM,, and SM,
bybb That of course depends also on the choice of À : Rdd x Rd - R. If we consider a (common) normalized setting where all input vectors have the same norm, we can rewrite A(x, y) ad X(0x., where Ox.y is an angle between x and y and ||x|| || = r. By our previous analysis we know that SMM,, +
trig
becomes perfect for 0 = and SM m becomes perfect for 0 = 0. That suggests particularly simple linear dependence of À on to guarantee vanishing variance for both the critical values: 0 0 and 0 π. It remains to show that such X-coefficient can be linearized. It turns out that this can be done with a particularly simple random fcature map mechanism, as we will show later, leading to the so-called angular hybrid variant (see: Section 2.4
2.3 CHOOSING BASE ESTIMATORS FOR HRFs: COMPLEX EXPONENTIAL ESTIMATORS
In this section, we explain how we choose base estimators in Equation g We denote by i a complex number such that 12 -1. The following lemma is gateway construct unbiased base estimators.
Lemma 2.4. Let w ~ (0, Then for every z (21, T E Cd the following holds:
E
Texp(wTer 2)) exp
21=1 2
(13)
2
For a complex vector z = (21, 2d) E Cd, we denote z² =
Σd= 22 Consider z = Ax +
(A 1AT)-' "y for an invertible (in (dxd) matrix A E Cddd Using Lemma 2.4
we get:
exp 2 exp ((A (1))
(Ax)²
SM(x,y) E(exp(w T (Ax + (AT)-Ty)) (14)
2
Thus for 1/m (u) def
1 Vm
exp(
(Mu)² ) (exp(wi Mu), .... exp(",, and wi NN(0
SM(x, y) = (x) ((AT) .y)].
(15)
4

---

Published as a conference paper at ICLR 2022
We obtain the new random feature map mechanism providing unbiased estimation of the softmax kernel. In general it is asymmetric since it applies different parameterization for x and y, i.c. one using A, the other one (A (AT)- Asymmetric random features is simple generalization of the cexp, model using the same & for both x and y, presented earlier. The resulting estimators SM,, (x,y), that we call complex exponential (CE) will serve as base estimators, and are defined as:
cexp SM, m
yx, def
(x) V (YAT)=(y)
(16)
For A = iLa, CE-estimator becomes trigonometric and for = Ia, becomes S.m. (x, y). Note also that estimator based on this mechanism has variance equal to zero if
x = -A (AAA)
(17)
In fact can relax that condition. is easy check that for the variance be zero, we only need:
Re(A)x + Re((A (e((A))-)))) 0 and Im (A) xx + Im((A") 0
(18)
2.4 ADAPTING HRF ESTIMATORS: HOW TO INSTANTIATE HYBRID RANDOM FEATURES
We will use complex exponential estimators from Section 2.3 as base estimators in different insanti- ations of HRF estimators (that by definition admit random feature map decomposition). The key to the linearization of the HRF estimators is then Equation 9 which implies linearization of the shifted versions of X-coefficients (Equation [0) and consequently desired random feature map decompo- sition via the mechanism of compositional kernels (details in the Appendix). The questions remains how to construct X-coefficients to reduce variance of the estimation in the desired regions of interest.
Accurate approximation of small and large softmax kernel values simultancously: As we explained earlier, in several applications of softmax estimators desired property is provide ac- curate estimation in two antipodal regions small and large softmax kernel values. This can be done
in particular by choosing p = 1, SM SM SM² = SM trig and A(x.y) = fxyy where 6x.y stands for an angle between x and y. The key observation that y(x,y) can be unbiasedly approx- imated via the mechanism of sgn random features (following from Goemans- Wiillimon algorithm (Goemans & Williamson, 20044) as follows (details in the Appendix, Sec. D
1
7(x, y) = + Px(x)) Py)(),
(19)
where Pl,n(2) = (2) = (z)
del i (sgg(( (Ti z), sgn(T,, T)) for TI....,T (O
Ven
and 12 -1. We call the resulting estimator of the softmax kernel angular hybrid. As we show in Sec. 3 this estimator is indeed particularly accurate in approximating both small and large softmax kernel values. For instance, for the prominent case when input vectors are of fixed length, its variance is zero for both the smallest (0x.y = R)) and largest (fx.y = 0) softmax kernel values (see: Fig. D.
Gaussian Lambda Coefficients: Another tempting option is to instantiate X-coefficients with approximate Gaussian kernel values (estimated either via positive or trigonometric random fea- tures). In that setting, functions X (x, are defined as (x = exp( | M.yll for some 20
MI..., Mp E Raxd and CI, Cp E IR. Since this time X-coefficients are the values of the Gaussian kernel between vectors ex and M&Y we can take ag 0 and define př...hi. as:
|x|| 20
||M,y113 20
( M&y
(x) = exp(
SM
and
P2. (y)
exp(
)
(20)
CH
Ck
where (SM is the random feature map corresponding to a particular estimator of the softmax kernel. Note that the resulting hybrid estimator, that we call Gaussian hybrid has nice necudo-combinatorial properties. The coefficients A E (0, 1] fire if x &3 -May (the firing window can be accurately controlled via hyperparameters cτ). Thus matrices Mj can be coupled with the base estimators particularly accurate in the regions, where x x -Ma. The mechanism can be thus trivially applied to the setting of accurate approximation of both small and large softmax kernel values for inputs of fixed length, yet it turns out to be characterized by the larger variance than its angular hybrid counterpart and cannot provide variance vanishing property for both 4xy 0 and 0x
5

---

Published as a conference paper at ICLR 2022
MSC
NSC
Figure 1: MSEs for three softmax kernel estimators (from left to right): SM SM and angular hybrid for m = 10, n = 1 and input of lenghts = 5. MSES are given as functions of: an angle 0x,y E [0, x] between and and r (symmetrized along 0 for length axis). For each plot, we marked in grey its slice for a fixed to illustrate that only for the angular hybrid estimator, the MSE goes zero for both 0x.y and 8x.y π.
Adaptation to data admitting clustering structure: Assume now that the inputs x come from a
distribution P(X) admitting certain clustering structure with clusters X,,...,,, C IR of centers X1, ..., Xy E Rd and that the analogous fact is true for y with the corresponding clusters Y1,.., ... Yu S Rd and centers y1, .... .......
Yw E Rd. For fixed pair of clusters' centers (x, yy), one can associate
cexp complex exponential estimator (Mx,,)) with the corresponding matrix A E Cdxd (see: Sec. 2.3
that minimizes the following loss, where |z| for z E Cd is defined as |z| = E i= 1212:
((A) = Ax + (AT)-
(21)
Note that, by Equation 17 if one can find A such that ((A) = O, then SM (x.,5) makes no error in approximating softmax kernel value on the centers Y. Furthermore, if sampled points are close to the centers and ((A) is small, the corresponding variance of the estimator will also be small. Thus one can think about SM (x:,y)) ) as estimator adapted to a pair of clusters (x,,)). We can add additional constraints regarding A, i.c. A Raxd A being diagonal (or both) for closed- form formulae of the optimal A, yet not necessarily zeroing out ((A) (details in the Appendix, Sec. Complex exponential estimators defined this way can be combined together, becoming base estimators in the HRF estimator. The corresponding X-coefficients can be chosen to be Gaussian as described in the previous paragraph (with optimized matrices M2 so that a coefficient fires for the corresponding pair of clusters' centers), but there is another much simpler choice of As. We index
pexp
H.3
different base estimators by pairs of clusters' centers (x) One can simply take p (x) (x) =
if belongs to the cluster of center x; and p1 (x.:)5) (x) O otherwise. Vectors P2 (x:35s) (y) (that also become scalars in this scenario) are defined analogously (the n symbol vanishes since p is determin- istic). Thus effectively the hybrid estimator activates precisely this random feature mechanism that is optimal for the particular pair of clusters. This scheme is feasible if gluster-mmbbery computationally acceptable, otherwise aforementioned Gaussian coefficients can be applied.
calculation
We conclude with an observation that in principle (to reduce the number of base estimators if needed) one can also construct HRF estimators that take base estimators defined in the above way only for the most massive pairs of clusters and add an additional default base estimators (for instance SM
THEORETICAL GUARANTEES
We provide here theoretical guarantees regarding HRF estimators. We focus on the bipolar setting and SM ++ The following is true:
with two base estimators: SM trig
Theorem 3.1 (MSE of the bipolar hybrid estimator). Take the bipolar hybrid estimator (x and SM (x are chosen independently i.e. their random projections are chosen independently (note that we always assume that Ax(x,)) y) is constructed in- -rig m
SM SMm, y), where SM,, trig
hyb
()) and SM, +
dependently from SM
()). Then the following holds:
byb MSE m.n (x,y)) = ET (x. MSE(SMM (x,y)) + B[(1 - 2, (x. MSEM trig yx) (22)

---

Published as a conference paper at ICLR 2022
trig Furthermore, SM (x, y) and SM, yx. apply the exact same sets of random projections, the mean squared error of the hybrid estimator is further reduced, namely we have:
lyb MSE(SM m.
1E M)M (xx))) + B[(1 7m y) 3)) trig,
(x, (
(x, y))
SSS (x, y) cos ||x |1> (x,) (x,
(23)
The exact formula on MSE(SM,, (x.y)) and MSE(SM, trig (x) is given in Lemma 2.2
We can apply this result directly to the angular hybrid estimator, since: Lemma 3.2. For the angular hybrid estimator the following holds:
(xx = 2x.y (
(0x.Y
0xy 0xy
2x. π
(24)
We see that, mentioned before, the variance of the angular hybrid estimator is zero for both 0x.y = 0 and 2xy = x if inputs y have the same length. We need one more definition.
Definition 3.3. Assume that the inputs to the estimators are taken from some given bounded set C C Rdd For a given estimator SM on feature vectors x, E C, we define its max-relaive-erro with respect to C as (c(M) = maxx.yEC 'xy (SM), where 6x.y (SM) //SE(SMMxx.))) Denote SM(x.y)
by S(r) a sphere centered at O and of radius r. Define cg.- def such that 0 = 0x (note that the mean squared errors of the considered estimators depend only on the angle 0x x,y for x, chosen from a fixed sphere).
=
(Sy( for x. E S(r) and
This definition captures the critical observation that in several applications of the softmax kernel estimation, e.g. efficient softmax sampling or Amear-attention Transformers (Choromanski et al.) 202Tb), small relative errors are a much more meaningful measure of the quality of the method than small absolute errors. It also enables us to find hidden symmetries between different estimators:
Lemma 3.4. The following holds:
6g SMM tris,
exp(2r² 21122>ssn?(] ()))() (1 exp(-4r² sin²( ())) (25)
/2m
and consequently for (r) = exp(2r²) (1 -exp(-4r3))
trig ES(r) (SM Es(r) (SM,
trig lim ca. (SM
lim co.- (SMM 0-
(
(26)
0-0
2m
Our main result shows that HRF estimators can be applied reduce the max-re of the previously applied estimators of the softmax kernel. In the next theorem we show that the max- relative-error of the angular hybrid estimator scales as & exp((r2) in the length r of its inputs as
opposed to exp(2r²) as it is the case for SM trig and SM (see: Lemma 3.4p. Furthermore, the max-eelative-erro scales as v and V - O as A - 0 and 0 + x respectively, in particular goes to 0 in both critical cases. This is not true for SM trig.
trig
nor for SM ++
Theorem 3.5. The max-relative-eroo of the angular hybrid estimator for the inputs x, y on the (ap(--42)))
sphere S()) of radius 2 1 satisfies for (()) exp(2r2)
anghyb
1
Es(r) (SM M
((r)
(27)
2m
nu
n π
<B. <8..(SM anghyh m. V
co (SM
anghyh
Furthermore, limg-s0
limg
m.n
Zurmn W (r).
vo-
Additional implications of Theorem 3.5 using the fact that ((m)) -immeniii HRFs can be con- structed in time O(nd + md + mn) (regular RFs need O(mnd)), are in the Appendix (Sec. D.4)
4 EXPERIMENTS
In this section we conduct exhaustive evaluation of the mechanism of HRFs on several tasks.

---

Published as a conference paper at ICLR 2022
4.1 POINTWISE SOFTMAX KERNEL ESTIMATION
We start with ablation studies regarding empirical relative errors of different softmax kernel estima- tors across different imputs' lengths and angles 8x y The results presented in Fig. 2 Ablations over more lengths are presented in the Appendix (Sec. G). The angular hybrid estimator most accu- rately approximates softmax kernel and has smaller max-relative-eroor than other estimators.
Trigonometric
+AVOR+
Angular Hybrid
0.3
0.4
5%
0.3
error
0.0
15
6.0
100
4.5
3.9
5-
-.0
0 x/2 11 0 m/2 1 m/2
Figure 2: Pointwise estimation of SM(x, for the Same-length 64-dim inputs (r 1.0 and r = 1.5) and various angles fx.y. Red-dotted lines are for marking zero-level. We used = 10000 estimated softmax values each subplot. The true value and the 5th and 95th quantile estimated values are shown by the left y-axis, and the empirical relative errors are shown by the right y-axis. Trigonometric estimator and FAVOR+ applied 128 random features. To make fair comparison, for the hybrid variant the configuration leading to the similar number of FLOPP operations per random feature map creation was applied. Similar gains as for the angular obtained by the Gaussian hybrid variant.
Comparison with QMC-methods: Even though, as we explained in Section I our algorithm does not compete with various methods for variance reduction based on different sampling techniques (since those methods, as orthogonal to ours, can be easily incorporated into our framework), we decided to compare HRFs also with them. We tested angular hybrid HRF variant well well- established QMC baselines: orthogonal random features (ORF) Yu et al. 2016) (regular variant ORF-reg as well the one applying Hadamard matrices ORF-Had) and QMCs based on random Halton sequences (aaaton-R)(Avin et al., 2016). For HRFs, we tested two sub-variants: with orthogonal (HRF-ort) and regular iid random features (HRF-idd). We computed empirical MSEs by averaging over 100 randomly sampled pairs of vectors from two UCI datasets: wine and Boston. HRFs outperform all other methods, as we see in Table [
Table 1: Comparison of different estimators of the softmax kernel on the datapoints from two UCT datasets: wine and Boston in terms of the MSF (measured in 10- units). The non-HRF estimators apply 512 random features and the HRF-ones comparison). We also reported standard deviations.
up match the non-HRFFs in terms of the number of FLOPS (for fair
Datasets
HRF-ort
HRF-id
ORF-reg
ORF-Had
Halton-R
Wine
0.70 0.08
0.85
0.05
1.00 + 0.04
1.10 + 0.06
4.02 0.1
Boston
0.72- + 0.06
0.79 + 0.08
1.05 + 0.03
1.14 +14+0.02
2.53 + 0.08
4.2 LANGUAGE MODELING
In this section, we apply HRFs to the language modeling task, training 2-layer LSTM with hidden size h = 200 and applying RFs for softmax sampling in training as described in (Rawat et al., 2019). Experimental results are obtained over 10 runs. Experimental details and additional validation re- sults are in the Appendix (Sec. Table 4..
3300
Le-iettersteiin
0.50
KS stotretics
0.2
2000
0.2%
Negative
0.1
2300
64
128
256
512
64
12.8
256
512
64
12.8
256
512
Number randem features
Number random features
Number
random features
FAVOR+
Trig.
Geuasien: hybrid
Anguler lrylrrdd
Figure 3: Statistical metrics measuring softmax matrix approximation quality on PennTree Bank. For standard estimators, the number of random features are 64, 128, 256, 512. To make fair comparison. for the hybrid variants, the configurations leading to the similar number of FLOPS operations per random feature map creation were applied. Negative fractions were not reported for FAVOR+ since by definition they are equal to zero.
Statistical Metrics, We compute 1d- Wasserstein distance and the Saaiooorro--Smimnoo (KS) met- ric (Ramdas al. 2017) between a true softmax distribution induced by query on all the classes
8

---

Published as a conference paper at ICLR 2022
and the approximate one given by different RF-mechanisms. We also compute the fraction of all probability estimates with undesired negative values. The results on the PennTree Bank Marcus et al. (1993) are presented in Fig. 3 In the Appendix (Sec. ! we present additional results for the WikiText22 dataset. We see that HRFs lead to most accurate softmax distribution estimation.
4.3 TRAINING SPEECH MODELS WITH CERF--OOOOOMRR-EERREEE
We also tested HRFs on speech models with LibriSpeech ASR corpus (Panayotov et al., 2015) We put implicit Performer's attention into 17-layer Conformer- Transsrurrr encoder (Gulati et al., 2020j and compared softmax kernel estimators in terms of word error rate (WER) metric, commonly used to evaluate speech models. We tested HRFs using angular hybrid variant well as those clustering-based For the latter ones, the clusters were created according to the K-means algorithm after first 1000 steps of training (and were frozen afterwards) and corresponding matrices A were constructed to minimize the loss given in Equation 21 The results are presented in Table 2 HRFs produce smallest WER models and the ceatterinn--ascer variants fully exercising the most general formulation on HRFs (with general complex estimators) turn out to be the best. Additional details regarding speech experiments as well as additional experiments with clustering- -based HRFs are presented in the Appendix (Sec. and Sec. H.I respectiveey.
Table 2: Comparison of WERs of Conformer-Transdueee applying different RF-mechanisms for the implicit attention. For methods other than clustering- HRFs (HRF-C), numbers next to method names define the values of m or (m,n), Method HRF-A stands for the angular hybrid variant. Numbers next to HRF-C correspond to the number of clusters constructed in the query and key space respectively. HRF-C uses 64
random features. We also report standard deviations averaged over 10 different training runs.
HRF-C(3,3)
HRF-C(2, 3)
HRF-C(3,2)
HRF-C(2,2)
HRF68
WER
1.72 + 0.02%
1.75 = 0.03%
1.83 + 0.03%
1.85 + 0.04%
2.03 + 0.08%
HRF-A(8, 8)
FAVOR+ 432
FAVOR+ 256
Trig 432
Trig 256
WER
2.05 1 0.05%
2.65 1 0.06%
2.77 1 0.04%
3.12 1 0.05%
3.3 1 0.06%
4.4 DOWNSTREAM ROBOTICS EXPERIMENTS
In Robotics, we leverage the accuracy of HRF's to obtain inference improvements, critical for on- robot deployment. To abstract from a particular hardware characteristic, we measure it in the number of FLOPS. We conduct two experiments targeting: (a) quadruped locomotion and (b) robotic-arm manipulation. In the former, HRFs are applied as a replacement of the default mechanism using positive random features in the class of implicit-attention architectures for vision processing called Implicit Attention Policies (or IAPs) (Choromanski et al.) 2021a). In the latter, HRFs become a part of the regular results are presented in Fig. 4 where HRFs need 3x fewer FLOPS to obtain same quality policies as baselines. All details regarding both experimental setups are given in the Appendix (Sec. K)
stack used to process high-resolution (500 x 500 3) input. The
Videos of the HRF- trained robotic policies are included in the supplementary material.
IAP reg 256
FAVOR+ 256 AA4AA
FAVOR+ 128
100
IAP HRF 8x8
Hybrid Bx8
FAVOR+ 64
400
Trig 256
302
Trig 128
Trig
100
-100
250
500
750
1000
1250
1530
1750
2000
0.0
0.2
0.4 0.6 0.6
ES Iterations
Iterations
Figure 4: Left: Step-stone locomotion task. Comparison of the training curves for: the IAP using angular hybrid estimator of m = = 8 and IAP applying regular FAVOR+ mechanism from (Choromanski et al. 2021b) with m = 256. Both final policies are of similar quality, yet HRF-method requires 3x+ fewer FLOPS run its trained policy. The visualization of the HRF policy in action and its attention (with filtered out pixels marked in red) in the bottom right corner. Right: Similar setting (and conclusions) but for the robotic-arm manipulation task. The additional five regular RF-configurations did not train by producing Nan loss due to large variance of the underlying softmax kernel estimators.
CONCLUSION
We presented a new class of random feature techniques, called hybrid random features for soft- max/Gaussian kernel estimation that more accurately approximate these kernels than previous al- gorithms. We also demonstrated their robustness in a wide range of applications from softmax sampling to Transformers/attention training, also for downstream Robotics tasks.

---

Published as a conference paper at ICLR 2022
Reproducibility Statement: Section 4 and the Appendix include all experimental details (e.g. hyperparameter setup) needed to reproduce all the results presented in the paper. The part of the code that we could make publicly available be found in the following github: lattps://githtn
com/HI.-han in /RRF_/FFFF 1n/HRF_TCHR2222222
HRFs can be used in principle to train massive Transformer models with large number of parameters and proportional compute resources. Thus they should be used responsibly, in particular given CO2 emission challenges that scientific community tries to address now.
Ethics Statement:
ACKNOWLEDGEMENTS
AW acknowledges support from Turing AI Fellowship under grant EP/V025379/1, The Alan Tur- ing Institute, and the Leverhulme Trust via CFI. AS acknowledges AWS compute resources from Onur Kara.
REFERENCES
Unitree Robotics. URL htt.htt:://www.nr http::/www.niitree.commm cc/
Haim Avron, Vikas Sindhwani, Jiyan Yang, and Michael W. Mahoney. Quasi-monte carlo feature maps for shift-invariant kernels. J. Mach. Learn. Res., 17:120:1-120:38, 2016. URL http:
7/ //imlr. wwtppppppss/v//////a.
html
Guy Blanc and Steffen Rendle. Adaptive sampled softmax with kernel based sampling. Jen- nifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference Machine Learning, ICML 2018, Stockholmsmüssan, Stockholm, Sweden, July 10-15, 2018, vol- 80 of Proceedings of Machine Learning Research, pp. 589-598. PMLR, 2018. URL http:/ http:://ptoceeeeeeeeee http://prceeee...cccccc .mm.rdiiinsmmmm............m.. m.lr. www.m.lrrrrrss..mmmmmmmmmmmms wwwps/vee00 www.baannaaa.cc...coeaaaalleeaae
Y. Cho and L. K. Saul. Analysis and extension of are-cosine kernels for large margin classification. Technical Report CS2012-0972, Department of Computer Science and Engineering, University of California, San Diego., 2012.
Krzysztof Choromanski and Vikas Sindhwani. Recycling randomness with structure for sublinear time kernel expansions. In Maria-Floring Balcan and Kilian Q. Weinberger (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICML. 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings. pp. 2502-2510. JMLR.org, 2016. URL http://proceecings wetpeeeeiiinns.m......mm. html
www. wellr.pres//aee www.pesss/v.... v48 htt.s//chooomaaaaiaiiccaaa m
Krzysztof Choromanski. Mark Rowland, Wenyu Chen, and Adrian Weller. Unifying orthog- onal Monte Carlo methods. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Pro- ceedings of the 36th International Conference on Machine Learning, ICMI. 2019, 9-15 June 2019, Long Beach, California, USA, volume of Proceedings of Machine Learning Re- search, pp. 1202-1212. PMLR, 2019. URL choromanski html
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, David Belanger, Lucy Colwell, and Adrian Weller. Masked language modeling for proteins via linearly scalable long-context transformers. arXiv preprint 035555,006.03555 2020.
Krzysztof Choromanski, Deepali Jain, Jack Parker-Holder, Xingyou Song, Valerii Likhosherstov, Anirban Santara, Aldo Pacchiano, Yunhao Tang, and Adrian Weller. Unlocking pixels for inforcement learning via implicit attention. CoRR, abs/2102.04353, 2021a. URI. https:
//axxiv.org/abb//21000 04353
Krzysztof Marcin Choromanski, Mark Rowland, and Adrian Weller. The unreasonable effective- ness of structured random orthogonal embeddings. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus. S. V. N. Vishwanathan, and Roman Gamett (eds.). Ad- vances Neural Information Processing Systems 30: Anmual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pp. 219-228. 2017.
10

---

Published as a conference paper at ICLR 2022
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with per- formers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021b. URL Utt.httts:///ppeereeee
Amit Daniely, Roy Frostig, Vineet Gupta, and Yoram Singer. Random features for compositional kernels. CoRR, abs/1703.07872. 2017. URL httt:::/aaxxv.orgg//oo
www.av.orr/gaaa wwt.:oggg//aaaooo 07872
Michel X. Goemans and David P. Williamson. Approximation algorithms for MAX-3-CUT and other problems via complex semidefinite programming. J. Comput. Syvl. Sci., 68(2):442-470. 2004. doi: jo.s
URL
2003. 02.3.77.022 0121
Arthur Gretton, Ralf Herbrich, Alexander J. Smola, Olivier Bousquet, and Bernhard Schölkopf. Kernel methods for measuring independence. J. Mach. Learn. Res., 6:2075 2129, 2005. URL http://
jm1 wrg. http://ppprss////g/peeeeeee
html
Anmol Gulati, James Qin, Chung- Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, and Ruoming Pang. Conformer: Convolution-augmeente transformer for speech recognition. In Helen Meng, Bo Xu, and Thomas Fang Zheng (eds.), Interspeech 2020, 21st Annual Conference of the International Speech Communication Associ- ation, Virtual Event, Shanghai, China, 25-29 October 2020. pp. 5036-5040. ISCA, 2020. doi: N0.21437/Intersppech.2222 2020-301.. URL https: / doi org 10 21 21437/ httt:///.....3../teeeeeeep..eee. Interspecch. 2010-3015
Hakan Inan, Khashayar Khosravi, and Richard Socher. Tying word vectors and word classifiers: A loss framework for language modeling. In Proceedings of the 5th International Conference Learning Representations, 2017.
Fanghui Liu, Xiaolin Huang, Yudong Chen, and Johan A. K. Suykens. Random features for kernel approximation: survey in algorithms, theory, and beyond. CoRR, abs/2004. 2020. URL
https::/arxxt org/ abs 2004. 11154
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. Advances in Neural Information Processing Systems, 34, 2021.
Mitchell P. Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building a large annotated corpus of English: The Penn Treebank. Computational Linguisties, 19(2):313330, 1993. URL
logy .rg. J93-2001
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer Sentinel Mixture Models. 5th International Conference on Learning Representations, abs/ htss/660.00843, 2017.
Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: An ASR corpus based public domain audio books. In 2015 IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2015, South Brishane, Queensland, Australia, April 19- 24, 2015, pp. 5206-5210. IEEE, 2015. doi: N0.1009CCASS222055777785
URL https:
10. 1ww.1/11110........ 2015. 7178964
Hao Peng, Nikolaos Pappas. Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention. 9th International Conference on Learning Represen- tations. ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL httpps:/openreview net atte:/oooummm---tteaaaaaaa-
Ofir Press and Lior Wolf. Using the Output Embedding to Improve Language Models, 2017.
Ali Rahimi and Benjamin Recht. Random features for large-scale kemnel machines. In John C. Platt, Daphne Koller, Yoram Singer, and Sam T. Roweis (eds.), Advances in Neural Infor- mation Processing Systems 20, Proceedings of the Twenty-Firs Annual Conference on Neural Information Processing Systems, Vancouver, British Columbia, Canada, December 3-6, 2007,
11

---

Published as a conference paper at ICLR 2022
pp. 1177-1184. Curran Associates, Inc., 2007. URL http:://appprr// www.rr...iii...... ct.pp.cc//ppp... 1182-randee features hate-aaargggaalll--eeaaaaaaaaaia kernel-machiner
scale-
Aaditya Ramdas, Nicolás García Trillos, and Marco Cuturi. On Wasserstein two-sample testing and 2017. doi: 10.3390/e URI.
related families of nonparametric tests. Entropy, 19(2):47, 2017.
https::/ddf .org/10........ .3900/39000000 e 9020047
Ankit Singh Rawat, Jiecao Chen, Felix X. Yu, Ananda Theertha Suresh, and Sanjiv Kumar. Sampled softmax with random Fourier features. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' d'Alché-Buc, Emily B. Fox, and Roman Garnett (eds.), Ad- vances in Neural Information Processing Systems 32: Anmual Conference on Neural Informa- tion Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, cctpppprep/119////////
pp. 13834-13844, 2019. URL https::/proceddd www.ggg.................. wwt.:iiww....ppeeeeeeeeeeeeeeeee
Mark Rowland, Krzysztof Choromanski, François Chalus, Aldo Pacchiano, Tamás Sarlós, Richard E. Turner, and Adrian Weller. Geometrically coupled Monte Carlo sampling. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Ro- man Gameet (eds.), Advances in Neural Information Processing Systems 31: Annual Confer- ence on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, nww.ips.cc... cc/ Montréal, Canada, pp. 195-205, 2018 URL oas html httpr//005/aaatte3333333333/3333
Abst
Alessandro Rudi, Raffaello Camoriano, and Lorenzo Rosasco. Less is more: Nyström computa- tional regularization. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 28: Annual Con- ference Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Que- bec, Canada, pp. 1657-1665, 2015. URL htt.ttttpstt.
http::::///occeeeeeeoe neea..neriii ps. cC cc/paper/ htm
Iaaol 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 Interna- tional Conference on Machine Learning. ICML 2021, 18-24 July 2021, Virtual Event, vol- ume 139 of Proceedings of Machine Learning Research, pp. 9355-9366 PMLR, 2021. URL http://proceedin
mlr. wwwss/ httpsss/vvv/////// hww.////chhhlll......
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. CoRR, abs/1 807.0070000 507.03748, 2018. URL httt:://p///iiii wwwp://wworrrooooooooo abs 807.807.00 03748
Christopher K. I. Williams and Matthias W. Seeger. Using the Nyström method to speed up kernel machines. In Todd K. Leen, Thomas G. Dietterich, and Volker Tresp (eds.), Advances in Neural Information Processing Systems 13, Papers from Neu- ral Information Processing Systems (NIPS) 2000, Denver, CO, USA pp. 682-688. neurips.. httpppps..///apee///////ssss Abstract. html
MIT Press, 2000. URL Atstrrraafaaaaeeeetatttttt.....
Tianbao Yang, Yu-Feng Li, Mehrdad Mahdavi, Rong Jin, and Zhi-Hua Zhou. Nyström method vs random Fourier features: A theoretical and empirical comparison. In Peter L. Bartlett, Fer- nando C. N. Pereira, Christopher J. C. Burges, Léon Bottou, and Kilian Q. Weinberger (eds.), Advances in Neural Information Processing Systems 25: 26th Annual Conference Neural Infor- mation Processing Systems 20 Proceedings of a meeting held December 3-6, 2012, Lake Tahoe, Nevada, United States, pp. 485 493, 2012. URL http:: htt.ttps:/pproceeeeeeee ngs newrips. cct httpa//00....aa22.a/aa//22222220
Felix X. Yu, Ananda Theertha Suresh, Krzysztof Marcin Choromanski, Daniel N. Holtmann-Rice and Sanjiv Kumar. Orthogonal random features. In Daniel D. Lee, Masashi Sugiyama, Ulrike von Luxburg, Isabelle Guyon, and Roman Garnett (eds.). Advances in Neural Information Processing Systems 29: Annual Conference Neural Information Processing Systems 2016, 5-100, 2016, Barcelona. Spain, pp. 1975-1983, 2016. URL http: papers nips cc cc/paper/ 6246 6 1al4-oorthoooal-ranngnoo features
12

---

Published as a conference paper at ICLR 2022
Qinyi Zhang, Sarah Filippi, Arthur Gretton, and Dino Sejdinovic. Large-scale kernel methods for independence testing. Stat. Comput. 28(1):113-330, 2018. doi: 10.1007/41222--11--777777 /0tp:///////00000////22222011110
13

---

Published as a conference paper at ICLR 2022
APPENDIX: HYBRID RANDOM FEATURES
A PROOF OF LEMMA 2.4
Proof. From the independence of w1, .... we get:
I [eppww> x)] II E [(xx((w;2)]]
(28)
Thus it suffices to show that for NN(,,) 1) and any E C the following holds:
22 )
Define f(2) = Efexp(gz)]- Note that f(z) exp( for ir, where 1 = -1 and x R which follows from the formula of the characteristic function of the Gaussian distribution. Similarly, = =
(29)
f()) = exp( =²2 for E IR which follows from the formula of the moment generating function of the Gaussian distribution. We now use the fact from complex analysis that if two analytic functions + C are identical on uncountably many points then they are equal on C to complete the proof (we leave to the reader checking that both f(2) and g(2) def exp( + are analytic).
B PROOF OF LEMMA 2.3
Proof. The following is true:
hyb SMm,((., y)
ak 1.k (i,m (x) * ... o(tnk(x)) (02 (y) (y) *... oltn (y))+
m
(P (x) @ of (x)) ( (y) @ i.k (y))+
k=
(30)
Σh
-k=1
0K
1.pt1 111m
(x)
*
Pi.
(x)) (x)
(02m
(ol.tt1( (y) *.. * 22.m 1 (y))
m
ΣΣ
(x) (x) s ((pp1(())) (p2)) (y) & diptt (y))
The statement of the lemma follows directly from that. Note that the key trick is an observation that expectation (i.e. rewritten with and
product of two functions of x, y that can be linearized disentangled) is itself function that can be lincarized on expectation. In the setting where the former are approximated by random features, the latter is approximated by their cartesian product (the random feature mechanism for the compositional kernels from (Daniely et al. 2017
C BIPOLAR HRF ESTIMATORS: PROOF OF THEOREM 3.1
Proof. The following holds:
Var(SM byb (x, Var (O)SM ++ (x,)) Var ((1 - (0) (0) SM hris (x. y)
(31)
trig 2Cov (O)SM, (x,y), 2Z, (x,
We will focus now on the covariance term. To simplify the notation, we will drop index n in X --trrg Assume first easier to analyze case where SM,, (xc y) and SM (x, y) are chosen independently.
-
Then the following is true:
14

---

Published as a conference paper at ICLR 2022
Cov (A(0)SM (x, (1 ((O)) ii (x. (x,))
B[²(0( - S( (()))SMM + (x, y) (x. - E(² (SM (x, - ((0)) trig (x,
=(B((()((1 - ((0))] - [[A - X(011)) BJSM,. + ( (x,
y) (((()(())
(32)
Arig Now assume that SMm (x, y) and SM, y) use the exact same random projections. Then, using similar analysis before, we get:
++
Cov (X(O)SM (x,)), (1 ((0))SM trig (x. (x.)s)Mm(((x)))))))))))
(33)
tris E|SM (x. S) (x. (EE
(SM((. (E( -1(01)
This time however it is no longer the case that E'SS,, (x, y)SM m (x, y [E()()) X()))] =
E[SM,, m (x, ESS, trig, (x, = (SM(x.y))* since SM,, (x, and SM, m (x, y) are no longer
trip,
{(0))). we will first trip (x, E()( independent. In order to compute E[SM,,, m (x, YSM, introduce useful denotation.
Denote by W1, üd (N Ia) the random projections sampled to construct both SM, (x,y) yx, Denote: Y; = cws (x + y)) def exp(w]
trig, and SM,, m
(x+y)) We have:
2
|x|| + 1y
Yi +... Y..
SM m (x.y) = exp
(34)
2
m
trig (x, y) as:
If we denote: Z, = cos(w, (x - -y),, then we can write SM,,
trig (x, = exp
|xx| + |1y1
21 + ... + Zm
SMm
(35)
2
We can then rewrite RSSM m (x, y YSM trig yx as:
E SMM
y)MM, trig (x,
1
Σ 2|, + Σ E|YZZ, ifj
m2
(36)
(m) m 2
(SM(X, y))² mmE[cosh(u T (x + y)) cos(w T (x - y)]
m
trig
where w ~ (0 I.). The equality follows from the unbiasedness of SM,, yx. and SM,, (x, v) and the fact that different w; are chosen independently. Thus remains to compute p B/coss(w" (x + y)) cos(w (x - y)]]. Note first that = R[exp(w> T (x + y)) cos(w (x y)) since -w (0, Ia) and cos an even function. Denote z = x + y + i(x y). We have:
Elexp(w T (x + y)) cos(w> T (x - y))] = Re [E(exp(w 2)]] Re
exp 2 =
d Σz= (ry + 2, 21(x² - w3) (xr, -,,)² ] -SMM(x cos(||x11 - 11x133)
Re
exp
2
(37)
Thus we conclude that:
+ E(SM, m (x, S)M, (x,y)) =
trig
1 101 -11 (SM(x))) + (SM(x, y))² cos(1|x11 - 11) (38)
m
15

---

Published as a conference paper at ICLR 2022
Therefore get the formulae for the covariance term in both: the setting where random projections
of SM,, trig (x,y) and SM,, m () are shared (variant II) and when they are not (variant 1). The following is true for Z 11(11 cos([| ((x))))((((0 X(0))):
-(SM(X, ()M(xx)))Vaar(Xll for variant -(SM(x tris (x, [- (aa( + 2) for variant II
(o)(S (I-
(39)
We also have the following:
Var(X(0)SM,, = (x, -B[(O)) (x,
[[({()))] (MSE(SM, (x.y) + (SM(x, - (SMX
(40)
(SM(x, y) Var (A(O)) +
(x,y))
and furthermore (by the analogous analysis):
Var((1 (()))M trig (x,
(EM(x, -X
trig (x. y)) (41)
By putting the derived formula for the above variance terms as well covariance terms back in the Equation 31 we complete the proof of the theorem (note that the mean squared error of the hybrid estimator is its variance since it is unbiased).
D ANGULAR HYBRID ESTIMATORS
D.1 PROOF OF LEMMA
Proof. The formula for the expectation is directly implied by the formula (1) from Sec. 2.1 oh &
Saul (2012) for the zeroth-order arc-cosine kernel kn (x. def 6xy It also follows directly from
Goemans- Williamson algorithm (Goemans & Williamson 2004 We will now derive the formula for e Ex² (x, y)]. Since depends only on the angle Oxxys we will refer as: X(0) and to
its estimator as X(0). Denote:
1 2),..., sgn(T,]]
gang (2) Vn
(42)
Denote: X;
[i]. We have:
( 1 Σ x
i()) 1
(43)
2
Note first that by the construction of 7((), have: B[x(O)] Therefore we conclude that:
and thus: E
X,]
20
27
20 +
==
E
1
1-2
2
X
X;
=
1 - 2(1
B[X +
Σ
EXX,,ELX
4
4
1=1
i=1
i=1
20
20
20
2(1
+ n
+
n
1)
(1
= 4
4
+ (1 )(
4
n
n2
0
(44)
16

---

Published as a conference paper at ICLR 2022
D.2 EXPLICIT FORMULA FOR THE MSE OF THE ANGULAR HYBRID ESTIMATOR
Theorem 3.1 and and Lemma 3.2 immediately imply the following result.
Theorem D.1 (MSE of the angular hybrid estimator). Take the angular hybrid estimator
tri anghyb (),,), where SM, yx, and SM (x.y) are chosen independently i.e. their random
SM
projections chosen independently (note that we always assume that y(x,)) is chosen indepen- tris dently from SMm (x,y) y) and m
++
(x,y)
Then the following holds:
anghyb m.n
0
0
0
1
1
MSE(SM
(x)
exp([|2|
()0[)=/)))MM((x)))( exp((11213))
nm
2m
0
0
0
1
1
(1X (45)
π
TU
2m
mrg
for A = y and z x +
y. Farthermore, if SM y) and SMm, yx apply the exact same sets of random projections, the mean squared error of the hybrid estimator is further reduced, namely we have:
anghyb
0
0
1
1
MSE(SM) m.
x))
+
-
exp(11z]
11/
n2
2m
0
0
1
1
n
2m
2
(1
1
0
SM2 y1 c
(1xl)))
+
117
(46)
Thus if ||x||2 = ||yl| = r then regardless of whether the same sets of random projections are applied or not, we get:
anghyb
0
0
exp(8r2 cos?( 0 -2
MSE SM m.n
(x,y))
=
+
π
2m
G))e
0 1
0
exp( 2n² sin 2 0 ( ())
(1 - exp( exp cos?
+
+
exp exp(2r2)
π n
n
nrT
2m
2
(47)
D.3 PROOF OF THEOREM 3.5
Proof. Note first that from the derived above formula of the MSE of the lambda-angular bipolar hybrid estimator and the definition of the max-relative-error, we obtain:
anghyb
exp(r2)
(s(r) SSM
m,n
V2m
max (o), (E(0.N)
(48)
where:
h-((0) = (0) am 8)
(49)
and a, (0) is defined as:
a
0
+ 2 n
exp(2r2 cos(O))
(1- expepp(-1)² cos? (
(50)
π
π
no
Therefore we have:
(r()) SMM (Mmma
anghyb
exp(r2)
max (#) @E(U,T]
(51)
vm
Notice that:
(0) < (,)( exp( ex)(-4r2))2
(52)
17

---

Published as a conference paper at ICLR 2022
where:
(0) = 12(0) + 02(00
(53)
and
(0) = (1 0²
exp(2r² cos(
(54)
π
0,) n π
exp(2r cos(
(55)
Therefore:
max (0) z max (0) @
max (0)
(56)
# [0
0 (0.
Denote:
b² = s0 (0) and 82 = ma [0 (2) Note that:
(0) exp cos(
- 20
sin)
(57)
d0
n 2
and
(0) = exp 2n (0)) T2IT
(1 27 sin(
(58)
d0
Thus, from the properties of function: 0 + θ (f) and the fact that 2 1, we conclude that both derivatives are first non non-negatie -engtvvae,ee, then non-positive and then non-negative and that the unique local maximum on the interval [0, T] is achieved for 0 I 3. Note also that (0). ()) 2 0 and (Tr)
exp(-2r")) We conclude that global maximum for by. on the interval [o, >] for i = 1,2 is achieved either in its unique local maximum on that interval or for и. Let us consider first: b/. In its local maximum on [0, π] we have:
67)(0) (0) = 0, (T) (1
cexp(-2r**
1 p2
M sin
(59)
Since 0 S sin( 3 on [0, Z), we get:
(0¹)² < क 1
(60)
2 p2'
i.e.:
1
0r S V r
(61)
Therefore:
1 1
(00) - (1
exp(2r2) 2 (T)
(62)
2
We thus conclude that:
max by (0) - (1 ) 6E|0,=]
2TT² exp(2r?)
(63)
By the completely analogous analysis applied to 63. we obtain:
max (0) S 2n 2Tr2 exp(2r?)
(64)
# [0
Now, using Equation 51 Equation
52 and Equation 56 we obtain:
exp(2r2)
(tr) (SM anghyb
(1 (xp(-447))
(65)
V2mr
n/a
and that completes the first part of the proof (proof of Inequality 27 directly implied by the fact that:
The equations on the limits
(s(r)) SMM anghyb
)= exp((222
(O)
(66)
22m
18

---

Published as a conference paper at ICLR 2022
D.4 DISCUSSION OF THEOREM 3.5
Theorem 3.3 leads to yet another important conclusion not discussed in the main body of the paper. 0 + 0 0 + n, the relative error decreases (as function of m _- (Z...) (ess) ), where dest stands for the number of random features of the resulting
Note that asymptotically and 12) as Vim
estimator. Note also that for the regular estimator the rate of decrease is also O ( (here we treat all other arguments of the formula for the relative error as constants; we already know that the dependence them of the relative error for the HRF estimators is superior to the one for regular estimators). The key difference from the computational point view is that for the HRF estimator, the random feature map be constructed in time Ond + md mn), whereas for the regular estimator it requires time ((mdd) (see: Sec. [). Thus for the regular estimator random feature map computation is much more expensive. An important application, where random feature map computations takes substantial time of the overall compute time are implicit-attention Transformers, such as Performers (Choromanski et al. 2021 This is also the setting, where an overwhelming fraction of the approximate softmax kernel values will be extreme (very small or large) since an overwhelming fraction of the attention matrix entries after some training will have extreme values (very small or large). In the setting, where the lengths of queries and keys do not vary much (for instance prominent case where they all have the same fixed length) that corresponds to angle values 0 or 0 + O. This is exactly the scenario from the second part of Theorem 3.5
E PROOF OF LEMMA 3.4
Notice that the relative error es,- (SM trig, is an increasing function of sin? sin?(g) and thus is largest for Plugging in this value into the formula of the relative error gives us the expression from the statement of the lemma. Completely analogous analysis holds for SM,, m
Proof. Equation 25
follows directly from Lemma 2.2
F GAUSSIAN HYBRID ESTIMATORS
In this section we will provide more results regarding Gaussian hybrid estimators. We will focus on the bipolar scenario and within this scenario the so-called normalized variant described below.
If x and y have fixed L2-norm ([|x||2 = ||y||a = r), we can propose a normalized bipolar Gaussian hybrid estimator such that x(x,y) 0 if 0x.y and A(x,y) 1 if 0x.y ग. Furthermore
the variance of that estimator, that we call shortly SM will be zeroed out for Ox.y = 0 or xx.y but not for both.
The coefficient-function A: Rdd x Rdd - IR is defined as:
x(x,y) = exp( 5" ||x y11
(67)
where p is given as:
= 1 20)2
(68)
The hyperparameter a controls the smoothness of the Gaussian kernel.
The exponential in can be estimated either by trigonometric or positive random features. It is easy to see that for the fixed input lengths. in the former case the variance of the estimator is zero for fx, 0 and in the latter it is zero for 6x.y = n, but the variance never zeroes out for both 0x. 0 and fxy = x. In principle, one can derive entire hierarchy of HRF estimators, where the coefficients in an HRF estimator from one level of hierarchy are estimated with the use of HRF estimators from the higher level. In that context, angular hybrid estimator can be applied to provide approximation of the coefficients of the normalized bipolar Gaussian hybrid estimator leading to variance zeroing out for both xxy 0 and 6x,y
However such nested constructions not the topic of this paper. From now on we will assume that trigonometric random features are applied to estimate X-coefficient, Therefore have:
19

---

Published as a conference paper at ICLR 2022
x(x,y)
(os(o (x- y)))
(69)
np
and that leads to the following choice of: a, P1,2, p2,n from Equation 10
{ - ())(z)= p2,n = vmp (sin(T) z), cos(T] z., sin(T,, 2), 2))
a
(70)
where Ti Tn NN Ia) and 1 -1
Using Theorem 3.1 we conclude that:
Theorem E1 (MSE of the normalized bipolar Gaussian estimator). Take the normalized bipolar trig Gaussian estimator SM kausshyb y), where SM, m yx and SM,, ++ (x are chosen indepen- dently iee. their random projections are chosen independently (note that we always assume that
A(x,y y) is chosen independently from SM, (x,y) and SM m the following holds:
(x, Denote: A x y. Then
MSE(SMM)) m.
gausshyb
(x,y) = M/SSSSMMM
(SM,,
crig
(x.y)) (71)
where
0² E)({(x,))) Tsl (1-epp( 1xx((==1111))
1
exp(
-1- -x-epp-oo"
(72)
12
2p
1 [[( -(xxx)))) p2 ((1 1eeppp(- -xp((-> 1)))) 1
0²
2p2n (1
exp( -xp(-
(73)
Furthermore, SM trig (x. y) and SM,,
m (x, y) apply the exact same sets of random projections, the mean squared error of the hybrid estimator remains the same.
Proof. We can calculate E[({(x,))) and B[(1 X(x, as follows:
[[[[ - y)) 5 Bl(p 1+
cos(ow T (x y)
(1-p)²
2(p
E Blcos(ow) T (x - yy)] + 122
E1( cos(ow T (x - y)))
12
(1 - p)²
2(p - 1)
a²
exp(
||) + 1
1xsp[-o211: exp(-o?|1A13 +
2p2n (1 - exp( - ( 1))
1
02
2
(11 - p) - -expl 0² ||(2))2 + 2p
(1 - exp( 0² ||| )2)
=
p
(74)
[[[[(x, v)) = 2E[1
7(x, + B[[1 X(x,
2_
2
0²
=-1+
exp(
| + E[(1 7(x, y)
(75)
P
1 0²
- (1 -xp((-
exp(- o2 -2|| 12))
20
By plugging in these two formulae in the expression for MSE, obtain the first half of the theorem.
20

---

Published as a conference paper at ICLR 2022
trig Note that if SM,, m (0.r) and SM, r0,r) apply the exact same sets of random projections, then kausshyb yx) becomes zero in Theorem F.1 therefore the
(1 c - ll)) in MSF(SM (a m.s MSE remains the same. We thus obtain second part of the theorem and that completes the proof.
Now let us assume the setting of the same-lengths inputs. We denote the length by r and an angle between inputs by 0. We can rewrite Eq. 67 and provide equivalent definition for À by replacing and y with their angle and norm r:
X(0,r) 1- exp(-222 20-0-22222 sin( )))
(76)
We obtain the following version of the above theorem:
Theorem F.2 (MSE of the normalized bipolar Gaussian hybrid estimator for same-length inputs). lyyll r and denote: 6 0xx- Take the normalized bipolar Gaussian gausehyb Assume that ||x|| =
++ (0,r) are chosen independently hybrid estimator SM m. (0,r), where SM, (0,r) and SMm
.rig
i.e. their random projections are chosen independently (note that we always assume that X(0, r)
chosen independently from SM, (0,r) and SM ++ (0, r)). Then the following holds:
Arig
MSE(SM (SMMaussyy
gausshyb
(0,r)) (0, MSEM ++ (0,r)) + B[(1 - r)) (MMS,EM,, trig
(0,r))
(77)
where
[[(A(0,r)))] = - exp( expp-22021= sin 2 (2)))2 + 1²
8
1
(1 --eeep-(o00- sin })) (78)
2p2n
E[(1- (((1)x(0))))] p² -expp) exp(-2o2p sin 2 )2)
1
1
exl-epp(-4o') 22 sin sin(
)3)) (79)
20nn
trig (0, r) and SMm () apply the exact same sets of random projections, the mean squared error of the hybrid estimator will remains the same.
Furthermore, if (SM
Proof. For |x|| ||yl2 = r2, the following holds for xx--v y:
0 g
|| 4r sin
(80)
And this theorem holds trivially by replacing || 2 with the above formula in Theorem F.1
G POINTWISE SOFTMAX KERNEL ESTIMATION EXPERIMENTS
In Fig. we present complete ablation studies regarding empirical relative errors of different RF- based estimators of the softmax kernel over different lengths of inputs r and angles 0x.y-
H COMPLEX EXPONENTIAL ESTIMATORS FOR CLUSTERED DATA
H.1 SYNTHETIC EXPERIMENTS ON DATA ADMITTING CLUSTERING STRUCTURE
Here we test HRF estimators customized to data with clustering structure, as described in Sec.
2.4 Inputs x are taken from two 50-dimensional 1000-clement Gaussian clusters and the inputs y from two other 50-dimensional 1000-element Gaussian clusters. We use empirical MSE metric and constrain matrix A to be real diagonal for a compact closed-form formulae (see: Sec. H.3). We
created four different clusters' configurations corresponding to different values of [(A) (see: Eq. 21) for all pairs of clusters and with different concentrations around centers controlled by the standard deviation (small/large values of and σ). As shown in Fig. 6 for all variants HRF estimators adapted to clustered data consistently provide notable accuracy gains, even for larger values of and less concentrated clusters. Additional experimental details given in Sec. H.
21

---

Published as a conference paper at ICLR 2022
Trigonometric
+
Angular Hybrid
35%
0.4
Relative Emor
5%
0.3
SW
erfor
0.2
0.1
Relative Error
1.2
0.8
2-
0.0
15
6.0
Relative Error
4.5
3.0
1.5
0
1d2
0
nd2 9
0
1d2
6
0
Figure 5: Pointwise estimation of SM(x, for the same-length 64-dim inputs (r = 1.0, 1.25 and r
1.5) and various angles 8xyy. We used = 10000 estimated softmax values in each subplot. The true softmax value and the 5th and 95th quantile estimated values are shown by the left y-axis, and the empirical relative errors are shown by the right y-axis. Trigonometric estimator and FAVOR+ applied 128 random features. To make fair comparison, for the hyhrid variant the configuration leading the similar number of FLOOP operations per random feature map creation was applied.
;-|0.08, 0.05, 0.04, 0
s=(0.11, (.00.,..00.000 0.08, 0.07].
5- 10.005 0.03, 0.03, 0.02) oF=10
5-10.07. 0.06, 0.05, 0.04),
0.05
Hyo id clitter FR/CR - 119
0.06
5.06-
0.05
0.00
5.04
0.01
0.00
0.02
0.02
5.22
0.01
0.01
0.00
0.00
5.00
50
10C
150
32D
52
100
150
202
50
-20
350
200
50
100
150
200
Number rendom Pratures
Namber random features
Numbbe randorr
Namberr moddmm
Figure 6: Comparing different estimators for data with clustering structure. The empirical MSE obtained by averaging over 20 trials. For the HRFs, we apply deterministic X-coefficients. The fact that the empirical error for FAVOR+ is not perfectly monotonic in m was first observed in (Luo et al. 2021) (see: Fig. 1: (b)).
H.2 INSTANTIATING COMPLEX EXPONENTIAL ESTIMATORS FOR CLUSTERED DATA
Assume that the inputs x (queries) can be modeled by ng clusters and that the inputs y (keys) can be modeled by ny clusters (e.g. via k-means clustering algorithm). Denote the center of each cluster as E Rdd (i 1, .... R& and rj E Rd. (j = 1,..., ng)) Then there exist 72,77 pairs of ry), 1, J I..; ng),, so we can construct n, softmax kernel estimators to estimate cross- group softmax kernel values.
Consider z Ar, + (AT)-Ir, for an invertible (in Cdxd) matrix A Cdxd An estimator based on this mechanism has variance equal to 0 if:
Ar, + (AT) r =
(81)
From now constrain A be diagonal,
A AT. We can rewrite the above equation
z Ar, + (A) -rr =
(82)
Since position (ke, k) in matrix A is complex number of the form ag + Bai, we need to satisfy the following equation for each & = 1,....d d:
(ax + Bki)raa
OKk + 3xi JiA
(83)
+ Og 3zi
(ax + BLi)rik
a 82 j.k 0
(84)
22

---

Published as a conference paper at ICLR 2022
where ri is the k-th entry of vector ry.
We can simplify this equation and separate into real and imaginary part:
{
Re: (a + Jarri + OnTjj Im (og + BETi, - PaTj,
(85)
Our goal now is to choose values for ag and Bx given r, and ry.
If J, O, we can set:
{ Bx VTak/Tink
ak 0
(86)
And if Fe
set:
{ OKk v - Ti 3x =0
(87)
When ri,k = ry,k = we can take og 8x 1. If ri, = 0, o 0 or the opposite, then we cannot satisfy the above equation perfectly. We can take ag to some large positive value and set 3x when TiKk = 0, rj # O, and set ag = 3r to some small positive value close to zero when Ti + 0, 0, = 0.
When restricting A diag(ar,.... .1,, aa) to Rdxd, it is easily seen that to minimize |Ar; + (A ) we can take
ag V Irja//aal
(88)
if riA # 0. An riA = 0,TJA # 0, dg can be set to large positive number. We can set og- = 1 if
Ti,k = 0, -5 =0
In the analysis below we denote matrix A calculated from Eq. 85 and Eq. 88 (depending whether we are using real or complex matrices) given r, and r, as
From the analysis in the main body of the paper we know that for arbitrary vectors x, y E Rdd
||A 5.5 x exp
[I ((A+))-")l' JSM(x. y) E[exp(w T ( + (A 'y)
exp(
(89)
2
2
Therefore, the softmax kernel can be estimated as:
SM(x, x (x) ((At)) (y)
(90)
with T (x) and T (AAa)-1 (y) given by:
(x) def 1
_exp(
-exp(-
|A x
(exp(w) A A A (91)
Vm
2
[1(A#))
1112 -
m (A:i)) (A iy)) (x)
def
exp( v m
(exp(ui (( expl(, (A (92)
2
for Wi, Wm NO.LI.)
Such a softmax kernel estimator has variance equal to zero if x = r, and y = r, if we are using the complex matrix and has small variance if are using the real matrix. For other pairs of vectors (x, y) close to r, and r,, the variance is close to zero, or relatively small.
We also take the Gaussian X-coefficient estimator that gets it maximum value when r.yy ry as X y). Such an estimator could be constructed with A s in a similar way.
It can be easily linearized since:
23

---

Published as a conference paper at ICLR 2022
| (A*J)
'y11
X (x. y) = exp(
2r
(93)
A 11 Jexp(
|1(A*) -(
-ly
x -y
exp(
JSM
2r2
2r²
T'
where (alxx/T.-y/11 could be estimated with similar random feature maps given by Eq. 90 Therefore, the Gaussian X-coefficients can be estimated as:
(x) = exp(- || 1. x + (Ald)
1 T.
(x)
(Aii))
T.
(94)
27
with v Ti (x) and T.
(AA-))- (-y) given by:
1
|
T. And
(x) =
exp(
(exp(w] A expl:)
(95)
mm
x2
((AA))))))) (cxp(- (A cxp(-wu (J) y010yyyy)))
T.m
(AtJ) ( (-y)=
exp(
y/T))
m
72
(96)
for wi, Wm L4) chosen independently and that were applied in base estimators.
We summarize this section with the following two constructions of the hybrid estimators for the clustered data that are implied by the above analysis.
Definition H.1 (Hybrid Coussian-Mixtures estimators for clustered data). Assume that inputs (queries) and y (keys) can be modeled by ny and n& clusters respectively with centers n, and r,
RR, (i = -I,..,, j = Bg). We denote A A.- as the complex matrix satisfying Eq. 85
with center r., Furthermore, we denote E (SM -id the softmax kernel SM and A = y) as list of estimators of (x, constructed independently from E. We also or SM (x) and denote the estimator of its softmax kernel as SM yx, and X-coefficient as
(x.y))?a)) as a list of estimators of
additional base estimator (e.g. SM trig y)
--D)
X (x. Then our hybrid estimator takes the following form:
12 ni
SM (x.y)
=
(x, (x,x) + (x,, (x,x)
(97)
with constraint:
T1 n4
+
(98)
==
and where the estimators of >-coefficients and base estimators are given as:
(x.(x,y)
T. Ald
(x) T T.
( (-y)
(99)
SM i) (x,y) T (x)
Tm
(
(100)
for m. ni (x), 1m ((sis) (y), A. (x), (T (-x) given by Eq.
97 92 95
96
24

---

Published as a conference paper at ICLR 2022
Definition H.2 (Hybrid Zero-One- One-Mixture estimators for clustered data). Assume that inputs (queries) and y (keys) can be modeled by ny and ny clusters respectively with centers n, and r, E RR, (i -_,..,, Ri. j = 1,...,RR.) Rg)) We denote A as the complex matrix satisfying Eq. 85)
with center I., Furthermore, we denote E (SM b.j (x,y)) as list of estimators of
the softmax kernel SM (x. y) and A = (x) Pant as a list of estimators of X X constructed independently from &
Then our hybrid estimator takes the following form:
ng nk
SM c.A (x,
Σ
(x.
Y)SM (x,
(101)
i=
with constraint:
(x, =1
(102)
and where the estimators of X-coefficients and base estimators are given as:
yx) (xx)),,()
(103)
SM 1
(x, y)
A (x)
Th (A:))--())
(104)
for M. (x), Vm ((A--)-1(y)) (y) given by Eq.
91 92 with (,(x) being a scalar indicating whether
belongs the i-th cluster of the ny clusters and similarly with yy(y) being scalar indicating whether belongs to the j-th of the ng clusters.
In our experiments used the formulation from the second definition with A constrained to be real diagonal. In other words, we are using Eq. 88 to construct A.
H.3 ADDITIONAL EXPERIMENTAL DETAILS
To create the synthetic data, we first generate random vectors Xo, Yo in R:50 Then we generate four random orthogonal matrices O E IR 50x50 where i = 1, and 1,2 and use X; = O Xo and Y; = O for 1, as the mean vectors for our Gaussian clusters. For each pair of clusters, the minimal value s of Eq. 21 may be non-zero if sign(x;,.)) usage of real diagonal matrices. To control the value of s, we manually adjust the sign of each element of X; and Y, The data for input is the combination of two clusters, each consisting 1000 data points following Gaussian distribution with mean vectors XI,X; and the common covariance matrix o²1. And similar constructions are done for input y. We create four synthetic data sets, with different values of s and magnitudes of the standard deviation o. After this step, we normalize the data points by controlling their & norms to make the true softmax value be within reasonable range. After the data is created, we first use K-means clustering algorithm with & = 2 to cluster x and y. and obtain centers (x; Vj vj) (. (1.2)- Then use the of real diagonal matrices to minimize Eq. 21 for all pairs of centers of x and y to generate four complex exponential base estimators. To choose coefficients, use the more efficient approach described above, i.e. using indicator vectors, and these are already computed in the clustering process. For each data set, we compare the performance of these estimators by calculating the mean square error of them, and we do this multiple times with different numbers of random features used. All the results are obtained by averaging over 20 repetitions. In addition to the average performance which can found in Fig. G we also record the maximal and minimal empirical MSE over 20 repetitions for all estimators in the Table where 8 = [0.08, 0.05, 0.04, 0.03), 82 = [0.11, 0.09, 0.08, 0.07), #3 [0.05, 0.03, 0.03, 0.02), 81 [0.07, 0.06, 0.05, 0.04] representing the lists of values for all pairs of clusters for different synthetic data sets.
= sign(y,..) for some k due to the
25

---

Published as a conference paper at ICLR 2022
Data Set
Number of Random Features
Hybrid -lusste [0.0016, 0.0183)
FAVOR+
Trigonometric
20
[0.0071, 0.1455]
[0.0039, 0.0808
50
[0.0011, 0.0081]
[0.0026, 0.2541]
[0.0031, 0.0510]
80
[0.0005, 0.0052]
[0.0032, 0.0282]
[0.0010, 0.0281]
81,
100
[0.0006, 0.0036]
[0.0022, 0.0528]
[0.0013, 0.01
120
[0.0004, 0.0025]
[0.0028, 0.0534]
[0.0012, 0.0135]
150
[0.0003, 0.0020]
[0.0014, 0.0582]
[0.0007, 0.0134]
200
[0.0002, 0.0010)
[0.0013, 0.0205]
[0.0004, 0.0087]
20
[0.0032, 0.0265]
[0.0095, 0.2071]
[0.0036, 0.0989]
50
[0.0012, 0.0079]
[0.0036, 0.2123]
[0.0019, 0.0565
80
[0.0007, 0.0059)
[0.0037, 0.0570]
[0.0008, 0.0208)
82,
100
[0.0007, 0.0040]
[0.0016, 0.0280]
[0.0010, 0.0164]
120
[0.0005, 0.0025]
[0.0013, 0.0262]
[0.0007, 0.0246]
150
[0.0004, 0.0024]
[0.0010, 0.0407]
[0.0004, 0.0190]
200
[0.0003, 0.0018]
[0.0018, 0.0236]
[0.0004, 0.0134]
[0.0083, 0.0300)
[0.0362, 0.1218]
[0.0174, 0.0490)
50
[0.0046, 0.0183]
[0.0125, 0.1262]
[0.0081, 0.0249]
[0.0029, 0.0105]
[0.0108, 0.0281]
[0.0044, 0.0143]
83, V10
100
[0.0024, 0.0056]
[0.0077, 0.0421]
10.0041,0.0099.
120
[0.0020, 0.0046]
[0.0077, 0.0402]
[0.0034, 0.0077]
150
[0.0015, 0.0041]
[0.0049, 0.0395]
[0.0025, 0.0073]
200
[0.0011, 0.0031]
[0.0037, 0.0150]
[0.0018, 0.0050]
[0.0097, 0.0294]
[0.0332, 0.1553]
[0.0179, 0.0485)
50
[0.0049, 0.0172]
[0.0176, O. 0.1345]
[0.0058, 0.0260)
[0.0031, 0.0113]
[0.0113, 0.0415]
[0.0038, 0.01
84 V10
100
[0.0028, 0.0061]
[0.0080, 0.0292]
[0.0031, 0.0100]
120
[0.0022, 0.0045]
[0.0065, 0.0185]
[0.0028, 0.0128]
150
[0.0016, 0.0044]
[0.0045, 0.0235]
[0.0021, 0.0094]
200
[0.0012, 0.0031]
[0.0046, 0.0166]
[0.0015, 0.0067]
Table 3: Minimal and maximal empirical MSE with 20 repetitions
26

---

Published as a conference paper at ICLR 2022
I LANGUAGE MODELING TRAINING DETAILS AND ADDITIONAL RESULTS
For the Language Modeling tasks, we trained a 2-layer LSTM of hidden sizes 200 and 650 on the PennTree Bank (Mareus al. 1993) and the WikiText2 dataset (Merity et al., 2017) respectively. We tied the weights of the word embedding layer and the decoder layer (Inan et al. 201 0117 Press & Wolf 2017). Thus we could treat the language modeling problem as minimizing the cross entropy loss between the dot product of the model output (queries) and the class embedding (keys) obtained from the embedding layer with the target word.
We now present detailed statistics of HRFs on the Penn Tree Bank dataset. The mean and the standard deviation of the statistical metric results in Fig. 3 Sec. 4.2 shown in Table The distribution of the lengths of the keys and the queries are shown in Figure 7 For our experiments with the Angular Hybrid and the Gaussian Hybrid, the number of random features are 8, 16, 32, 64 for the base estimators, and 8 for the coefficient estimators.
Table 4: Results computed over 10 runs Penntree Bank Dataset. Boldface denotes the hest one, and underline denotes the second best. Negative fractions for Favor+ reported it produces positive random features.
1d- Wassersteinn
RF types
d= 64
d
d= 256
d= 512
FAVOR+
3171.35 + 69.04
3142.65 + 85.55
3050.58 + 108.25
2985.18 + 73.15
Trig.
1577.44 + 13.59
1582.50 + 13.42
1
1576.50 5.83
Gaussian
1520.32 + 11.65
1521.10 + 8.79
1516.50 6.66
1515.41 5.96
Angular
1395.03 +: 79.91
1445.24
58.27
1394.97
1425.88 65.82
KS statistics
FAVOR+
0.573 = 0.0073
0.569 + 0.0098
0.558 + 0.013
0.551 + 0.0092
Trig.
0.424 0.0040
0.427 + 0.0041
0.425 + 0.0024
0.424 + 0.0023
Gaussian
0.411 0.0029
0.410 - 0.0019 0.0019
0.410 0.0013
0.410 0.0014
Angular
0.404 0.017
0.419 0.010
0.418 0.0089
0.0099
Negative fraction
Trig.
0.257 0.0057
0.258 1 0.0039
0.258 1 0.0021
0.257 0.0017
Gaussian
0.155 0.0039
0.159 1 0.0019
0.157 0.0011
1
Angular
0.152 0.0041
0.158 = 0.0022
0. 153 0.0018
0.157
Keys length
Queries length
1200
2000
1750
1000
1500
800
Count
1250
Count
600
1000
750
400
500
200
2.5
5.0
7.5
10.0
1
2
4
Figure 7: Distribution of the lengths of the keys and queries in the Penn'Tre Bank dataset
27

---

Published as a conference paper at ICLR 2022
Since the HRF mechanisms can be sensitive the lengths of the keys and the queries, for the language modeling task on the WikiText2 dataset, we added a regularizer to constrain the lengths of the keys and the queries to be close to 1. However, constraining the lengths close to hurts the model performance and so we chose to use a temperature scaling as (Rawat et al. 2019). Thus before passing the computed dot product the softmax layer, scaled the dot product by 1. Our final loss function for the language modeling task was:
()) = LCF X B(luul2 - 1)a + A2B([||:|22 - 1)2
(105)
where |1.112 is the row-wise Le nomm of the appropriate matrices, and Log is the cross-entropy loss. The distribution of the lengths of the keys and the queries are shown in Figure
Finally present some additional results the Wiki'eex22 dataset. For experiments the WikiText2 dataset, we chose A1 = A2 = and r = 6. Our model trained with these hyperparameters achieve perplexity score of 105.35 the test set after 50 epochs.
We compute the 1-dimensional Wasserstein distance and the Kolmogrov-Sminet (KS) met- ric (Ramdas al.] (017)) between the approximated softmax distribution and the true softmax dis- tribution. The mean and the standard deviation of the statistical metric results in Fig. Bis shown in Table 55] For our experiments with the Angular Hybrid and the Gaussian Hybrid, the number of random features are 8, 16, 32, 64 for the base estimators, and 8 for the coefficient estimators.
14000
Jumo
en Dimension
Feakir
Figure 8: Statistical metrics measuring softmax matrix approximation quality on the WikiText2 dataset. For standard estimators, the number of random features are 64, 128, 256, 512. To make fair comparison, for the hybrid variants, the configurations leading to the similar number of FLOPS operations per random feature map creation were applied. Results reported over 10 runs.
Table 5: Results are computed over 10 runs on the Wikitext2 Dataset. Boldface denotes the best onc, and underline denotes the second best.
Wasserstein
RF types
d= 64
d= 128
= 256
d=
FAVOR+
12970.38 + 119.49
12971.67 + 108. 79
12965.82 8: 89.62
12954.17 83.95
Trig.
12913.66 + 59.22
12912.65 = 57.42
12910.91 - 37.23
12913.25 35.12
Gaussian
12809.03 97.2
12805.71 78.91
12808.02
+ 84.57
12809.40
+ 65.88
Angular
12833.56 + 41.87
12765.44 + 38.91
12561.26
26.48
12534.48
25.88
KS statistics
FAVOR+
0.687 + 0.0043
0.687 + 0.0051
0.687 + 0.0019
0.686 + 0.0067
Trig.
0.685 + 0.0071
0.685 + 0.0053
0.685 + 0.0029
0.685 + 0.0014
Gaussian
0.683 0.0092
0.683 + 0.0073
0.683 1 0.0049
0.682 + 0,0089
Angular
0.682 1 0.0079
0.682 1 0.0066
0.682 1 0.0037
0.678 1 0.0064
The above tables (Table and Table 5 show that the our hybrid variants consistently outperform Favort and the trigonometric random features.
I.1 LANGUAGE MODELING USING HRF
In this subsection, we will describe how one can use HRF to train a LSTM on the language modeling task the PennTree Bank dataset, similar to the experiments carried out in (Rawat et al. 2019
28

---

Published as a conference paper at ICLR 2022
Keys length
Queries length
4000
2500
3000
Count
Count
1500
2000
1000
1000
500
0
0
1
Figure 9: Distribution of the lengths of the keys and queries in the WikiText2 dataset
We trained 2-layer LSTM model with hidden and output sizes of 200, and used the output as input embedding for sampled softmax. We sampled 40 negative (other than true) classes out of 10000 classes approximate the expected loss in each training epoch. As observed in Fig. - the relative error all three estimators grow exponentially fast with embedding norm r. Therefore, keep un-normalized embeddings during sampling, then even though we could get unbiased estimation of the loss, the variance could be high. Such bias-variance trade off also mentioned in paper (Rawat al. 2019). To solve this issue, used normalized input and class embeddings during sampling to generate biased sampling distribution with lower variance, while keeping un- normalized embeddings calculate the loss. We trained our model for 80 epochs, with batch size equal to 20, dropout ratio in LSTM equal to 0.5. Implementation details could be seen in our Github repository.
185.0
151
Surpled
FWOR+ mm1004
Trigonomatic, mr 1024
140
Peploay
157.5
Angular Hybrd, m=1128. 8)
148
1525
140
145
147.5
14.0
50
70 80 50
85
70
Epoche
Epochs
Figure 10: Language Modeling with softmax sampling: training results on PennTree Bank dataset (50- 80 epoch window zoomed in on the right subfigure). The solid line for each method is estimated over 25 independent runs. The shaded areas represent perplexity within standard deviation of the average. FA- Vii+/trigooometric mechanisms used 1024 random features. To make fair comparison, for HRFs, the figurations leading to similar number of FLOPS operations per random feature map creation were applied. 128 and random features is used to estimate the softmax estimators and X-coefficient in HRFs respectively. Even though statistical metrics for HRFs are better in Fig. 3 the difference between different random feature estimators is very significant.
We compared our hybrid estimator with Trigonometric, FAVOR+, the uniform method which sam- pled 40 negative classes with equal probability, well as the sampled softmax method which uses the unbiased true probability calculated from full softmax as the sampling probability. Trigonomet- ric and FAVOR+ mechanisms used 1024 random features. To make comparison fair, hybrid variants were using (m, cooooiggurat characterized by the similar number FLOPS needed to con- struct their corresponding random features regular RF-estimators. We reported comparison
29

---

Published as a conference paper at ICLR 2022
results averaged over 25 independent runs in Fig. 10 The perplexity score in Fig. 10 is defined as 2 entropy loss, We could conclude that even though the statistical metrics for HRFs are better in Fig. 3 the difference between HRFs and other random feature estimators in this specific softmax sampling downstream task not very significant.
J SPEECH EXPERIMENTS: ADDITIONAL DETAILS
Tested Conformer-Perforeee models consisted of / = 17 conformer layers. Each attention layer used H! 8 heads. The embedding dimensionality was p 512 and since dimensions were split equally among different heads, query/key dimensionality was set up to dQk = 64. Each input sequence was of length 500. Padding mechanism was applied for all tested variants.
K DOWNSTREAM ROBOTICS EXPERIMENTS
In both Robotics experiments we found query/key lyy-normalization technique particularly effective (queries and keys of L2 Lo-norm equal one) thus we applied it in all the experiments.
K.1 VISUAL LOCOMOTION WITH QUADRUPED ROBOTS
We evaluate HRF-based attention RL policies in a robotic task of learning locomotion on uneven terrains from vision input. We use the quadruped from Unitree called Laikago (lai). It has 12 ac- tuated joints, 3 per leg. The task is to walk forward on randomized uneven terrain that requires careful foot placement planning based on visual feedback. The ground made of a series of step- stones with gaps in between. The step stones widths are fixed at 50 cm, the lengths are between [50, 80] cm in length, and the gap size between adjacent stones are between [10, 20] cm. It perceives the ground through depth cameras attached to its body, one the front and other the belly facing downwards. We use Implicit Attention Policy (IAP) architecture (masking variant) described in Choromanski et al. (2021a) which uses tefoormer-bseed attention mechanism to process 32 24 depth images from the 2 cameras.
We apply hierarchical setup solve this task. The high level uses IAP-rank with masking to process camera images and output the desired foot placement position. The low level employs position-based swing leg controller, and an model predictive control (MPC) based stance leg con- troller, to achieve the foot placement decided by high level. The policies are trained with evolution- ary strategies (ES). In Fig. (left) we compare FAVOR+ with angular hybrid approximation in IAP. HRF with 8 8 random projections (m random projections. A series of image frames along the episode of a learned locomotion HRF-based IAP policy is shown bottom right of Fig. On the top-left corner of the images, the input camera images are attached. The red part of the camera image is the area masked out by self-attention. The policy learns to pay attention to the gaps in the scene in order to avoid stepping them.
8) performs as well as the softmax kernel with 256
The training curves in Fig. (left) are obtained by averaging over top runs. The shaded regions shows the standard deviation multiple runs.
K.2 ROBOTIC MANIPULATION: BI-MANUAL SWEEPING
Here we consider the bi-manual sweep robotic-arm manipulation. The two-robotic-arm system needs to solve the task of placing red balls into green bowls (all distributions are equally rewarded, so in particular robot can just use one bowl).
In this bi-manual sweeping task, the reward per step between 0 and 1, the fraction of the blocks within the green bowls. The loss for this task is the negative log likelihood between the normalized softmax probability distribution of sampled actions (one positive example, and all others uniform negative Dounter-examplesy and the ground truth one-hot probability distribution for a given obser- vation from the training dataset. For more details see (van den Oord et al.) 2018). This is a 12-DoF cartesian control problem, the state observation consists of the current cartesian pose of the arms and 500 x 500 x 3 RGB image. An example of the camera image is given in Fig. []
The training curves Fig. 4 (right) are obtained by averaging over 3 learning rates: 10-3 10-1, 10-5
30

---

Published as a conference paper at ICLR 2022
Figure 11: An example of the camera image for the bi-manual sweep robotic-arm manipulation task.
L COMPUTATIONAL GAINS OF HYBRID ESTIMATORS
Consider the general hybrid estimator of SM(x. y) from Equation using p + 1 base estimators.
Assume that the ttth base estimator for k 1,.... P + 1 applies m random features each of the corresponding & random maps and that random features used in each of the corresponding we conclude
1x, random maps to approximate 1 X-coefficient for 1, 1,-... p. From Equation that time complexity of constructing (z) is:
T=o (4 - + 1p-1)md + (nd + md + mn)
p+1 P
(106)
Thus the resulting ((mmm)-imeenioonl random feature map V (z) can be constructed in time (nnd+ mn + mn) opposed (mndd computational gains. This is true the (1mm)mdmmmnniion) HRFs can provide similar quality as
the case for the regular estimator, providing substantial
their mn- cin-eiinniional regular counterparts. It turns out that this is often the case, see: Section D.4
31


Reviews:
[{"review_id": "sXk-e-k_XVK", "reviewer": "Reviewer_Ks6B", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n\n1, The motivation of this work is clear: current random features based algorithms are difficult to approximate softmax kernels with very small/large values. I agree with the authors on this issue.\n\n2, the algorithm to approximate softmax kernel in Sec. 2.3.1 can be regarded as a combination of SM^{++} and SM^{trig}, with an extra cosine similarity function.\nThis function is in essence close to the zero-order arc-cosine kernel [Cho and Saul, 2009], and thus is actually approximated by the standard random features with the Hesivide activiation function. The algorithm enjoys theoretical guarantees on smaller worst-case relative errors.\n\n3, Experiments on various applications are good to support their findings.\n\nCons:\n\n1, Personally, this paper is not easy to follow even though I’m quite familiar with random features. For example, I understand the first part in Sec. 2.3, e.g., Eq. (14), (15) till reading Sec. 2.3.1; The notation SM^{hyb}_{m,n}(x,y) is not defined until Lemma 2.4; Some notations are unclear. I suggest the author carefully polish this paper.\n\n2, the relationship between Sec 2.2 and the algorithm in Sec. 2.3 is not clear to me. The authors attempt to unify trigonometric and positive random features under a same framework, i.e., SM^{cexp}. However, the proposed algorithm in Sec. 2.3 to approximate softmax kernel is the exact combination of them. The complex exponential estimator appears unused for Gaussian/softmax kernel approximation. More importantly, this estimator is asymmetric, which appears strange. This only holds for two special cases (symmetric due to A=(A^T)^{-1}) by A:= iI_d or I_d. This estimator appears invalid/asymmetric for general case. In fact, the asymmetric property destroyed the foundation of kernel associated with RKHS. In this case, this estimator (the approximated kernel) under a general matrix A is not well-defined in kernel methods.\n\n3, In theory, Theorem 3.1/3.5 is not clear to show that whether SM^{anghyb} achieves variance reduction or smaller MSE when compared to SM^{trig} and SM^{++}. Refer to a survey with discussion on the variance reduction: \n\n[S1] Liu, F., Huang, X., Chen, Y., & Suykens, J. A.K. (2020). Random Features for Kernel Approximation: A Survey on Algorithms, Theory, and Beyond. TPAMI 2021.\n\n4, it’s good to see many experiments but evaluation on kernel approximation error is needed. Experimental results in Figure 3 is not enough. In fact, kernel approximation error (as well as unbiasedness) is quite important in random features based algorithms. To validate the reduction on approximation error, it would be better to compare the proposed algorithms with SM^{trig} and SM^{++} as well as other representative algorithms including ORF and QMC on typical datasets.\n(ORF and QMC can be also used for softmax/Gaussian kernel approximation. The proposed algorithm might be inferior to ORF and QMC due to the standard sampling strategies. It's ok for me but at least it can beat SM^{trig} and SM^{++} on typical datasets.)\n\nBesides, regarding to the experiments on Transformers for language modelling, the following refs is missing.\n\n[S2] Peng, H., Pappas, N., Yogatama, D., Schwartz, R., Smith, N., & Kong, L. Random Feature Attention. ICLR2021.\n", "summary_of_the_paper": "Author rebuttal checked. I increase my score to 8. BTW, I strongly suggest the authors polish this paper as the current version appears uneasy to follow, especifically for researchers unfamiliar with RFF.\n=============================================================\nThis paper proposes a hybrid random features that combines the classical trigonometric and positive random features, which aims to achieve good approximation performance for softmax kernel with both small and large values. This framework is also applicable to Gaussian kernels. Numerical experiments on softmax kernel approximation in linear attention and robotics are conducted to support their algorithm and theory.", "summary_of_the_review": "In sum, this paper provides a novel random features algorithm with theoretical guarantees and has good applications. Nevertheless, some parts on algorithm and experiments are unclear, which would devalue this paper. I would like to see the response from the authors", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "aLoEsc5FJ8v", "reviewer": "Reviewer_YmNP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main idea of the paper – to try to get the “best of both worlds” when different types of random features are useful in different regimes – is very nice. The theoretical results show that HRFs can be superior to either of the “base” random features used. The experiments in section 4.1 on pointwise kernel estimation also nicely support this narrative. The other experiments suggest at least some benefits, with perhaps the strongest results being for the speech model example in section 4.4. However, there are a few serious issues that need to be addressed.\n\n## Writing and presentation\n\nThis was a very confusing paper to read.\n\n### Notation \n\nThere is an incredible amount of very dense notation, much of which it should be possible to drastically simplify with a little thoughtful effort. For example, there are many cases where to approximate a function $f(x, y)$, a random feature representation of the form $\\mathbb E[\\psi^{(1)}(x)^\\top\\psi^{(2)}(y)]$ is used. But rather than, as I have done, reuse the same symbol for both and distinguish the random functions with an index, different symbols are used (e.g., $\\alpha$ and $\\beta$). This forces the author(s) to write about $\\phi$, which can be either $\\alpha$ or $\\beta$. Thus, whereas a single symbol plus an index would do, the reader instead has to juggle three symbols. For a particularly egregious example of this, see the sentence that begins before eq. (15). \n\nAnother reason for the difficult notation is the author(s) try to make the methods presented as general as possible, which puts a heavy burden on the reader. However, often this generality is now actually used very much. For example:\n\n* Nearly all of the focus is on “bipolar” HRFs with just $p + 1 = 2$ “base” random features. However, the general case is covered, which requires, e.g., to define $\\lambda^k$ for $k=1,\\dots,p$, whereas for the most part only a single $\\lambda$ function is needed. \n* $\\lambda^k$ can take values in $\\mathbb R$, but in all examples it is constrained to [0,1] \n\n\n### Paper structure \n\nThe paper needs to be restructured on both local and global scales. At the local scale, definitions and semantic explanations often come long after something is used. For example, what is the purpose of $\\lambda^k$? It is unclear until a few pages after it is defined. Or, the definitions of $\\otimes$ and $\\prod^\\star$ in Lemma 2.4 are at the very end, leaving the reader unsure what any of the four displayed equations that come before actually mean. Sometimes objects are never defined (e.g., MSE and $\\mathcal P(\\mathbb R^d)$) or only long after they are first used (e.g., $\\boldsymbol i$ is finally defined right before Lemma 2.4). Also, global definitions are put inside lemmas and theorems (e.g., Lemma 2.2). \n\nAt the global scale, the main contribution of trading off whether the trig or positive random features are “active” is in the paragraph at the top of p. 5. If some version of this came, say, after section 2.1, followed by the material in section 3 (theoretical guarantees), the reader would have everything they need to understand all but the experiments in section 4.2, which in any case aren’t very interesting. \n\n### Use of colons\n\nThere is liberal use of colons in the middle of sentences. Please remove these. \n\n## Scope of contribution\n\nThere are a number of places where I don’t think the paper properly describes the scope of its contribution. \n\n1. Abstract: “By generalizing Bochner’s Theorem for softmax/Gaussian kernels and leveraging random features for compositional kernels, the HRF- mechanism provides strong theoretical guarantees”. Given the generality of Bochner’s theorem, this sounds like a major contribution and an important part of what makes HRFs work. However, I believe the generalization of Bochner’s theorem is Lemma 2.3, which seems to only be used for the HRFs for clustering structure; these make only a brief appearance in the paper: there is no theory for them and they earn a small synthetic experiment (Section 4.2), but there practical utility is never shown. Thus, Lemma 2.3 appears to have nothing to do with the “strong theoretical guarantees” developed in section 3, which are only for the bipolar kernel with trig and positive random features. \n\n2. Related work at the end of section 1: “rather than focusing on improving [the] sampling mechanism for a given approximation algorithm, [HRFs] provide a completely new algorithm.” I find this claim to be misleading. HRFs provide a way to combine existing random feature algorithms. \n\n3. Generality. As I’ve mentioned above, the paper presents HRFs in great generality, but very little of that generality is actually used in a consequential way. It may be that this generality is quite useful! I really like the idea, and could imagine othe HRFs will be developed or the cluster-based HRF will find use cases. But the theory and experiments given in the paper don’t support such grand claims about “strong theoretical guarantees” for the “HRF mechanism.” On the contrary, the paper shows that one particular instantiation of the HRF mechanism has smaller relative MSE than the underlying components. \n\n## Experiments\n\nSection 4.1: This was a nice clean demonstration of the theory. I think the red dotted lines in Fig. 2 are for zero. Please clarify this in the caption. \n\nSection 4.2: As mentioned above, unless this experiment is supported by some practical application, it should be moved to the appendix as a “proof-of-concept”.\n\nSection 4.3: It seems like none of the random feature approaches give very good approximations (Fig. 4). There are no error bars in Fig. 5, so I’m unsure if the differences are meaningful. Moreover, I’m not convinced the ordering would remain the same if the experiment were run for, say, 10 more epochs. \n\nSection 4.4: This is a potentially strong experiment. However, the statement that “ ≥ 0.2% are considered to be statistically significant for that task” makes no sense. The differences are either statistically significant or not, given a particular test and choice of test size. \n\nSection 4.5: The results here are promising. A 3x improvement in FLOPS certainly is something. But the lower variation in reward, particular in the step-stone locomotion task, is also noteworthy.  \n\n## Recommendations \n\n1. Simplify the notation \n\n2. Focus only on bipolar random features in the main text and move the other types and the general definitions  (e.g., material around eq. (14), Gaussian lambda coefficients, adaptation data admitting clustering structure, section 4.2) to the SM or to a separate section that can be skipped on first reading. \n\n3. Restructure the paper so motivation and semantic interpretations come before definitions\n\n4. Clarify scope and contributions. \n\n5. Strengthen experiments by including error bars for all performance metrics and running the experiments in section 4.3 for at least 10 more epochs, or until the perplexities stabilize. \n", "summary_of_the_paper": "The paper makes two methodological contributions: a new approach to constructing randomized approximations to Gaussian and softmax kernels and a proposal to combine multiple randomized approximations to create “hybrid random features” (HRFs). The idea of the hybrid random features is to activate a particular approximation for pairs (x,y) for which that approximation is accurate. Some theoretical results support the use of a particular instance of HRFs that combine trigonometric and positive random features, showing the hybrid has lower mean squared error than either type of random feature. Experiments are presented to verify the benefits of HRFs on a variety of tasks, with a focus on their use in neural network settings. ", "summary_of_the_review": "A potentially useful method with a nice idea and some promising results. But the paper structure and writing need to be improved, the scope and contributions clarified, and the experiments tightened up. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "14b1hghv4V", "reviewer": "Reviewer_iboB", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**The contributions of the paper are original and very well motivated**, as they aim at providing a scalable and unbiased estimator of the widely used softmax and gaussian kernels which also has low variance.\nNamely, the principal contributions are:\n1. Hybrid Random Features as composition of base estimators, unbiased, that adapt to data points and can be linearized;\n2. Generic form of base estimators, named complex exponential estimator;\n3. Three instantiations of HRF that provide low variance estimators.\n\nThe proposed HRF can benefit the ICLR community at large and in particular the fervent research around transformers.\nThe theoretical results seem sound, however I haven't checked in details the proofs of Lemma 2.4, Lemma 3.2, Lemma 3.4 and Theorem 3.4.\n\n**The paper itself is hard to read** in particular because intuitions and justifications for the derivations and definitions are rarely provided.\nI have few suggestions for improving its presentation:\n\n1. The ordering of the exposition of the contributions does not help the reader understand the train-of-thoughts and intuitions. I'd suggest to first introduce HRF (Section 2.3) to then show how base estimators can be chosen (Section 2.2) and how $\\lambda$-coefficients can be defined accordingly.\nMoreover, it would help the presentation of the complex exponential estimators to start from defining the new feature map $\\Phi_M^m(u)$'s components $\\exp(wMu)$ to clarify first the goal of defining a generalized form of random features that can be specialized to well known ones depending on $M$. In the current presentation the new choice of $z$ comes out of the blue.\n\n2. When presenting HRF, the choice of having kernels as $\\lambda$-coefficients and not e.g. constants is not motivated and not sufficiently highlighted. I understood only later, when examples of instantiations are provided, the interest of this choice which is to reduce the variance of the estimator. Still it is not clear to me why they need to be shifted (with the offset parameter $a_k$) and whether in practice this choice provides better results than fixing these coefficients to constants, e.g. learned from a subset of data. \n\n3. Maths and notation could be simplified. Namely, Section 2.3 could be written for $p=1$ (2 base estimators), and the general form for $p > 1$ could be reported in the appendix.\n\n4. When presenting the theoretical results, it should be made more explicit which results are new and which are not. This is not clear especially for Lemma 2.3. A reference for Definition 2.1 should be provided and in general pointers to the proofs in the appendix should be added. \n\n5. The experimental section of the main text is too packed. It is valuable to have carried out the empirical analysis on multiple tasks and contexts, but reporting them all comes at the cost of losing clarity, as experiments and results are not thoroughly described and commented.  \nMore precisely:\n(a) Figure 1: colorbar is missing and it is not clear why the light blue is both for low and high MSE values.\n(b) Section 4.1: What study is referred to as complete ablations? Why the results for gaussian $\\lambda$-coefficients are not reported? What is the number of random features for Angular Hybrid?\n(c) Section 4.2: Why is $s$ a vector of size 4 when $l(A)$ should be a scalar?\n(d) Section 4.3: Which task is performed here and what is it meant by negative classes? Are the differences between methods for the three metrics significant? What is the perplexity score? A note should be added on why negative fractions are not reported for FAVOR+.\n(e) Section 4.5: Why is training with state-of-the-art estimators unstable on this task (loss going to Nan for particular numbers of random features)? What are the \"drastic compute improvements\" given that score differences between methods are not significant for the same number of iterations?\nI'd suggest to select the most important experiments and present them thoroughly in the main text, while reporting the rest in the appendix. \n\nMinor:\n1. MSE in this context should be defined in the introduction.\n2. Gaussian kernel's inputs are missing in Eq. (1).", "summary_of_the_paper": "The paper proposes a new type of estimators of softmax and gaussian kernels based on random features, consisting of compositions of base estimators such as the trigonometric and positive random features. The composition is a linear function, such that the new estimator (called hybrid) is unbiased, whose coefficients (called $\\lambda$-coefficients) are independent of the base estimators and are also kernel functions (to better adapt the estimator to the compared data) hence also need to be linearized for scalability. \n\nThree different instantiations of the hybrid estimator are proposed using base estimators of the form of the defined complex exponential one (an asymmetric generalization of random feature estimators). In particular the angular one (i.e. $\\lambda$-coefficients depending on the angle between the pair of data points) is shown to have low variance and low maximal relative error for small and large kernel values.\nHRF has also a lower computational complexity in the total number of random features w.r.t. non compositional estimators.\n\nExperiments are carried out on multiple applications, first to verify the approximation capability of the hybrid estimator w.r.t true kernel, and trigonometric and positive estimators in simple settings, then to show the improvement (either in quality scores or computational) of using the proposed estimator in models requiring the use of softmax at scale.", "summary_of_the_review": "The paper makes strong contributions with potential high impact, however its presentation is over-complicated and lacks clarity. I suggest to improve on this point and I am willing to increase my score accordingly.\n\n## UPDATE AFTER REBUTTAL\nThe presentation of the paper has improved and, while some notation are still hard to read, I believe that the paper's contributions are now more accessible. I am hence increasing my score to 8. I am looking forward to seeing how the paper will be presented in video lectures.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Kernel methods are among the most flexible and powerful approaches of our times. Random features (RF) provide a recent mechanism to also make them scalable due to the associated finite (and often small)-dimensional approximate feature map (in the paper referred to as linearization). The focus of the submission is the linearization of the softmax kernel (defined in (1)) while making sure that the obtained RF approximation is accurate simultaneously for the small and the large kernel values. The authors present a hybrid random feature (HRF, defined in (8)) construction parameterized by base estimators and weights, and show that specific choice of these parameters is capable of implementing the goal. Some of the HRF estimators are also accompanied by theoretical guarantees (Section 3). Their numerical efficiency is illustrated (Section 4) on synthetic examples and in the context of natural language and speech modelling, and in robotics.

Scaling up kernel methods is a fundamental task of machine learning. The authors present a nice and valuable construction in this direction which can be of both theoretical and practical interest to the community.

The submission would benefit from implementing the remarks of the reviewers to improve its clarity.
Decision: Accept (Poster)