SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Valerii Likhosherstov\*†  
University of Cambridge  
v.lihosherstov@gmail.com
Krzysztof Choromanski*  
Google DeepMind & Columbia University  
kchoro@gmail.com
Avinava Dubey* Google Research
Frederick Liu Google Research
Tamas Sarlos  
Google Research
Adrian Weller  
University of Cambridge & The Alan Turing Institute
The problem of efficient approximation of a linear operator induced by the Gaussian or softmax kernel is often addressed using random features (RFs) which yield an unbiased approximation of the operator's result. Such operators emerge in important applications ranging from kernel methods to efficient Transformers. We propose parameterized, positive, non-trigonometric RFs which approximate Gaussian and softmax-kernels. In contrast to traditional RF approximations, parameters of these new methods can be optimized to reduce the variance of the approximation, and the optimum can be expressed in closed form. We show that our methods lead to variance reduction in practice ( $e^{10}$ -times smaller variance and beyond) and outperform previous methods in a kernel regression task. Using our proposed mechanism, we also present FAVOR#, a method for self-attention approximation in Transformers. We show that FAVOR# outperforms other random feature methods in speech modelling and natural language processing.
Random feature decomposition is an important technique for the linearization of nonlinear kernel functions with theoretical guarantees such as unbiasedness and concentration around the true kernel value. Linearization allows a significant reduction in computations from quadratic to linear complexity in the size of the operator induced by the kernel. The technique emerged under the name of random kitchen sinks (RKS) introduced in [39, 40, 41] and was used in many applications such as kernel SVM [43, 29, 38, 3], dimensionality reduction [18, 1], neural networks [12, 22, 49, 24, 14], function-to-function regression [34], kernel regression [28, 2], nonparametric adaptive control [6], differentially-private ML algorithms [11], operator-valued kernels [32, 8] and semigroup kernels [51]. An in-depth theoretical analysis of random features was performed by [30, 52, 44, 42].
An exciting recent application of random features is in the area of scalable Transformer networks [16, 13, 17, 26], where the self-attention matrix is approximated as a low-rank matrix when the sequence is long. However, the RKS family of methods relies on the Fourier transform, resulting in sin and cos types of random features, which were shown to be unsuitable for application in Transformers due to negative values in the low-rank matrix. [16] proposed a solution in the form of positive-valued random features relying on the exponential function (positive random features, PosRFs), yielding
* Equal contribution
† Work done while at University of Cambridge. The author is at Waymo now.
a method they called Fast Attention Via Orthogonal positive Random features (FAVOR+) for self-attention approximation. This solution was improved by [31] by means of a careful choice of the linear combination parameters under the exponent, and the so-called homogeneity heuristic, which allows a choice of one set of parameters for all approximated values. The resulting random features were called generalized exponential random features (GERFs), and the corresponding self-attention approximation method was termed FAVOR++.
Contributions: In this paper, we make a leap forward in the design of positive-valued random features by proposing dense exponential random features (DERFs) which contain both PosRFs and GERFs as special cases. Instead of scalar parameters as in GERFs, DERFs rely on matrix parameters and dense quadratic forms inside the exponent. We show how to select parameters of the new random features efficiently without harming the overall subquadratic complexity.
More technically, our contributions are as follows:
1. We show that the homogeneity heuristic of [31] may in fact be viewed not as a heuristic, but as a closed-form optimum of the shifted log-variance objective.
2. We introduce DERFs and three special instantiations: asymmetric DERFs (ADERFs), symmetric DERFs (SDERFs), and simplified ADERFs (SADERFs). All these instantiations contain GERFs as a special case (Figure 1, left). For each instantiation we prove that a closed-form optimum of the shifted log-variance objective can be found efficiently.
3. We show that our new variants result in lower variance than GERFs and other previous methods in practice (e.g. up to  $e^{10}$  times variance improvement as in Figure 1, right). Further, we show that DERFs outperform other random features in kernel regression and Transformer setups (speech modelling and natural language processing). We refer to the DERF-based self-attention approximation method as FAVOR#.
2 Prerequisites
2.1 Scaled softmax kernel and random features
By the scaled softmax kernel  $K^{(\alpha)}: \mathbb{R}^d \times \mathbb{R}^d \to (0, +\infty)$ , where  $\alpha \in \mathbb{R}$ , we denote a mapping defined as  $K^{(\alpha)}(\mathbf{x}, \mathbf{y}) = \exp(\alpha \| \mathbf{x} \|^2 + \mathbf{x}^\top \mathbf{y} + \alpha \| \mathbf{y} \|^2)$  for all  $\mathbf{x}, \mathbf{y} \in \mathbb{R}^d$  where  $\| \cdot \|$  is an  $L_2$ -norm. Two important special cases of the scaled softmax kernel are 1) the Gaussian kernel  $K^{(-1/2)}(\mathbf{x}, \mathbf{y}) = \exp(-\| \mathbf{x} - \mathbf{y} \|^2/2)$  and 2) the softmax kernel  $K^{(0)}(\mathbf{x}, \mathbf{y}) = \exp(\mathbf{x}^\top \mathbf{y})$ . For two sets of vectors  $\mathcal{X} = \{\mathbf{x}^{(i)} \in \mathbb{R}^d\}_{i=1}^L$  and  $\mathcal{Y} = \{\mathbf{y}^{(j)} \in \mathbb{R}^d\}_{j=1}^L$ , by  $\mathcal{K}(\mathcal{X}, \mathcal{Y}) \in \mathbb{R}^{L \times L}$  we denote a matrix where  $\mathcal{K}^{(\alpha)}(\mathcal{X}, \mathcal{Y})_{i,j} = K^{(\alpha)}(\mathbf{x}^{(i)}, \mathbf{y}^{(j)})$  for all  $1 \leq i, j \leq L$ .
In this paper, we will be interested in the problem of computing  $\mathcal{K}^{(\alpha)}(\mathcal{X},\mathcal{Y})\mathbf{C}$  where  $\mathcal{X},\mathcal{Y}$  and a matrix  $\mathbf{C}\in \mathbb{R}^{L\times n}$  are provided as an input. A naive solution requires  $O(L^2 (d + n))$  computations for constructing  $\mathcal{K}^{(\alpha)}(\mathcal{X},\mathcal{Y})$ $(O(L^{2}d))$  and computing the matrix multiplication  $\mathcal{K}^{(\alpha)}(\mathcal{X},\mathcal{Y})\times \mathbf{C}$ $(O(L^{2}n))$ . Instead, we will use an efficient Monte-Carlo approximation of  $\mathcal{K}^{(\alpha)}(\mathcal{X},\mathcal{Y})\times \mathbf{C}$  using the following notion of random features (RFs) for the scaled softmax kernel:
Figure 1: (left) Venn diagram of the new types of random features (green) we propose. (right) Logarithm of the relative variance of different random feature maps on pairs of vectors sampled from CIFAR10 and MNIST/CIFAR10. A new random feature map SDERF results in a consistent variance reduction of the previous best method GERF, up to  $\approx e^{10}$  and  $\approx e^5$  times. Figure 2 extends this plot.
Definition 2.1. By random features for the scaled softmax kernel  $K^{(\alpha)}$ ,  $\alpha \in \mathbb{R}$ , we denote a triple  $\mathcal{T} = \langle \nu, f^{(1)}, f^{(2)} \rangle$  where  $\nu$  is a probability distribution over random objects  $\omega \in \Omega$  and  $f^{(i)}: \Omega \times \mathbb{R}^d \to \mathbb{R}$ ,  $i \in \{1, 2\}$ , are such mappings that, for all  $\mathbf{x}, \mathbf{y} \in \mathbb{R}^d$ ,
$$
K ^ {(\alpha)} (\mathbf {x}, \mathbf {y}) = \mathbb {E} _ {\nu} \left[ f ^ {(1)} (\boldsymbol {\omega}, \mathbf {x}) f ^ {(2)} (\boldsymbol {\omega}, \mathbf {y}) \right]. \tag {1}
$$
The decomposition of type (1) can be used for an efficient unbiased approximation of  $\mathcal{K}^{(\alpha)}(\mathcal{X},\mathcal{Y})\mathbf{C}$ . Let  $\omega^{(1)},\ldots ,\omega^{(M)}\in \Omega$  be i.i.d. samples from  $\nu$ . Define matrices  $\mathbf{P},\mathbf{S}\in \mathbb{R}^{L\times M}$  where for all  $1\leq i,j\leq L$
$$
\mathbf {P} _ {i,:} = M ^ {- 1 / 2} \left(f ^ {(1)} \left(\boldsymbol {\omega} ^ {(m)}, \mathbf {x} ^ {(i)}\right)\right) _ {m = 1} ^ {M}, \quad \mathbf {S} _ {j,:} = M ^ {- 1 / 2} \left(f ^ {(2)} \left(\boldsymbol {\omega} ^ {(m)}, \mathbf {y} ^ {(j)}\right)\right) _ {m = 1} ^ {M}, \tag {2}
$$
where  $\mathbf{P}_{i,:},\mathbf{S}_{j,:}\in \mathbb{R}^d$  are column vectors corresponding to the rows of  $\mathbf{P},\mathbf{S}$ . Then according to (1),  $\widehat{\mathcal{K}} = \mathbf{PS}^\top$  is an unbiased Monte Carlo (MC) approximation of  $\kappa^{(\alpha)}(\mathcal{X},\mathcal{Y})$  on  $M$  samples. The variance  $\mathrm{Var}\widehat{\mathcal{K}}_{i,j} = M^{-1}\mathrm{Var}_{\nu}f^{(1)}(\pmb {\omega},\mathbf{x}^{(i)})f^{(2)}(\pmb {\omega},\mathbf{y}^{(j)})$  of this approximation is inversely proportional to  $M$ , hence  $M$  is a tradeoff parameter between computations and precision. Now,  $\widehat{\kappa}\mathbf{C}$  is an unbiased approximation of  $\kappa^{(\alpha)}(\mathcal{X},\mathcal{Y})\mathbf{C}$  but  $\widehat{\kappa} = \mathbf{PS}^\top$  is a rank-  $M$  matrix, hence computing  $\widehat{\kappa}\mathbf{C}$  has  $O(LMn)$  complexity. Assuming that sampling each  $\pmb{\omega}^{(m)}$  and computing  $f^{(\cdot)}(\cdot ,\cdot)$  are  $O(d)$  operations, which is usually the case, precomputing  $\mathbf{P}$  and  $\mathbf{S}$  takes  $O(LMd)$  computations, resulting in a total  $O(LM(d + n))$  computational complexity. By choosing  $M\ll L$ , we obtain a significant reduction in computations compared to the exact variant:  $O(LM(d + n))\ll O(L^{2}(d + n))$ .
Operations of type  $\mathcal{K}^{(\alpha)}(\mathcal{X},\mathcal{Y})\mathbf{C}$ , especially for the Gaussian kernel  $\alpha = -1 / 2$ , emerge in kernel SVM [40], kernel regression [33, 47] and in physics in the form of the Gauss transform [50]. Another important application is in the area of efficient Transformers and is discussed in the next section [16].
2.2 Random features for efficient Transformers
RFs found a prominent application in the area of efficient long-sequence Transformers [16]. Transformers rely on a self-attention block for propagating information between elements of the sequence. If the sequence length is  $L$  and input matrices are denoted as  $\mathbf{Q}$ ,  $\mathbf{K}$ ,  $\mathbf{V} \in \mathbb{R}^{L \times d}$  (queries, keys and values), then self-attention outputs the following matrix:
$$
\mathbf {Y} = \operatorname {d i a g} \left(\mathcal {K} ^ {(0)} (\mathcal {X}, \mathcal {Y}) \mathbf {1} _ {L}\right) ^ {- 1} \mathcal {K} ^ {(0)} (\mathcal {X}, \mathcal {Y}) \mathbf {V} \in \mathbb {R} ^ {L \times d}, \tag {3}
$$
where  $\mathbf{1}_L\in \mathbb{R}^L$  is a vector of all ones,  $\mathrm{diag}(\cdot)$  returns a diagonal  $(L\times L)$ -sized matrix with the argument on the diagonal,  $\mathcal{X} = \{\mathbf{x}^{(i)} = d^{-1 / 4}\mathbf{Q}_{i,:}\in \mathbb{R}^d\}$ , and  $\mathcal{Y} = \{\mathbf{y}^{(j)} = d^{-1 / 4}\mathbf{K}_{j,:}\in \mathbb{R}^d\}$ . Hence, substitution of  $\widehat{\kappa}$  instead of  $\kappa^{(0)}(\mathcal{X},\mathcal{Y})$  in (3) reduces computational complexity from  $O(L^{2}d)$  to  $O(LMd)$  ( $n = d + 1$ ).  $\mathrm{diag}(\mathcal{K}^{(0)}(\mathcal{X},\mathcal{Y})\mathbf{1}_L)^{-1}\mathcal{K}^{(0)}(\mathcal{X},\mathcal{Y})$  is the result of a softmax operation performed on rows of  $d^{-1 / 2}\mathbf{Q}\mathbf{K}^\top$ .
2.3 Existing random features for the softmax kernel
Representation (1) is not unique and different RFs can be proposed for a single  $K^{(\alpha)}$ . Note that if  $\langle \nu, f^{(1)}, f^{(2)} \rangle$  are RFs for  $K^{(0)}$ , then  $\langle \nu, \widehat{f}^{(1)}, \widehat{f}^{(2)} \rangle$  are RFs for  $K^{(\alpha)}$  for  $\alpha \in \mathbb{R}$  where  $\widehat{f}^{(k)}(\omega, \mathbf{x}) = \exp(\alpha \| \mathbf{x} \|^2) f^{(k)}(\omega, \mathbf{x})$ . Hereafter we focus on the softmax kernel  $K^{(0)}$  without loss of generality.
[15] proposed to use trigonometric random features (TrigRFs) from [40] for efficient Transformers:
$$
\begin{array}{l} \Omega_ {\mathrm {t r i g}} = \mathbb {R} ^ {d + 1}, \quad \nu_ {\mathrm {t r i g}} = \operatorname {U n i f} ([ 0, 2 \pi ]) \times \mathcal {N} (\mathbf {0} _ {d}, \mathbf {I} _ {d}) ^ {d}, \quad f _ {\mathrm {t r i g}} ^ {(1)} ((\theta , \widetilde {\boldsymbol {\omega}}), \mathbf {x}) = (4) \\ = \sqrt {2} \exp \left(\| \mathbf {x} \| ^ {2} / 2\right) \cos \left(\widetilde {\boldsymbol {\omega}} ^ {\top} \mathbf {x} + \theta\right), \quad f _ {\mathrm {t r i g}} ^ {(2)} \left(\left(\theta , \widetilde {\boldsymbol {\omega}}\right), \mathbf {y}\right) = \sqrt {2} \exp \left(\| \mathbf {y} \| ^ {2} / 2\right) \cos \left(- \widetilde {\boldsymbol {\omega}} ^ {\top} \mathbf {y} + \theta\right), (5) \\ \end{array}
$$
where  $\omega = (\theta, \widetilde{\omega})$ ,  $\mathrm{Unif}(\cdot)$  denotes a uniform distribution on the argument set,  $\mathcal{N}(\mathbf{0}_d, \mathbf{I}_d)$  is a multivariate Gaussian distribution with mean  $\mathbf{0}_d$  (vector of  $d$  zeros) and covariance matrix  $\mathbf{I}_d$  (identity matrix of size  $d \times d$ ).
The next iteration of efficient attention approximators [16] observed a problem with TrigRFs (4-5). The attention matrix  $\mathrm{diag}(\mathcal{K}^{(0)}(\mathcal{X},\mathcal{Y})\mathbf{1}_L)^{-1}\mathcal{K}^{(0)}(\mathcal{X},\mathcal{Y})$  from (3) is right stochastic meaning that its
entries are nonnegative and each row sums to 1 due to the normalizing term  $\mathrm{diag}(\mathcal{K}^{(0)}(\mathcal{X},\mathcal{Y})\mathbf{1}_L)^{-1}$ . However, since  $f_{\mathrm{trig}}^{(\cdot)}$  can be arbitrary real numbers,  $\mathbf{P},\mathbf{S}$  (2) and, therefore,  $\widehat{\kappa}$  can take negative values. Hence,  $\mathrm{diag}(\widehat{\kappa}\mathbf{1}_L)^{-1}\widehat{\kappa}$  is not right stochastic in general and entries of  $\widehat{\kappa}\mathbf{1}_L$  can take very small and/or negative values resulting in unstable behaviour when inverting  $\mathrm{diag}(\widehat{\kappa}\mathbf{1}_L)^{-1}$ . [16] therefore proposed a new type of positive random features (PosRFs) which have the form:
$$
\Omega_ {\mathrm {p o s}} = \mathbb {R} ^ {d}, \quad \nu_ {\mathrm {p o s}} = \mathcal {N} (0, 1) ^ {d}, \quad f _ {\mathrm {p o s}} ^ {(1)} (\boldsymbol {\omega}, \mathbf {x}) = f _ {\mathrm {p o s}} ^ {(2)} (\boldsymbol {\omega}, \mathbf {x}) = \exp (\boldsymbol {\omega} ^ {\top} \mathbf {x} - \| \mathbf {x} \| ^ {2} / 2).
$$
It is clear that such  $f_{\mathrm{pos}}^{(\cdot)}$  only take strictly positive values resulting in the right stochastic  $\mathrm{diag}(\widehat{\mathcal{K}}\mathbf{1}_L)^{-1}\widehat{\mathcal{K}}$  and a stable Transformer training procedure.
[31] extended PosRFs, proposing generalized exponential random features (GERFs) for  $K^{(0)}$ :
$$
\Omega_ {\mathrm {G E}} = \mathbb {R} ^ {d}, \nu_ {\mathrm {G E}} = \mathcal {N} (0, 1) ^ {d}, f _ {\mathrm {G E}} ^ {(1)} (\boldsymbol {\omega}, \mathbf {x}) = f _ {\mathrm {G E}} ^ {(2)} (\boldsymbol {\omega}, \mathbf {x}) = D \exp \left(A \| \boldsymbol {\omega} \| ^ {2} + B \boldsymbol {\omega} ^ {\top} \mathbf {x} + C \| \mathbf {x} \| ^ {2} / 2\right) \tag {6}
$$
where  $A, B, C, D$  are real numbers satisfying:  $1 - 8A > 0$ ,  $B = \sqrt{1 - 4A}$ ,  $C = -0.5$ ,  $D = (1 - 4A)^{d/4}$ . [31] express  $B, C, D$  through  $A$  and find a closed-form equation for the variance of (1):
$$
\operatorname {V a r} _ {\nu_ {\mathrm {G E}}} f _ {\mathrm {G E}} ^ {(1)} (\boldsymbol {\omega}, \mathbf {x}) f _ {\mathrm {G E}} ^ {(2)} (\boldsymbol {\omega}, \mathbf {y}) = e ^ {\mathcal {L} _ {\mathrm {G E}} (A, \mathbf {x}, \mathbf {y})} - K ^ {(0)} (\mathbf {x}, \mathbf {y}) ^ {2},
$$
$$
\mathcal {L} _ {\mathrm {G E}} (A, \mathbf {x}, \mathbf {y}) = d \log \left(\frac {1 - 4 A}{\sqrt {1 - 8 A}}\right) + \frac {2 (1 - 4 A)}{1 - 8 A} \| \mathbf {x} + \mathbf {y} \| ^ {2} - \| \mathbf {x} \| ^ {2} - \| \mathbf {y} \| ^ {2}. \tag {7}
$$
The minimum variance corresponds to the minimum  $\mathcal{L}(A, \mathbf{x}, \mathbf{y})$  since  $K^{(0)}(\mathbf{x}, \mathbf{y})^2$  does not depend on  $A$ . Since  $\mathcal{L}(A, \mathbf{x}, \mathbf{y})$  is defined for a single pair of  $\mathbf{x}$ ,  $\mathbf{y}$  and not for sets  $\mathcal{X}, \mathcal{Y}$ , [31] propose a homogeneity heuristic when they replace  $\| \mathbf{x} + \mathbf{y} \|^2$ ,  $\| \mathbf{x} \|^2$ ,  $\| \mathbf{y} \|^2$  in (7) with averages over  $\mathcal{X}, \mathcal{Y}$ :  $L^{-2} \sum_{i,j} \| \mathbf{x}^{(i)} + \mathbf{y}^{(j)} \|^2$ ,  $L^{-1} \sum_{i} \| \mathbf{x}^{(i)} \|^2$  and  $L^{-1} \sum_{j} \| \mathbf{y}^{(j)} \|^2$  respectively. This heuristic is based on the assumption that  $\{\mathbf{x}^{(i)}\}$  and  $\{\mathbf{y}^{(j)}\}$  are homogeneous and their statistics are tightly concentrated around the mean. After this, the minimum of (7) with respect to  $A$  can be found in closed form.
3 Dense-exponential random features (DERFs)
We prove that the homogeneity heuristic (Section 2.3) corresponds to a certain minimization problem. Then, we present DERFs which generalize GERFs and provide a tighter solution of that problem.
3.1 The objective minimized by GERFs
Our first contribution is showing that the homogeneity heuristic adopted in GERFs is actually an analytic solution of a certain optimization problem. Define
$$
\bar {\mathcal {L}} (\boldsymbol {\theta}; \mathcal {X}, \mathcal {Y}, \mathcal {T}) = L ^ {- 2} \sum_ {1 \leq i, j \leq L} \log \left(\operatorname {V a r} _ {\nu} \left[ f ^ {(1)} \left(\boldsymbol {\omega}, \mathbf {x} ^ {(i)}\right) f ^ {(2)} \left(\boldsymbol {\omega}, \mathbf {y} ^ {(j)}\right) \right] + K ^ {(0)} \left(\mathbf {x} ^ {(i)}, \mathbf {y} ^ {(j)}\right) ^ {2}\right), \tag {8}
$$
where  $\mathcal{T} = \langle \nu, f^{(1)}, f^{(2)} \rangle$  are RFs for the kernel  $K^{(0)}$  and  $\pmb{\theta}$  are their parameters appearing implicitly in  $\nu, f^{(1)}, f^{(2)}$ . (8) is a mean log-variance shifted by  $K^{(0)}(\mathbf{x}^{(i)}, \mathbf{y}^{(j)})^2$ . The best possible value of (8) is  $\log K^{(0)}(\mathbf{x}^{(i)}, \mathbf{y}^{(j)})$  which corresponds to all variances  $\mathrm{Var}f^{(1)}(\omega, \mathbf{x}^{(i)})f^{(2)}(\omega, \mathbf{y}^{(j)})$  being zero, meaning that RFs provide exact kernel estimation. Hence, minimization of (8) leads to more precise estimators on  $\mathcal{X}, \mathcal{Y}$ . We call the loss function  $\overline{\mathcal{L}}(\pmb{\theta}; \mathcal{X}, \mathcal{Y}, \mathcal{T})$  the shifted log-variance objective. Since  $\log$  is concave loss  $\overline{\mathcal{L}}$  is conceptually similar to relative standard deviation: the smaller  $K^{(0)}(\mathbf{x}^{(i)}, \mathbf{y}^{(j)})$ , the higher  $\overline{\mathcal{L}}$  is for the same amount of variance  $\mathrm{Var}f^{(1)}(\omega, \mathbf{x}^{(i)})f^{(2)}(\omega, \mathbf{y}^{(j)})$ .
If  $\mathcal{T}_{\mathrm{GE}} = \langle \nu_{\mathrm{GE}},f_{\mathrm{GE}}^{(1)},f_{\mathrm{GE}}^{(2)}\rangle$  are taken in (8), then  $\pmb{\theta}_{\mathrm{GE}} = \{A,B,C,D\}$  and  $\overline{\mathcal{L}} (\pmb{\theta}_{\mathrm{GE}};\mathcal{X},\mathcal{Y},\mathcal{T}_{\mathrm{GE}}) = L^{-2}\sum_{i,j}\mathcal{L}_{\mathrm{GE}}(A;\mathbf{x}^{(i)},\mathbf{y}^{(j)})$ . Using (7), we get:  $\overline{\mathcal{L}} (\pmb{\theta}_{\mathrm{GE}};\mathcal{X},\mathcal{Y},\mathcal{T}_{\mathrm{GE}}) =$
$$
d \log \left(\frac {1 - 4 A}{\sqrt {1 - 8 A}}\right) + \frac {2 - 8 A}{1 - 8 A} \frac {1}{L ^ {2}} \sum_ {i, j} \| \mathbf {x} ^ {(i)} + \mathbf {y} ^ {(j)} \| ^ {2} - \frac {1}{L} \sum_ {i} \| \mathbf {x} ^ {(i)} \| ^ {2} - \frac {1}{L} \sum_ {j} \| \mathbf {y} ^ {(j)} \| ^ {2}.
$$
[31] define these RFs for  $K^{(-1/2)}$  but we adapt them for  $K^{(0)}$  using the trick mentioned above.
[231] consider a more generalized form when  $A, B, C, D$  are complex with an additional parameter  $s = \pm 1$ , however only the subfamily (6) with  $s = 1$  is proposed for use in the Transformer application.
That is,  $\mathcal{L}(A; \mathcal{X}, \mathcal{Y})$  coincides with (7) when  $\| \mathbf{x} + \mathbf{y} \|^2$ ,  $\| \mathbf{x} \|^2$ ,  $\| \mathbf{y} \|^2$  are replaced by their average statistics computed on  $\mathcal{X}$ ,  $\mathcal{Y}$ . Hence, the homogeneity heuristic (Section 2.3) is nothing but minimization of (8). While in general it's unclear how to find a closed-form optimum of  $\operatorname{Var}(\widehat{\mathcal{K}})$  or  $\operatorname{Var}(\widehat{\mathcal{K}}\mathbf{C})$ , the global minimum of (8) is feasible and can be computed in  $O(1)$  time. Further, in the case of GERF, the values minimizing (8) lead to very good results in large-scale applications of efficient Transformers as shown by [31] without knowing about (8) objective. In the next section we present extensions of GERFs which aim to minimize (8) in closed form.
4 Towards DERFs
Dense-exponential random features (DERFs) are an extension of GERFs where scalars  $A, B, C$  are replaced with dense matrices. DERFs may be viewed as a generalization that contain the previously introduced classes as special cases. We define DERFs as follows:  $\Omega_{\mathrm{DE}} = \mathbb{R}^d$ ,  $\nu_{\mathrm{DE}} = \mathcal{N}(0,1)^d$  and for  $k \in \{1,2\}$ :
$$
f _ {\mathrm {D E}} ^ {(k)} (\boldsymbol {\omega}, \mathbf {x}) = D \exp (\boldsymbol {\omega} ^ {\top} \mathbf {A} \boldsymbol {\omega} + \boldsymbol {\omega} ^ {\top} \mathbf {B} ^ {(k)} \mathbf {x} + \mathbf {x} ^ {\top} \mathbf {C} ^ {(k)} \mathbf {x}),
$$
where  $\mathbf{B}^{(k)},\mathbf{C}^{(k)}\in \mathbb{R}^{d\times d}$ $D\in \mathbb{R}$ $\mathbf{A}\in \mathbb{S}_d$  (a set of  $d\times d$  real symmetric matrices). Clearly, GERFs with parameters  $A,B,C,D$  can be expressed via DERFs with parameters  $\mathbf{A} = A\mathbf{I}_d$ $\mathbf{B}^{(1)} = \mathbf{B}^{(2)} = B\mathbf{I}_d$ $\mathbf{C}^{(1)} = \mathbf{C}^{(2)} = C\mathbf{I}_d$ $D$  is unchanged. Our first theoretical result is giving the conditions when  $\mathcal{T}_{\mathrm{DE}} = \langle \nu_{\mathrm{DE}},f_{\mathrm{DE}}^{(1)},f_{\mathrm{DE}}^{(2)}\rangle$  are valid RFs:
Theorem 4.1. Let the following conditions hold:  $8\mathbf{A} \prec \mathbf{I}_d$ ,  $(\mathbf{B}^{(1)})^\top (\mathbf{I}_d - 4\mathbf{A})^{-1}\mathbf{B}^{(2)} = \mathbf{I}_d$ ,  $\mathbf{C}^{(k)} = -\frac{1}{2} (\mathbf{B}^{(k)})^\top (\mathbf{I}_d - 4\mathbf{A})^{-1}\mathbf{B}^{(k)}$ ,  $D = \operatorname*{det}(\mathbf{I}_d - 4\mathbf{A})^{1/4}$  where  $k \in \{1,2\}$ . Then  $\mathcal{T}_{\mathrm{DE}}$  are RFs for  $K^{(0)}$  and for all  $\mathbf{x}, \mathbf{y} \in \mathbb{R}^d$ :  $\operatorname{Var}_{\nu_{\mathrm{DE}}} f_{\mathrm{DE}}^{(1)}(\boldsymbol{\omega}, \mathbf{x}) f_{\mathrm{DE}}^{(2)}(\boldsymbol{\omega}, \mathbf{y}) =$
$$
\begin{array}{l} D ^ {4} \det (\mathbf {I} _ {d} - 8 \mathbf {A}) ^ {- 1 / 2} \exp \left(2 \mathbf {x} ^ {\top} \left(\mathbf {C} ^ {(1)} + (\mathbf {B} ^ {(1)}) ^ {\top} (\mathbf {I} _ {d} - 8 \mathbf {A}) ^ {- 1} \mathbf {B} ^ {(1)}\right) \mathbf {x} \right. \\ + 2 \mathbf {y} ^ {\top} \left(\mathbf {C} ^ {(2)} + \left(\mathbf {B} ^ {(2)}\right) ^ {\top} \left(\mathbf {I} _ {d} - 8 \mathbf {A}\right) ^ {- 1} \mathbf {B} ^ {(2)}\right) \mathbf {y} + 4 \mathbf {x} ^ {\top} \left(\mathbf {B} ^ {(1)}\right) ^ {\top} \left(\mathbf {I} _ {d} - 8 \mathbf {A}\right) ^ {- 1} \mathbf {B} ^ {(2)} \mathbf {y}\left. \right) - K ^ {(0)} (\mathbf {x}, \mathbf {y}) ^ {2}. \tag {9} \\ \end{array}
$$
Our ultimate goal is to find optimal parameters  $\mathbf{A}$ ,  $\mathbf{B}^{(k)}$ ,  $\mathbf{C}^{(k)}$  and  $D$  minimizing the variance of the low-rank approximation of  $\mathcal{K}^{(0)}(\mathcal{X},\mathcal{Y})$  where sets  $\mathcal{X},\mathcal{Y}$  are provided. Our first observation is that we can assume that  $\mathbf{A} \in \mathbb{D}_d$  (a set of  $d \times d$  real diagonal matrices). Indeed, any symmetric  $\mathbf{A}$  can be expressed as  $\mathbf{Q}\widetilde{\mathbf{A}}\mathbf{Q}^{\top}$  where  $\mathbf{Q} \in \mathbb{O}_d$  (a set of orthogonal matrices  $\{\mathbf{Z} \in \mathbb{R}^{d \times d} \mid \mathbf{Z}^{\top}\mathbf{Z} = \mathbf{I}_d\}$ ) and  $\widetilde{\mathbf{A}} \in \mathbb{D}_d$ . Let  $\omega \sim \mathcal{N}(\mathbf{0}_d,\mathbf{I}_d)$ . Then, for any  $\mathbf{x} \in \mathbb{R}^d$ ,  $k \in \{1,2\}$ ,