Hybrid Random Features

1 INTRODUCTION & RELATED WORK . 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 : 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 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 . 2 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 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 ) 2.2 THE ALGORITHM . 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 ) .

Summary: 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.