HYBRID RANDOM FEATURES
Year: 2022

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., 2021b). By generalizing Bochner’s Theorem for softmax/Gaussian kernels and leveraging random features for compositional kernels, the HRFmechanism 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.

heading: 1 INTRODUCTION & RELATED WORK
text: Consider the softmax and Gaussian kernel functions K : Rd×d → R defined as follows:
SM(x,y) def = exp(x⊤y), Kgauss(x,y) def = exp(−∥x− y∥ 2 2
2 ). (1)
These two are prominent examples of functions used in the so-called kernel methods (Gretton et al., 2005; Zhang et al., 2018) and beyond, i.e. in softmax-sampling (Blanc & Rendle, 2018). Random features (RFs, Rahimi & Recht, 2007; Liu et al., 2020; Peng et al., 2021) yield a 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 ϕ : Rd → Rm. The tight relationship between softmax and Gaussian kernels given by the transformation Kgauss(x,y) = exp(−∥x∥ 2 2 2 )SM(x,y) exp(− ∥y∥22 2 ) provides a mapping of any random feature vector ϕSM(u) for the softmax kernel to the corresponding one ϕgauss(u) = exp(−∥u∥ 2 2
2 )ϕSM(u) for the Gaussian kernel, thus we will focus on the former kernel. The classic random feature map mechanism ϕtrigm for the softmax kernel, obtained from Bochner’s Theorem applied to the Gaussian kernel (Rahimi & Recht, 2007), is of the form:
ϕtrigm (u) = 1√ m exp( ∥u∥2 2 ) ( sin(ω⊤1 u), ..., sin(ω ⊤ mu), cos(ω ⊤ 1 u), ..., cos(ω ⊤ mu) )⊤ , (2)
where m stands for the number of random features and ω1, ..., ωm iid∼ N (0, Id).
The above (common in most downstream use-cases) method for linearizing softmax/Gaussian kernels was recently shown to fail in some of the new impactful applications of scalable kernel methods such as implicit-attention Transformer architectures called Performers (Choromanski et al., 2021b). We denote by MSE the mean squared error of the estimator (i.e. its variance since all estimators considered in this paper are 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. †Authorship in alphabetical order ‡Independent researcher. Work done during postdoc at Yale University
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 a new positive random feature map for unbiased softmax kernel estimation was applied:
ϕ++m (u) = 1√ 2m exp(−∥u∥ 2 2 ) ( exp(ω⊤1 u), ..., exp(ω ⊤ mu), exp(−ω⊤1 u), ..., exp(−ω⊤mu) )⊤ . (3) 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-zero ones or softmax sampling) the algorithm needs to process simultaneously very small and large kernel values. The following natural questions arise:

heading: 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 ?
text: We give affirmative answers to both of the above questions by constructing a new class of random feature maps techniques called hybrid random features or 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., 2021b) are just special corollaries of the more general result from complex analysis, (b) integrate in the original way 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 computational gains. We also conduct their thorough empirical evaluation on tasks ranging from pointwise kernel estimation to downstream problems involving training Transformer-models or even end-toend robotic controller-stacks including attention-based 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 variance 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 as orthogonality (Choromanski et 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 in 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 kernel in Transformers proposed in (Choromanski et al., 2020; 2021b) led to fruitful research on the extensions and limitations of these methods. Schlag et al. (2021) replaced random features by sparse deterministic constructions (no longer approximating softmax kernel). Luo et al. (2021) observed that combining L2-normalization of queries and keys for variance reduction of softmax kernel estimation with FFT-based implementations of relative position encoding and FAVOR+ mechanism from Performers helps in training. Trigonometric random features were applied for softmax sampling in (Rawat et al., 2019). Several other techniques such as Nyström method (Yang et al., 2012; Williams & Seeger, 2000; Rudi et al., 2015) were proposed to construct data-dependent feature representations. Even though, as we show in Sec. 2.4, certain instantiations of the HRF mechanism 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.

heading: 2 HYBRID RANDOM FEATURES

heading: 2.1 PRELIMINARIES
text: 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 a random feature (RF) map representation if it can be written as
K(x,y) = Eω∼Ω
[ l∑
i=1
ξi(x, ω)ξi(y, ω)
] , (4)
for some ξi : Rd × Rd → R, and where ω is sampled from some probabilistic distribution Ω ∈ P(Rd). The corresponding random feature map, for a given m ∈ N, is defined as
ϕm(u) = 1√ m ϕ1m(u) ⋆ ... ⋆ ϕ l m(u) ∈ Rml, (5)
where ϕim = (ξi(u, ω1), ..., ξi(u, ωm)) ⊤, ⋆ stands for vertical concatenation, and ω1, ..., ωm iid∼ Ω.
Random feature maps can be used to unbiasedly approximate corresponding kernels, as follows:
K̂(x,y) = ϕm(x) ⊤ϕm(y). (6)
Using the above notation, a trigonometric random feature ϕtrigm from Equation 2 can be encoded as applying l = 2, ξ1(u, ω) = sin(ω⊤u), ξ2(u, ω) = cos(ω⊤u). Similarly, positive random features can be encoded as taking l = 2 and ξ1(u, ω) = 1√2 exp(ω ⊤u), ξ2(u, ω) = 1√2 exp(−ω ⊤u).
The following result from (Choromanski et al., 2021b) shows that the mean squared error (MSE) of the trigonometric estimator is 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 ŜM trig
m (x,y) an estimator of SM(x,y) for x,y ∈ Rd using trigonometric RFs and ω1, ..., ωm iid∼ N (0, Id). Denote by ŜM ++
m (x,y) its analogue using positive RFs. We have: Lemma 2.2 (positive versus trigonometric RFs). Take ∆ = x − y, z = x + y, f1(u) = (2m)−1 exp(u2)SM−2(x,y), f2(u) = (2m)−1 exp(u2)SM2(x,y), f3(u) = (1 − exp(−u2))2. The MSEs of these estimators are:
MSE(ŜM trig m (x,y)) = f1(∥z∥2)f3(∥∆∥2), MSE(ŜM ++ m (x,y)) = f2(∥z∥2)f3(∥z∥2). (7)

heading: 2.2 THE ALGORITHM
text: We are ready to present the mechanism of Hybrid Random Features (HRFs). Denote by E = (ŜM
k (x,y))p+1k=1 a list of estimators of SM(x,y) (the so-called base estimators) and by Λ =
(λ̂k(x,y))pk=1 a list of estimators of {λk(x,y)} p k=1 for some functions λ k : Rd × Rd → [0, 1], constructed independently from E . Take the following estimator of SM(x,y):
ŜM E,Λ (x,y) = p∑ k=1 λ̂k(x,y)ŜM k (x,y) +
( 1−
p∑ k=1 λ̂k(x,y)
) ŜM p+1 (x,y) (8)
In the next section, we explain in detail how base estimators are chosen. We call ŜM E,Λ
(x,y) a hybrid random feature (HRF) estimator of SM(x,y) parameterized by E ,Λ. The role of the λcoefficients is to dynamically (based on the input (x,y)) prioritize or deprioritize certain estimators to promote those which are characterized by lower variance for a given input. Note that if elements of E are unbiased estimators of SM(x,y), then trivially ŜM E,Λ (x,y) is also an unbiased estimator of SM(x,y). Assume that each ŜM k (x,y) is of the form ŜM k (x,y) = (ϕk1,m(x))
⊤ϕk2,m(y) for ϕkj,m(u) = 1√ m ϕ1,kj,m(u) ⋆ ... ⋆ ϕ tk,k j,m (u), tk > 0 and ϕ 1,k j,m, ..., ϕ tk,k j,m : Rd → Rm, where j ∈ {1, 2}. Assume also that λk(x,y) can be written as:
λk(x,y) = ak + Eτ∼Ω[ lk∑ i=1 f i1,k(x, τ)f i 2,k(y, τ)] (9)
for some scalars ak ∈ R, distribution Ω ∈ P(Rd) (where P(Rd stands for the set of probabilistic distributions on Rd), mappings ξik, ηik : Rd × Rd → R and that the corresponding estimator λ̂k(x,y) = λ̂kn(x,y) of λ k(x,y) is of the form:
λ̂kn(x,y) = ak + (ρ k 1,n(x)) ⊤ρk2,n(y) (10)
for ρkj,n(u) = 1√ n ρ1,kj,n(u) ⋆ ... ⋆ ρ lk,k j,n (u), ρ i,k j,n(u) = (f i j,k(u, τ1), ..., f i j,k(u, τn)) ⊤, τ1, ..., τn ∼ Ω, and where j ∈ {1, 2}. Linearization of the λ-coefficients given by Equation 9 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 approximate λ-coefficients as ŜM hyb
m,n. Furthermore, for two vectors u,v, we denote their vectorized outerproduct by u⊗ v. Finally, we denote by ∏⋆ vectors-concatenation. Estimator ŜMhybm,n(x,y) can be rewritten as a dot-product of two (hybrid) random feature vectors, as the next lemma shows.
Lemma 2.3. The HRF estimator ŜM hyb m,n(x,y) satisfies ŜM hyb m,n(x,y) = Ψ1(x) ⊤Ψ2(y), where Ψj for j ∈ {1, 2} is given as Ψj(z) = Ψ1j (z) ⋆Ψ2j (z) ⋆Ψ3j (z) ⋆Ψ4j (z) and:
Ψ1j (z) = ⋆∏ k=1,...,p √ ak m ϕ1,kj,m(z) ⋆ ... ⋆ ϕ tk,k j,m (z)
Ψ2j (z) = 1√ mn ⋆∏ k=1,...,p ⋆∏ i,j∈{1,...,lk}×{1,...,tk} ρi,kj,n(z)⊗ ϕ j,k j,m(z)
Ψ3j (z) =
√ 1− ∑p k=1 ak m ϕ1,p+1j,m (z) ⋆ ... ⋆ ϕ tp+1,p+1 j,m (z)
Ψ4j (z) = i√ mn ⋆∏ k=1,...,p ⋆∏ i,j∈{1,...,lk}×{1,...,tp+1} ρi,kj,n(z)⊗ ϕ j,p+1 j,m (z)
(11)
Bipolar estimators: A prominent special case of the general hybrid estimator defined above is the one where: E = (ŜM ++ (x,y), ŜM trig (x,y)). Thus consider the following estimator ŜM hyb
m,n:
ŜM hyb m,n(x,y) = λ̂n(x,y)ŜM ++ m (x,y) + (1− λ̂n(x,y))ŜM trig m (x,y). (12)
The question arises whether ŜM hyb m,n defined in such a way can outperform both ŜM ++ m and ŜM trig
m . That of course depends also on the choice of λ : Rd × Rd → R. If we consider a (common) normalized setting where all input vectors have the same norm, we can rewrite λ(x,y) ad λ(θx,y, r), where θx,y is an angle between x and y and ∥x∥ = ∥y∥ = r. By our previous analysis we know that ŜM ++ m becomes perfect for θ = π and ŜM trig
m becomes perfect for θ = 0. That suggests particularly simple linear dependence of λ on θ to guarantee vanishing variance for both the critical values: θ = 0 and θ = π. It remains to show that such a λ-coefficient can be linearized. It turns out that this can be done with a particularly simple random feature map mechanism, as we will show later, leading to the so-called angular hybrid variant (see: Section 2.4).

heading: 2.3 CHOOSING BASE ESTIMATORS FOR HRFS: COMPLEX EXPONENTIAL ESTIMATORS
text: In this section, we explain how we choose base estimators in Equation 8. We denote by i a complex number such that i2 = −1. The following lemma is a gateway to construct unbiased base estimators. Lemma 2.4. Let ω ∼ N (0, Id). Then for every z = (z1, ..., zd)⊤ ∈ Cd the following holds:
E [ exp(ω⊤z) ] = exp (∑d i=1 z 2 i
2
) . (13)
For a complex vector z = (z1, ..., zd) ∈ Cd, we denote z2 = ∑d i=1 z 2 i . Consider z = Ax + (A⊤)−1y for an invertible (in Cd×d) matrix A ∈ Cd×d. Using Lemma 2.4, we get:
exp
( (Ax)2
2
) exp ( ((A⊤)−1y)2
2
) SM(x,y) = E[exp(ω⊤(Ax+ (A⊤)−1y)] (14)
Thus for ΨmM(u) def = 1√ m exp(− (Mu)
2
2 )(exp(ω ⊤ 1 Mu), ..., exp(ω ⊤ mMu)) ⊤ and ωi ∼ N (0, Id):
SM(x,y) = E[ΨmA(x)⊤Ψm(A⊤)−1(y)]. (15)
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.e. one using A, the other one (A⊤)−1. Asymmetric random features is a simple generalization of the model using the same ϕ for both x and y, presented earlier. The resulting estimators ŜM cexp
m (x,y), that we call complex exponential (CE) will serve as base estimators, and are defined as:
ŜM cexp m (x,y) def = ΨmA(x) ⊤Ψm(A⊤)−1(y) (16) For A = iId, CE-estimator becomes trigonometric and for A = Id, it becomes ŜM ++
m (x,y). Note also that an estimator based on this mechanism has variance equal to zero if
x = −(A−1)(A⊤)−1y. (17) In fact we can relax that condition. It is easy to check that for the variance to be zero, we only need:
Re(A)x+Re((AT )−1)y = 0 and Im(A)x+ Im((AT )−1)y = 0 (18)

heading: 2.4 ADAPTING HRF ESTIMATORS: HOW TO INSTANTIATE HYBRID RANDOM FEATURES
text: We will use complex exponential estimators from Section 2.3 as base estimators in different insantiations 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 λ-coefficients (Equation 10) and consequently - desired random feature map decomposition via the mechanism of compositional kernels (details in the Appendix). The questions remains how to construct λ-coefficients to reduce variance of the estimation in the desired regions of interest.
Accurate approximation of small and large softmax kernel values simultaneously: As we explained earlier, in several applications of softmax estimators a desired property is to provide accurate estimation in two antipodal regions - small and large softmax kernel values. This can be done in particular by choosing p = 1, ŜM 1 = ŜM ++ , ŜM 2 = ŜM
trig and λ(x,y) = θx,yπ , where θx,y
stands for an angle between x and y. The key observation is that λ(x,y) can be unbiasedly approximated via the mechanism of sgn random features (following from Goemans-Willimson algorithm (Goemans & Williamson, 2004)) as follows (details in the Appendix, Sec. D):
λ̂(x,y) = 1
2 + ρn(x)
⊤ρn(y), (19)
where ρ1,n(z) = ρ2,n(z) = ρn(z) def = i√ 2n (sgn(τ⊤1 z), ..., sgn(τ ⊤ n z)) ⊤ for τ1, ..., τn ∼ N (0, Id) and i2 = −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 (θx,y = π) and largest (θx,y = 0) softmax kernel values (see: Fig. 1).
Gaussian Lambda Coefficients: Another tempting option is to instantiate λ-coefficients with approximate Gaussian kernel values (estimated either via positive or trigonometric random features). In that setting, functions λ̂k(x,y) are defined as λ̂k(x,y) = exp(−∥x+Mky∥ 2 2
2c2k ), for some
M1, ...,Mp ∈ Rd×d and c1, ..., cp ∈ R. Since this time λ-coefficients are the values of the Gaussian kernel between vectors xck and − Mky ck , we can take ak = 0 and define ρk1,·,ρ k 2,· as:
ρk1,·(x) = exp(− ∥x∥22 2c2k )ϕSM( x ck ) and ρk2,·(y) = exp(− ∥Mky∥22 2c2k )ϕSM(−Mky ck ), (20)
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 pseudo-combinatorial properties. The coefficients λk ∈ (0, 1] fire if x ≈ −Mky (the firing window can be accurately controlled via hyperparameters ck). Thus matrices Mk can be coupled with the base estimators particularly accurate in the regions, where x ≈ −Mky. 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 θx,y = 0 and θx,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 X1, ..., Xu ⊆ Rd of centers x1, ...,xu ∈ Rd and that the analogous fact is true for y with the corresponding clusters Y1, ..., Yw ⊆ Rd and centers y1, ...,yw ∈ Rd. For a fixed pair of clusters’ centers (xi,yj), one can associate a complex exponential estimator ŜM cexp
(xi,yj) with the corresponding matrix A ∈ C d×d (see: Sec. 2.3) that minimizes the following loss, where |z| for z ∈ Cd is defined as |z| = √∑d
i=1 |zi|2:
l(A) = |Axi + (A⊤)−1yj |2 (21)
Note that, by Equation 17, if one can find A such that l(A) = 0, then ŜM cexp
(xi,yj) makes no error in approximating softmax kernel value on the centers xi,yj . Furthermore, if sampled points are close to the centers and l(A) is small, the corresponding variance of the estimator will also be small. Thus one can think about ŜM cexp
(xi,yj) as an estimator adapted to a pair of clusters (xi,yj). We can add additional constraints regarding A, i.e. A ∈ Rd×d or A being diagonal (or both) for closedform formulae of the optimal A, yet not necessarily zeroing out l(A) (details in the Appendix, Sec. H.3). Complex exponential estimators defined in this way can be combined together, becoming base estimators in the HRF estimator. The corresponding λ-coefficients can be chosen to be Gaussian as described in the previous paragraph (with optimized matrices Mk so that a coefficient fires for the corresponding pair of clusters’ centers), but there is another much simpler choice of λs. We index different base estimators by pairs of clusters’ centers (xi,yj). One can simply take ρ (xi,yj) 1 (x) = 1 if x belongs to the cluster of center xi and ρ (xi,yj) 1 (x) = 0 otherwise. Vectors ρ (xi,yj) 2 (y) (that also become scalars in this scenario) are defined analogously (the n symbol vanishes since ρ is deterministic). 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 cluster-membership calculation is computationally acceptable, otherwise aforementioned Gaussian coefficients can be applied. 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 ŜM ++ ).

heading: 3 THEORETICAL GUARANTEES
text: We provide here theoretical guarantees regarding HRF estimators. We focus on the bipolar setting with two base estimators: ŜM trig and ŜM ++
. The following is true: Theorem 3.1 (MSE of the bipolar hybrid estimator). Take the bipolar hybrid estimator ŜM hyb m,n(x,y), where ŜM trig m (x,y) and ŜM ++
m (x,y) are chosen independently i.e. their random projections are chosen independently (note that we always assume that λ̂n(x,y) is constructed independently from ŜM trig m (x,y) and ŜM ++ m (x,y)). Then the following holds:
MSE(ŜM hyb m,n(x,y)) = E[λ̂2n(x,y)]MSE(ŜM ++ m (x,y)) + E[(1− λ̂n(x,y))2]MSE(ŜM trig
m (x,y)) (22)
Furthermore, if ŜM trig m (x,y) and ŜM ++
m (x,y) apply the exact same sets of random projections, the mean squared error of the hybrid estimator is further reduced, namely we have:
MSE(ŜM hyb m,n(x,y)) = E[λ̂2n(x,y)]MSE(ŜM ++ m (x,y)) + E[(1− λ̂n(x,y))2]MSE(ŜM trig m (x,y))
− 2 m SM2(x,y)(1− cos(∥x∥22 − ∥y∥22))E[λ̂n(x,y)(1− λ̂n(x,y))]
(23)
The exact formula on MSE(ŜM ++ m (x,y)) and MSE(ŜM trig m (x,y)) 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:
E[λ̂2n(x,y)] = θx,y π ( θx,y π − θx,y nπ + 1 n ) ,E[λ̂n(x,y)] = θx,y π . (24)
We see that, as mentioned before, the variance of the angular hybrid estimator is zero for both θx,y = 0 and θx,y = π if inputs x,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 ⊆ Rd. For a given estimator ŜM on feature vectors x,y ∈ C, we define its max-relative-error with respect to C as ϵC(ŜM) = maxx,y∈C ϵx,y(ŜM),where ϵx,y(ŜM) = √ MSE(ŜM(x,y))
SM(x,y) . Denote
by S(r) a sphere centered at 0 and of radius r. Define ϵθ,r(ŜM) def = ϵx,y(ŜM) for x,y ∈ S(r) and such that θ = θx,y (note that the mean squared errors of the considered estimators depend only on the angle θx,y for x,y chosen from a fixed sphere).
This definition captures the critical observation that in several applications of the softmax kernel estimation, e.g. efficient softmax sampling or linear-attention Transformers (Choromanski et al., 2021b), 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:
ϵθ,r(ŜM trig m ) = ϵπ−θ,r(ŜM ++ m ) = 1√ 2m exp(2r2 sin2( θ 2 ))
( 1− exp(−4r2 sin2(θ
2 ))
) , (25)
and consequently for W (r) = exp(2r2) ( 1− exp(−4r2) ) :
ϵS(r)(ŜM trig m ) = ϵS(r)(ŜM ++
m ) = lim θ→π
ϵθ,r(ŜM trig
m ) = lim θ→0
ϵθ,r(ŜM ++ m ) =
√ 1
2m W (r) (26)
Our main result shows that HRF estimators can be applied to reduce the max-relative-error of the previously applied estimators of the softmax kernel. In the next theorem we show that the maxrelative-error of the angular hybrid estimator scales as 1r exp(2r 2) in the length r of its inputs as opposed to exp(2r2) as it is the case for ŜM trig and ŜM ++
(see: Lemma 3.4). Furthermore, the max-relative-error scales as √ θ and √ π − θ as θ → 0 and θ → π respectively, in particular goes to 0 in both critical cases. This is not true for ŜM trig nor for ŜM ++
. Theorem 3.5. The max-relative-error of the angular hybrid estimator for the inputs x,y on the sphere S(r) of radius r ≥ 1 satisfies for W (r) = exp(2r2) ( 1− exp(−4r2) ) :
ϵS(r)(ŜM anghyb m,n ) ≤ 1
r
√ 1
2m W (r)
√ 1
π − 1 nπ + 1 n √ π
(27)
Furthermore, limθ→0 ϵθ,r(ŜM
anghyb
m,n )√ θ = limθ→π ϵθ,r(ŜM
anghyb
m,n )√ θ−π =
√ 1
2πmnW (r).
Additional implications of Theorem 3.5, using the fact that Θ(mn)-dimensional HRFs can be constructed in time O(nd +md +mn) (regular RFs need O(mnd)), are in the Appendix (Sec. D.4).

heading: 4 EXPERIMENTS
text: In this section we conduct exhaustive evaluation of the mechanism of HRFs on several tasks.

heading: 4.1 POINTWISE SOFTMAX KERNEL ESTIMATION
text: We start with ablation studies regarding empirical relative errors of different softmax kernel estimators across different inputs’ lengths r and angles θx,y. The results are presented in Fig. 2. Ablations over more lengths are presented in the Appendix (Sec. G). The angular hybrid estimator most accurately approximates softmax kernel and has smaller max-relative-error than other estimators.
Comparison with QMC-methods: Even though, as we explained in Section 1, 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 as well as wellestablished QMC baselines: orthogonal random features (ORF)(Yu et al., 2016) (regular variant ORF-reg as well as the one applying Hadamard matrices ORF-Had) and QMCs based on random Halton sequences (Halton-R)(Avron et al., 2016). For HRFs, we tested two sub-variants: with orthogonal (HRF-ort) and regular iid random features (HRF-iid). 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 1.

heading: 4.2 LANGUAGE MODELING
text: In this section, we apply HRFs to the language modeling task, training a 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 results are in the Appendix (Sec. I, Table 4).
Statistical Metrics. We compute 1d-Wasserstein distance and the Kolmogorov-Smirnov (KS) metric (Ramdas et al., 2017) between a true softmax distribution induced by a query on all the classes
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. I) we present additional results for the WikiText2 dataset. We see that HRFs lead to most accurate softmax distribution estimation.

heading: 4.3 TRAINING SPEECH MODELS WITH HRF-CONFORMERS-PERFORMERS
text: 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-Transducer encoder (Gulati et al., 2020) 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 as 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 clustering-based 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. J and Sec. H.1 respectively).

heading: 4.4 DOWNSTREAM ROBOTICS EXPERIMENTS
text: In Robotics, we leverage the accuracy of HRFs to obtain inference improvements, critical for onrobot 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 Vision-Performer stack used to process high-resolution (500 x 500 x 3) input. The 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). Videos of the HRF-trained robotic policies are included in the supplementary material.

heading: 5 CONCLUSION
text: We presented a new class of random feature techniques, called hybrid random features for softmax/Gaussian kernel estimation that more accurately approximate these kernels than previous algorithms. We also demonstrated their robustness in a wide range of applications from softmax sampling to Transformers/attention training, also for downstream Robotics tasks.
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 can be found in the following github: https://github. com/HL-hanlin/HRF_ICLR2022. Ethics Statement: 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.

heading: ACKNOWLEDGEMENTS
text: AW acknowledges support from a Turing AI Fellowship under grant EP/V025379/1, The Alan Turing Institute, and the Leverhulme Trust via CFI. AS acknowledges AWS compute resources from Onur Kara.

heading: APPENDIX: HYBRID RANDOM FEATURES

heading: A PROOF OF LEMMA 2.4
text: Proof. From the independence of ω1, ..., ωd, we get:
E[exp(ω⊤z)] = d∏
i=1
E[exp(ωizi)] (28)
Thus it suffices to show that for g ∼ N (0, 1) and any z ∈ C the following holds:
E[exp(gz)] = exp( z2
2 ) (29)
Define f(z) = E[exp(gz)]. Note that f(z) = exp( z 2
2 ) for z = ix, where i 2 = −1 and x ∈ R
which follows from the formula of the characteristic function of the Gaussian distribution. Similarly, f(z) = exp( z 2
2 ) for z ∈ R 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 → 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(z) and g(z) def= exp( z 2
2 ) are analytic).

heading: B PROOF OF LEMMA 2.3
text: Proof. The following is true:
ŜM hyb
m,n(x,y) = p∑ k=1 ak m (ϕ1,k1,m(x) ⋆ ... ⋆ ϕ tk,k 1,m (x)) ⊤(ϕ1,k2,m(y) ⋆ ... ⋆ ϕ tk,k 2,m (y))+
1
mn p∑ k=1 lk∑ i=1 tk∑ j=1 (ρi,k1,n(x)⊗ ϕ j,k 1,m(x)) ⊤(ρi,k2,n(y)⊗ ϕ j,k 2,m(y))+
1− ∑p
k=1 ak m (ϕ1,p+11,m (x) ⋆ ... ⋆ ϕ tp+1,p+1 1,m (x)) ⊤(ϕ1,p+12,m (y) ⋆ ... ⋆ ϕ tp+1,p+1 2,m (y))−
1
mn p∑ k=1 lk∑ i=1 tp+1∑ j=1 (ρi,k1,n(x)⊗ ϕ j,p+1 1,m (x)) ⊤(ρi,k2,n(y)⊗ ϕ j,p+1 2,m (y)).
(30)
The statement of the lemma follows directly from that. Note that the key trick is an observation that a product of two functions of x,y that can be linearized on expectation (i.e. rewritten with x and y disentangled) is itself a function that can be linearized 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)).

heading: C BIPOLAR HRF ESTIMATORS: PROOF OF THEOREM 3.1
text: Proof. The following holds:
Var(ŜM hyb m,n(x,y)) = Var ( λ̂n(θ)ŜM ++ m (x,y) ) +Var ( (1− λ̂n(θ))ŜM trig m (x,y) ) +
2Cov ( λ̂n(θ)ŜM ++ m (x,y), (1− λ̂n(θ))ŜM trig m (x,y) ) (31)
We will focus now on the covariance term. To simplify the notation, we will drop index n in λ̂n. Assume first easier to analyze case where ŜM trig m (x,y) and ŜM ++
m (x,y) are chosen independently. Then the following is true:
Cov ( λ̂(θ)ŜM ++ m (x,y), (1− λ̂(θ))ŜM trig m (x,y) ) =
E[λ̂(θ)(1− λ̂(θ))ŜM ++ m (x,y)ŜM trig m (x,y)]− E[λ̂(θ)ŜM ++ m (x,y)]E[(1− λ̂(θ))ŜM trig m (x,y)] = ( E[λ̂(θ)(1− λ̂(θ))]− E[λ̂(θ)]E[(1− λ̂(θ))] ) E[ŜM ++ m (x,y)]E[ŜM trig m (x,y)]
= −(SM(x,y))2Var(λ̂(θ)) (32)
Now assume that ŜM trig m (x,y) and ŜM ++
m (x,y) use the exact same random projections. Then, using similar analysis as before, we get:
Cov ( λ̂(θ)ŜM ++ m (x,y), (1− λ̂(θ))ŜM trig m (x,y) ) =
E[ŜM ++ m (x,y)ŜM trig m (x,y)]E[λ̂(θ)(1− λ̂(θ))]− (SM(x,y))2E[λ̂(θ)]E[(1− λ̂(θ))] (33)
This time however it is no longer the case that E[ŜM ++ m (x,y)ŜM trig
m (x,y)]E[λ̂(θ)(1 − λ̂(θ))] = E[ŜM ++ m (x,y)]E[ŜM trig m (x,y)] = (SM(x,y)) 2 since ŜM ++ m (x,y) and ŜM trig m (x,y) are no longer independent. In order to compute E[ŜM ++ m (x,y)ŜM trig
m (x,y)]E[λ̂(θ)(1 − λ̂(θ))], we will first introduce useful denotation.
Denote by ω1, ..., ωm iid∼ N (0, Id) the random projections sampled to construct both ŜM
++ m (x,y)
and ŜM trig
m (x,y). Denote: Yi = cosh(ω ⊤ i (x+ y)) def =
exp(ω⊤i (x+y))+exp(−ω ⊤ i (x+y))
2 . We have:
ŜM ++
m (x,y) = exp
( −∥x∥ 2 + ∥y∥2
2
) Y1 + ...+ Ym
m (34)
If we denote: Zi = cos(ω⊤i (x− y)), then we can write ŜM trig m (x,y) as:
ŜM trig
m (x,y) = exp
( ∥x∥2 + ∥y∥2
2
) Z1 + ...+ Zm
m (35)
We can then rewrite E[ŜM ++ m (x,y)ŜM trig m (x,y)] as:
E[ŜM ++ m (x,y)ŜM trig m (x,y)] = 1
m2 ∑ i ̸=j E[YiZj ] + m∑ i=1 E[YiZi]  = 1
m2
[( m
2
) (SM(x,y))2 +mE[cosh(ω⊤(x+ y)) cos(ω⊤(x− y))] ] ,
(36)
where ω ∼ N (0, Id). The equality follows from the unbiasedness of ŜM trig m (x,y) and ŜM ++
m (x,y) and the fact that different ωi are chosen independently. Thus it remains to compute ρ = E[cosh(ω⊤(x + y)) cos(ω⊤(x − y))]. Note first that ρ = E[exp(ω⊤(x + y)) cos(ω⊤(x − y))] since −ω ∼ N (0, Id) and cos is an even function. Denote z = x+ y + i(x− y). We have:
E[exp(ω⊤(x+ y)) cos(ω⊤(x− y))] = Re [ E[exp(ω⊤z)] ] = Re[ d∏ i=1 exp( z2i 2 )] =
Re [ exp (∑d j=1(xj + yj) 2 + 2i(x2j − y2j )− (xj − yj)2
2
)] = (SM(x,y))2 cos(∥x∥22 − ∥y∥22)
(37)
Thus we conclude that:
E[ŜM ++ m (x,y)ŜM trig m (x,y)] = (1− 1
m )(SM(x,y))2 +
1 m (SM(x,y))2 cos(∥x∥22 − ∥y∥22) (38)
Therefore we get the formulae for the covariance term in both: the setting where random projections of ŜM trig m (x,y) and ŜM ++
m (x,y) are shared (variant II) and when they are not (variant I). The following is true for Z = 1m (1− cos(∥x∥ 2 2 − ∥y∥22)E[λ̂(θ)(1− λ̂(θ))]:
Cov(λ̂(θ)ŜM ++ m (x,y), (1−λ̂(θ))ŜM trig m (x,y)) =
{ −(SM(x,y))2Var(λ̂(θ)) for variant I
−(SM(x,y))2(Var(λ̂(θ)) + Z) for variant II (39)
We also have the following:
Var(λ̂(θ)ŜM ++ m (x,y)) = E[(λ̂(θ))2(ŜM ++ m (x,y)) 2]− (E[λ̂(θ)ŜM ++ m (x,y)]) 2 = E[(λ̂(θ))2] ( MSE(ŜM ++ m (x,y)) + (SM(x,y)) 2 ) − (E[λ̂(θ)])2(SM(x,y))2 =
(SM(x,y))2Var(λ̂(θ)) + E[(λ̂(θ))2]MSE(ŜM ++
m (x,y))
(40)
and furthermore (by the analogous analysis):
Var((1−λ̂(θ))ŜM trig
m (x,y)) = (SM(x,y)) 2Var(λ̂(θ))+E[(1−λ̂(θ))2]MSE(ŜM
trig m (x,y)) (41)
By putting the derived formula for the above variance terms as well as 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).

heading: D ANGULAR HYBRID ESTIMATORS
text: D.1 PROOF OF LEMMA 3.2
Proof. The formula for the expectation is directly implied by the formula (1) from Sec. 2.1 of Cho & Saul (2012) for the zeroth-order arc-cosine kernel k0(x,y) def = 1− θx,yπ . It also follows directly from Goemans-Williamson algorithm (Goemans & Williamson, 2004). We will now derive the formula for c = E[λ̂2(x,y)]. Since λ depends only on the angle θ = θx,y, we will refer to it as: λ(θ) and to its estimator as λ̂(θ). Denote:
ϕangn (z) = 1√ n (sgn(τ⊤1 z), ..., sgn(τ ⊤ n z)) ⊤ (42)
Denote: Xi = (ϕangn (x))[i]ϕ ang n (y)[i]. We have:
λ̂(θ) = 1
2
( 1−
n∑ i=1 Xi
) . (43)
Note first that by the construction of λ̂(θ), we have: E[λ̂(θ)] = θπ and thus: E[ ∑n i=1 Xi] = 1− 2θ π . Therefore we conclude that:
c = 1
4 E 1− 2 n∑ i=1 Xi + ( n∑ i=1 Xi )2 = 1 4 1− 2(1− 2θ π ) + n∑ i=1 E[X2i ] + ∑ i ̸=j E[Xi]E[Xj ]  = 1
4
( 1− 2(1− 2θ
π ) + n · 1 n2 + n(n− 1) · 1 n2 (1− 2θ π )2 ) = 1 4 ( 4 θ π + (1− 1 n )(1− 2θ π )2 )
= θ
π
( θ
π − θ nπ + 1 n ) (44)
D.2 EXPLICIT FORMULA FOR THE MSE OF THE ANGULAR HYBRID ESTIMATOR
Theorem 3.1 and Lemma 3.2 immediately imply the following result. Theorem D.1 (MSE of the angular hybrid estimator). Take the angular hybrid estimator ŜM anghyb m,n (x,y), where ŜM trig m (x,y) and ŜM ++
m (x,y) are chosen independently i.e. their random projections are chosen independently (note that we always assume that λ̂(x,y) is chosen independently from ŜM trig m (x,y) and ŜM ++ m (x,y)). Then the following holds:
MSE(ŜM anghyb m,n (x,y)) = θ
π
( θ
π − θ nπ + 1 n
) 1
2m exp(∥z∥2)SM2(x,y)(1− exp(−∥z∥2))2+
= ( 1− θ
π
)( 1− θ
π +
θ
nπ
) 1
2m exp(∥z∥2)SM−2(x,y)(1− exp(−∥∆∥2))2
(45)
for ∆ = x − y and z = x + y. Furthermore, if ŜM trig m (x,y) and ŜM ++
m (x,y) apply the exact same sets of random projections, the mean squared error of the hybrid estimator is further reduced, namely we have:
MSE(ŜM anghyb m,n (x,y)) = θ
π
( θ
π − θ nπ + 1 n
) 1
2m exp(∥z∥2)SM2(x,y)(1− exp(−∥z∥2))2+(
1− θ π
)( 1− θ
π +
θ
nπ
) 1
2m exp(∥z∥2)SM−2(x,y)(1− exp(−∥∆∥2))2
− 2 m SM2(x,y)(1− cos(∥x∥22 − ∥y∥22)) θ π
( 1− 1
n − θ π + θ nπ ) (46)
Thus if ∥x∥2 = ∥y∥2 = r then regardless of whether the same sets of random projections are applied or not, we get:
MSE(ŜM anghyb m,n (x,y)) = θ
π
( θ
π − θ nπ + 1 n
) 1
2m exp(8r2 cos2(
θ 2 )− 2r2)·
(1− exp(−4r2 cos2(θ 2 )))2 + θ π
( 1− 1
n − θ π + θ nπ
) 1
2m exp(2r2)(1− exp(−4r2 sin2(θ 2 )))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:
ϵS(r)(ŜM anghyb m,n ) = exp(r2)√
2m
√ max θ∈[0,π] hr(θ), (48)
where:
hr(θ) = ar(θ) + ar(π − θ) (49)
and ar(θ) is defined as:
ar(θ) = θ π ( θ π − θ nπ + 1 n ) exp(2r2 cos(θ))
( 1− exp(−4r2 cos2(θ
2 ))
)2 . (50)
Therefore we have:
ϵS(r)(ŜM anghyb m,n ) ≤ exp(r2)√
m √ max θ∈[0,π] ar(θ) (51)
Notice that: ar(θ) ≤ br(θ)(1− exp(−4r2))2 (52)
where: br(θ) = b 1 r(θ) + b 2 r(θ) (53)
and
b1r(θ) = (1− 1 n ) θ2 π2 exp(2r2 cos(θ)), (54)
b2r(θ) = 1 n · θ π exp(2r2 cos(θ)). (55)
Therefore: max θ∈[0,π] br(θ) ≤ max θ∈[0,π] b1r(θ) + max θ∈[0,π] b2r(θ) (56)
Denote: b1 = maxθ∈[0,π] b1r(θ) and b 2 = maxθ∈[0,π] b 2 r(θ). Note that:
db1r(θ)
dθ = exp(2r2 cos(θ))(1− 1 n ) 2θ π2 (1− r2θ sin(θ)) (57)
and db2r(θ)
dθ = exp(2r2 cos(θ))
1
nπ (1− 2r2θ sin(θ)) (58)
Thus, from the properties of function: θ → θ sin(θ) and the fact that r ≥ 1, we conclude that both derivatives are first non-negative, then non-positive and then non-negative and that the unique local maximum on the interval [0, π] is achieved for θ ≤ π2 . Note also that b 1 r(θ), b 1 r(θ) ≥ 0 and b1r(0) = b 2 r(0) = 0, b 1 r(π) = (1 − 1n ) exp(−2r 2), b2r(π) = 1 n exp(−2r
2). We conclude that global maximum for bir on the interval [0, π] for i = 1, 2 is achieved either in its unique local maximum on that interval or for θ = π. Let us consider first: b1r . In its local maximum on [0, π] we have:
θ∗ sin(θ∗) = 1
r2 (59)
Since θ ≤ sin(θ) · π2 on [0, π 2 ], we get:
(θ∗)2 ≤ π 2 1 r2 , (60)
i.e.:
θ∗ ≤ √ π
2
1 r (61)
Therefore: b1r(θ
∗) ≤ (1− 1 n ) 1 2πr2 exp(2r2) ≥ b1r(π) (62)
We thus conclude that: max θ∈[0,π] b1r(θ) ≤ (1− 1 n ) 1 2πr2 exp(2r2) (63)
By the completely analogous analysis applied to b2r , we obtain:
max θ∈[0,π]
b1r(θ) ≤ 1
2n √ πr2 exp(2r2) (64)
Now, using Equation 51, Equation 52, and Equation 56, we obtain:
ϵS(r)(ŜM) anghyb m,n ≤ exp(2r2)√ 2mr (1− exp(−4r2))
√ 1
π − 1 nπ + 1 n √ π
(65)
and that completes the first part of the proof (proof of Inequality 27). The equations on the limits are directly implied by the fact that:
ϵS(r)(ŜM anghyb m,n ) = exp(2r2)√
2m
√ hr(θ). (66)
D.4 DISCUSSION OF THEOREM 3.5
Theorem 3.5 leads to yet another important conclusion not discussed in the main body of the paper. Note that asymptotically as θ → 0 or θ → π, the relative error decreases (as a function of m and n) as 1√
mn = Θ( 1dest ), where dest stands for the number of random features of the resulting
estimator. Note also that for the regular estimator the rate of decrease is also Θ( 1dest ) (here we treat all other arguments of the formula for the relative error as constants; we already know that the dependence on them of the relative error for the HRF estimators is superior to the one for regular estimators). The key difference from the computational point of view is that for the HRF estimator, the random feature map can be constructed in time O(nd + md + mn), whereas for the regular estimator it requires time O(mnd) (see: Sec. L). 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., 2021b). 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 a prominent case where they all have the same fixed length) that corresponds to angle values θ → π or θ → 0. This is exactly the scenario from the second part of Theorem 3.5.

heading: E PROOF OF LEMMA 3.4
text: Proof. Equation 25 follows directly from Lemma 2.2. Notice that the relative error ϵθ,r(ŜM trig m ) is an increasing function of sin2( θ2 ) 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 ŜM ++
m .

heading: F GAUSSIAN HYBRID ESTIMATORS
text: In this section we will provide more results regarding Gaussian hybrid estimators. We will focus on the bipolar scenario and within this scenario on the so-called normalized variant described below.
If x and y have fixed L2-norm (∥x∥2 = ∥y∥2 = r), we can propose a normalized bipolar Gaussian hybrid estimator such that λ(x,y) = 0 if θx,y = 0 and λ(x,y) = 1 if θx,y = π. Furthermore the variance of that estimator, that we call shortly ŜM gausshyb
m,n , will be zeroed out for θx,y = 0 or θx,y = π, but not for both.
The coefficient-function λ : Rd × Rd → R is defined as:
λ(x,y) = 1− exp(−σ
2
2 ∥x− y∥ 2) ρ (67)
where ρ is given as: ρ = 1− exp(−2σ2r2) (68)
The hyperparameter σ 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 θx,y = 0 and in the latter it is zero for θx,y = π, but the variance never zeroes out for both θx,y = 0 and θx,y = π. 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 θx,y = 0 and θx,y = π. However such nested constructions are not the topic of this paper. From now on we will assume that trigonometric random features are applied to estimate λ-coefficient, Therefore we have:
λ̂(x,y) = 1 ρ − 1 nρ n∑ i=1 cos(σ(ω⊤i (x− y))) (69)
and that leads to the following choice of: a, ρ1,2, ρ2,n from Equation 10:{ a = 1ρ ρ1,n(z) = ρ2,n(z) = i√ nρ (sin(τ ⊤ 1 z), cos(τ ⊤ 1 z), ..., sin(τ ⊤ n z), cos(τ ⊤ n z)) ⊤ (70)
where τ1, ..., τn ∼ N (0, Id) and i2 = −1. Using Theorem 3.1, we conclude that:
Theorem F.1 (MSE of the normalized bipolar Gaussian estimator). Take the normalized bipolar Gaussian estimator ŜM gausshyb m,n (x,y), where ŜM trig m (x,y) and ŜM ++
m (x,y) are chosen independently i.e. their random projections are chosen independently (note that we always assume that λ̂(x,y) is chosen independently from ŜM trig m (x,y) and ŜM ++
m (x,y)). Denote: ∆ = x − y. Then the following holds:
MSE(ŜM gausshyb m,n (x,y)) = E[λ̂2(x,y)]MSE(ŜM ++ m (x,y))+E[(1−λ̂(x,y))2]MSE(ŜM trig
m (x,y)) (71)
where
E[(λ̂(x,y))2] = 1 ρ2 (1− exp(−σ
2
2 ∥∆∥2))2 + 1 2ρ2n (1− exp(−σ2∥∆∥2))2 (72)
E[(1− λ̂(x,y))2] = 1 ρ2 ((1− ρ)− exp(−σ 2 2 ∥∆∥2))2 + 1 2ρ2n (1− exp(−σ2∥∆∥2))2. (73)
Furthermore, if ŜM trig m (x,y) and ŜM ++
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,y))2] and E[(1− λ̂(x,y))2] as follows:
E[(1− λ̂(x,y))2] = 1 ρ2 E[(ρ− 1 + 1 n n∑ i=1 cos(σω⊤(x− y)))2]
= (1− ρ)2
ρ2 + 2(ρ− 1) ρ2 E[cos(σω⊤(x− y))] + 1 ρ2 E[( 1 k n∑ i=1 cos(σω⊤(x− y)))2]
= (1− ρ)2
ρ2 + 2(ρ− 1) ρ2 exp(−σ 2 2 ∥∆∥2) + 1 ρ2 exp(−σ2∥∆∥2) + 1 2ρ2n (1− exp(−σ2∥∆∥2))2
= 1 ρ2 ((1− ρ)− exp(−σ
2
2 ∥∆∥2))2 + 1 2ρ2n (1− exp(−σ2∥∆∥2))2
(74)
E[(λ̂(x,y))2] = 1− 2E[1− λ̂(x,y)] + E[(1− λ̂(x,y))2]
= −1 + 2 ρ − 2 ρ exp(−σ
2
2 ∥∆∥2) + E[(1− λ̂(x,y))2]
= 1 ρ2 (1− exp(−σ
2
2 ∥∆∥2))2 + 1 2ρ2n (1− exp(−σ2∥∆∥2))2
(75)
By plugging in these two formulae in the expression for MSE, we obtain the first half of the theorem.
Note that if ŜM trig m (θ, r) and ŜM ++ m (θ, r) apply the exact same sets of random projections, then (1 − cos(∥x∥22 − ∥y∥22)) in MSE(ŜM gausshyb
m,n (x,y)) becomes zero in Theorem F.1, therefore the MSE remains the same. We thus obtain the 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 θ. We can rewrite Eq. 67 and provide equivalent definition for λ by replacing x and y with their angle θ and norm r:
λ(θ, r) = 1− exp(−2σ2r2 sin( θ2 ) 2)
ρ (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). Assume that ∥x∥2 = ∥y∥2 = r and denote: θ = θx,y. Take the normalized bipolar Gaussian hybrid estimator ŜM gausshyb m,n (θ, r), where ŜM trig m (θ, r) and ŜM ++
m (θ, r) are chosen independently i.e. their random projections are chosen independently (note that we always assume that λ̂(θ, r) is chosen independently from ŜM trig m (θ, r) and ŜM ++ m (θ, r)). Then the following holds:
MSE(ŜM gausshyb m,n (θ, r)) = E[λ̂2(θ, r)]MSE(ŜM ++ m (θ, r)) + E[(1− λ̂(θ, r))2]MSE(ŜM trig
m (θ, r)) (77)
where
E[(λ̂(θ, r))2] = 1 ρ2 (1− exp(−2σ2r2 sin(θ 2 )2))2 + 1 2ρ2n (1− exp(−4σ2r2 sin(θ 2 )2))2 (78)
E[(1−λ̂(θ, r))2] = 1 ρ2 ((1−ρ)−exp(−2σ2r2 sin(θ 2 )2))2+ 1 2ρ2n (1−exp(−4σ2r2 sin(θ 2 )2))2 (79)
Furthermore, if ŜM trig m (θ, r) and ŜM ++
m (θ, r) apply the exact same sets of random projections, the mean squared error of the hybrid estimator will remains the same.
Proof. For ∥x∥2 = ∥y∥2 = r2, the following holds for ∆ = x− y:
∥∆∥2 = 4r2 sin(θ 2 )2 (80)
And this theorem holds trivially by replacing ∥∆∥2 with the above formula in Theorem F.1.

heading: G POINTWISE SOFTMAX KERNEL ESTIMATION EXPERIMENTS
text: In Fig. 5 we present complete ablation studies regarding empirical relative errors of different RFbased estimators of the softmax kernel over different lengths of inputs r and angles θx,y.

heading: H COMPLEX EXPONENTIAL ESTIMATORS FOR CLUSTERED DATA
text: 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-element 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 s = l(A) (see: Eq. 21) for all pairs of clusters and with different concentrations around centers controlled by the standard deviation σ (small/large values of s 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 s and less concentrated clusters. Additional experimental details are given in Sec. H.3.
H.2 INSTANTIATING COMPLEX EXPONENTIAL ESTIMATORS FOR CLUSTERED DATA
Assume that the inputs x (queries) can be modeled by nq clusters and that the inputs y (keys) can be modeled by nk clusters (e.g. via k-means clustering algorithm). Denote the center of each cluster as ri ∈ Rd, (i = 1, ..., nk) and rj ∈ Rd, (j = 1, ..., nq). Then there exist nqnk pairs of (ri, rj), (i = 1, ..., nq, j = 1, ..., nk), so we can construct nqnk softmax kernel estimators to estimate crossgroup softmax kernel values.
Consider z = Ari + (A⊤)−1rj for an invertible (in Cd×d) matrix A ∈ Cd×d. An estimator based on this mechanism has variance equal to 0 if:
z = Ari + (A ⊤)−1rj = 0 (81)
From now on we constrain A to be diagonal, so A = A⊤. We can rewrite the above equation as: z = Ari + (A) −1rj = 0 (82)
Since position (k, k) in matrix A is a complex number of the form αk + βki, we need to satisfy the following equation for each k = 1, ..., d:
(αk + βki)ri,k + 1
αk + βki rj,k = 0 (83)
(αk + βki)ri,k + αk − βki α2k + β 2 k rj,k = 0 (84)
where ri,k is the k-th entry of vector ri.
We can simplify this equation and separate into real and imaginary part:{ Re : (α2k + β 2 k)αkri,k + αkrj,k = 0
Im : (α2k + β 2 k)βkri,k − βkrj,k = 0
(85)
Our goal now is to choose values for αk and βk given ri and rj . If ri,krj,k > 0, we can set: { αk = 0
βk = √ rj,k/ri,k (86)
And if ri,krj,k < 0, we can set: { αk = √ −rj,k/ri,k
βk = 0 (87)
When ri,k = rj,k = 0, we can take αk = βk = 1. If ri,k = 0, rj,k ̸= 0 or the opposite, then we cannot satisfy the above equation perfectly. We can take αk to some large positive value and set βk = 0 when ri,k = 0, rj,k ̸= 0, and set αk = βk to some small positive value close to zero when ri,k ̸= 0, rj,k = 0.
When restricting A = diag(a1, . . . , ad) to Rd×d, it is easily seen that to minimize |Ari + (A⊤)−1rj |2, we can take
ak = √ |rj,k/ri,k| (88)
if ri,k ̸= 0. And if ri,k = 0, rj,k ̸= 0, ak can be set to a large positive number. We can set ak = 1 if ri,k = 0, rj,k = 0.
In the analysis below we denote matrix A calculated from Eq. 85 and Eq. 88 (depending on whether we are using real or complex matrices) given ri and rj as Ai,j .
From the analysis in the main body of the paper we know that for arbitrary vectors x,y ∈ Rd:
exp( ∥Ai,jx∥2
2 )exp(
∥(Ai,j)−1y∥2
2 )SM(x,y) = E[exp(ω⊤(Ai,jx+ (Ai,j)−1y)] (89)
Therefore, the softmax kernel can be estimated as:
SM(x,y) ≈ ΨmAi,j (x) ⊤Ψm(Ai,j)−1(y) (90)
with ΨmAi,j (x) and Ψ m (Ai,j)−1(y) given by:
ΨmAi,j (x) def = 1√ m exp(−∥A i,jx∥2 2 )(exp(ω⊤1 A i,jx), ..., exp(ω⊤mA i,jx))⊤ (91)
Ψm(Ai,j)−1(y) def = 1√ m exp(−∥(A i,j)−1y∥2 2 )(exp(ω⊤1 (A i,j)−1y), ..., exp(ω⊤m(A i,j)−1y))⊤ (92)
for ω1, ..., ωm ∼ N (0, Id). Such a softmax kernel estimator has variance equal to zero if x = ri and y = rj if we are using the complex matrix and has small variance if we are using the real matrix. For other pairs of vectors (x,y) close to ri and rj , the variance is close to zero, or relatively small.
We can also take the Gaussian λ-coefficient estimator that gets it maximum value when x = ri,y = rj as λi,j(x,y). Such an estimator could be constructed with Ai,j in a similar way.
It can be easily linearized since:
λi,j(x,y) = exp(−∥A i,jx+ (Ai,j)−1y∥2
2τ2 ) =
exp(−∥A i,jx∥2 2τ2 )exp(−∥(A i,j)−1y∥2 2τ2 )SM( x τ , −y τ ),
(93)
where SM(x/τ,−y/τ) could be estimated with similar random feature maps as given by Eq. 90. Therefore, the Gaussian λ-coefficients can be estimated as:
λi,j(x,y) = exp(−∥A i,jx+ (Ai,j)−1y∥2
2τ2 ) ≈ Ψτ,mAi,j (x) ⊤Ψτ,m(Ai,j)−1(y) (94)
with Ψτ,mAi,j (x) and Ψ τ,m (Ai,j)−1(−y) given by:
Ψτ,mAi,j (x) = 1√ m exp(−∥A i,jx∥2 τ2 )(exp(ω⊤1 A i,jx/τ), ..., exp(ω⊤mA i,jx/τ))⊤ (95)
Ψτ,m(Ai,j)−1(−y) = 1√ m exp(−∥(A i,j)−1y∥2 τ2 )(exp(−ω⊤1 (Ai,j)−1y/τ), ..., exp(−ω⊤m(Ai,j)−1y/τ))⊤
(96)
for ω1, ..., ωm ∼ N (0, Id) 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 Gaussian-Mixtures estimators for clustered data). Assume that inputs x (queries) and y (keys) can be modeled by nq and nk clusters respectively with centers ri and rj ∈ Rd, (i = 1, ..., nk, j = 1, ..., nq). We denote Ai,j as the complex matrix satisfying Eq. (85) with center ri, rj . Furthermore, we denote E = (ŜM i,j (x,y)) nq,nk i=1,j=1 as a list of estimators of the softmax kernel SMi,j(x,y) and Λ = λ̂i,j(x,y))nq,nki=1,j=1 as a list of estimators of λ i,j(x,y) constructed independently from E . We also use one additional base estimator (e.g. ŜM trig (x,y) or ŜM ++ (x,y)) and denote the estimator of its softmax kernel as ŜM 0 (x,y) and λ-coefficient as
λ̂0(x,y). Then our hybrid estimator takes the following form:
ŜM E,Λ (x,y) = nq∑ i=1 nk∑ j=1 λ̂i,j(x,y)ŜM i,j (x,y) + λ̂0(x,y)ŜM 0 (x,y) (97)
with constraint:
nq∑ i=1 nk∑ j=1 λ̂i,j(x,y) + λ̂0(x,y) = 1 (98)
and where the estimators of λ-coefficients and base estimators are given as:
λ̂i,j(x,y) = Ψτ,mAi,j (x) ⊤Ψτ,m(Ai,j)−1(−y) (99)
ŜM i,j (x,y) = ΨmAi,j (x) ⊤Ψm(Ai,j)−1(y) (100)
for ΨmAi,j (x),Ψ m (Ai,j)−1(y),Ψ τ,m Ai,j (x),Ψ τ,m (Ai,j)−1(−y) given by Eq. (91, 92, 95, 96).
Definition H.2 (Hybrid Zero-One-Mixtures estimators for clustered data). Assume that inputs x (queries) and y (keys) can be modeled by nq and nk clusters respectively with centers ri and rj ∈ Rd, (i = 1, ..., nk, j = 1, ..., nq). We denote Ai,j as the complex matrix satisfying Eq. (85) with center ri, rj . Furthermore, we denote E = (ŜM i,j (x,y)) nq,nk i=1,j=1 as a list of estimators of the softmax kernel SMi,j(x,y) and Λ = λ̂i,j(x,y))nq,nki=1,j=1 as a list of estimators of λ i,j(x,y) constructed independently from E . Then our hybrid estimator takes the following form:
ŜM E,Λ (x,y) = nq∑ i=1 nk∑ j=1 λ̂i,j(x,y)ŜM i,j (x,y) (101)
with constraint:
nq∑ i=1 nk∑ j=1 λ̂i,j(x,y) = 1 (102)
and where the estimators of λ-coefficients and base estimators are given as:
λ̂i,j(x,y) = Ψi(x)Ψj(y) (103)
ŜM i,j (x,y) = ΨmAi,j (x) ⊤Ψm(Ai,j)−1(y) (104)
for ΨmAi,j (x),Ψ m (Ai,j)−1(y) given by Eq. (91, 92), with Ψi(x) being a scalar indicating whether x belongs to the i-th cluster of the nq clusters and similarly with Ψj(y) being a scalar indicating whether y belongs to the j-th of the nk clusters.
In our experiments we 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 two random vectors X0, Y0 in R50. Then we generate four random orthogonal matrices Oi,j ∈ R50×50 where i = 1, 2 and j = 1, 2 and use Xi = Oi,1X0 and Yi = Oi,2Y0 for i = 1, 2 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(xi,k) = sign(yj,k) for some k due to the usage of real diagonal matrices. To control the value of s, we manually adjust the sign of each element of Xi and Yi. The data for input x is the combination of two clusters, each consisting 1000 data points following Gaussian distribution with mean vectors X1, X2 and the common covariance matrix σ2I . 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 σ. After this step, we normalize the data points by controlling their ℓ2 norms to make the true softmax value be within a reasonable range. After the data is created, we first use K-means clustering algorithm with k = 2 to cluster x and y, and we obtain centers (xi,yj)i,j∈{1,2}. Then we use the set 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, we 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 be found in Fig. 6, we also record the maximal and minimal empirical MSE over 20 repetitions for all estimators in the Table 3, where s1 = [0.08, 0.05, 0.04, 0.03], s2 = [0.11, 0.09, 0.08, 0.07], s3 = [0.05, 0.03, 0.03, 0.02], s4 = [0.07, 0.06, 0.05, 0.04] representing the lists of s values for all pairs of clusters for different synthetic data sets.

heading: I LANGUAGE MODELING TRAINING DETAILS AND ADDITIONAL RESULTS
text: For the Language Modeling tasks, we trained a 2-layer LSTM of hidden sizes 200 and 650 on the PennTree Bank (Marcus et 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., 2017; 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 is shown in Table 4. 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.
Since the HRF mechanisms can be sensitive to 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 1 hurts the model performance and so we chose to use a temperature scaling as in (Rawat et al., 2019). Thus before passing the computed dot product to the softmax layer, we scaled the dot product by τ > 1. Our final loss function for the language modeling task was:
L(θ) = LCE + λ1E(||q||2 − 1)2 + λ2E(||k||2 − 1)2 (105)
where || · ||2 is the row-wise L2 norm of the appropriate matrices, and LCE is the cross-entropy loss. The distribution of the lengths of the keys and the queries are shown in Figure 9.
Finally we present some additional results on the WikiText2 dataset. For our experiments on the WikiText2 dataset, we chose λ1 = λ2 = 2 and τ = 6. Our model trained with these hyperparameters achieve a perplexity score of 105.35 on the test set after 50 epochs.
We compute the 1-dimensional Wasserstein distance and the Kolmogrov-Smirnov (KS) metric (Ramdas et al., 2017) between the approximated softmax distribution and the true softmax distribution. The mean and the standard deviation of the statistical metric results in Fig. 8 is shown in Table 5. 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.
The above tables (Table 4 and Table 5) show that the our hybrid variants consistently outperform Favor+ 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 on the PennTree Bank dataset, similar to the experiments carried out in (Rawat et al., 2019).
We trained a 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 to approximate the expected loss in each training epoch. As observed in Fig. 2, the relative error of all three estimators grow exponentially fast with embedding norm r. Therefore, if we keep un-normalized embeddings during sampling, then even though we could get an unbiased estimation of the loss, the variance could be high. Such bias-variance trade off is also mentioned in paper (Rawat et al., 2019). To solve this issue, we used normalized input and class embeddings during sampling to generate biased sampling distribution with lower variance, while keeping unnormalized embeddings to 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.
We compared our hybrid estimator with Trigonometric, FAVOR+, the uniform method which sampled 40 negative classes with equal probability, as well as the sampled softmax method which uses the unbiased true probability calculated from full softmax as the sampling probability. Trigonometric and FAVOR+ mechanisms used 1024 random features. To make comparison fair, hybrid variants were using (m,n)-configurations characterized by the similar number of FLOPS needed to construct their corresponding random features as regular RF-estimators. We reported our comparison
results averaged over 25 independent runs in Fig. 10. The perplexity score in Fig. 10 is defined as 2cross 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 is not very significant.

heading: J SPEECH EXPERIMENTS: ADDITIONAL DETAILS
text: Tested Conformer-Performer models consisted of l = 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 L ≈ 500. Padding mechanism was applied for all tested variants.

heading: K DOWNSTREAM ROBOTICS EXPERIMENTS
text: In both Robotics experiments we found query/key L2-normalization technique particularly effective (queries and keys of L2-norm equal to 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 actuated joints, 3 per leg. The task is to walk forward on a randomized uneven terrain that requires careful foot placement planning based on visual feedback. The ground is made of a series of stepstones 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 2 depth cameras attached to its body, one on the front and other on the belly facing downwards. We use Implicit Attention Policy (IAP) architecture (masking variant) described in Choromanski et al. (2021a) which uses Performer-based attention mechanism to process 32× 24 depth images from the 2 cameras.
We apply a hierarchical setup to 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 a position-based swing leg controller, and an model predictive control (MPC) based stance leg controller, to achieve the foot placement decided by high level. The policies are trained with evolutionary strategies (ES). In Fig. 4 (left) we compare FAVOR+ with angular hybrid approximation in IAP. HRF with 8 x 8 random projections (m = n = 8) performs as well as the softmax kernel with 256 random projections. A series of image frames along the episode of a learned locomotion HRF-based IAP policy is shown bottom right of Fig. 4. 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 on them.
The training curves in Fig. 4 (left) are obtained by averaging over 5 top runs. The shaded regions shows the standard deviation over 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 is 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 counter-examples), and the ground truth one-hot probability distribution for a given observation 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 a 500× 500× 3 RGB image. An example of the camera image is given in Fig. 11. The training curves Fig. 4 (right) are obtained by averaging over 3 learning rates: 10−3, 10−4, 10−5.

heading: L COMPUTATIONAL GAINS OF HYBRID ESTIMATORS
text: Consider the general hybrid estimator of SM(x,y) from Equation 8, using p + 1 base estimators. Assume that the kth base estimator for k = 1, ..., p + 1 applies m random features in each of the corresponding tk random maps and that n random features are used in each of the corresponding lλk random maps to approximate k
th λ-coefficient for k = 1, ..., p. From Equation 11, we conclude that time complexity of constructing Ψ(z) is:
T = O ( (t1 + ...+ tp+1)md+ (nd+md+mn)
p+1∑ k=1 p∑ r=1 tklr
) (106)
Thus the resulting Θ(mn)-dimensional random feature map Ψ(z) can be constructed in time O(nd+ md + mn) as opposed to O(mnd) as it is the case for the regular estimator, providing substantial computational gains. This is true if the (θ)(mn)-dimensional HRFs can provide similar quality as their mn-dimensional regular counterparts. It turns out that this is often the case, see: Section D.4.

Review summaries:
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.   Three 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. HRF has also a lower computational complexity in the total number of random features w.r.t. non compositional estimators.  Experiments 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.

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. 

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. ============================================================= This 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.