All Random Features Representations are Equivalent

Kernel methods have a long and illustrious history in machine learning [1]. Nowadays one of their most widespread applications is the softmax kernel inside the attention layers of a transformer. Unfortunately, kernel methods scale linearly in the size of the dataset, which is impractical for very large datasets.

In their seminal paper, Rahimi and Recht [2] show that this linear cost can be avoided by replacing the kernel with a randomized approximation that is correct in expectation. They did this by means of a random feature representation.

To see why this is useful, consider a simple kernel estimator over some dataset $\{(x_{1},y_{1}),(x_{2},y_{2}),...,(x_{n},y_{n})\}$.

As written, this kernel estimator requires linear time to evaluate, since for every new $x$ it’s necessary to compute a dot product with every element in the dataset. But if the kernel admits a random features representation, we may write

This expectation can be approximated via sampling,

where $\omega_{1},\omega_{2},\dots,\omega_{k}$ are samples drawn from $\Omega$. As long as $k$ is less than $n$, we can save time by precomputing the sum on the right. The principle challenge is finding an estimator that

Choromanski et al. [3] were the first to apply this idea to attention. They did so by finding a choice of $\phi$ that was numerically stable and had low sample variance. This kicked off a flurry of research ([4] [5] [6]) into better choices of $\phi$ with ever lower sample variance.

We provide a potential conclusion to this competition. Rather than optimizing $\phi$, we simply chose our sampling strategy carefully. By approximating the expectation using an optimal importance sampling strategy, we are able to achieve a sample error that is as low as possible.

Remarkably, this sample variance does not depend on $\phi$, meaning that when sampled optimally, all random features schemes have the same variance.

Ours is not the first scheme to consider importance sampling, but in contrast to others that sample based on kernel polarization heuristics [7] or quadrature rules [8], we directly solve the variance minimization problem. To our knowledge we are the first to obtain global optima that are independent of the choice of $\phi$.

Importance sampling [9] is a well known technique for reducing variance when estimating an expectation via sampling. It’s based on the observation that some samples have a larger effect on the expectation than others. Instead of naively sampling from the given distribution, we sample from some new distribution $\Psi$, rescaling the samples such that the correct final estimate is preserved.

In our case, we wish to estimate (1) via sampling. Let $\Psi$ be some new distribution with the same support as $\Omega$, and suppose $\Psi$ and $\Omega$ have densities $p_{\Psi}$ and $p_{\Omega}$, respectively. Observe that

In other words, we are free to sample from $\Psi$ without changing the value of the expectation. On the other hand, the sample variance, given by

does depend on $\Psi$. Notice that the mean squared error of our sampling procedure is precisely this variance. All we need to do to minimize the per-sample error is find a choice of $\Psi$ to make this quantity small.

Choosing an optimal $\Psi$ is an intrinsically data-driven question, as the variance depends on the input. We’d like to minimize the expected sample variance over all inputs, as captured in the following definition.

The $\Psi$ that minimizes $\mathcal{V}_{\Psi}$ can be found analytically.

This result tells us how to minimize the variance for a particular choice of $\phi$. What is less clear in this form is that it also tells us something about optimizing over every $\phi$. We will explore this in the next section.

As written, Theorem 1 appears to be a statement about how to get the most out of any particular choice of $\phi$. While this is useful, the real insight is what it can tell us about every possible $\phi$. We are ready now to begin exploring our main point: the optimal sample variance is often entirely independent of $\phi$.

Our first step in this direction is a tight upper bound on the sampling variance, and therefore the error, that can be stated without reference to $\phi$.

This follows from Theorem 1 by applying the Cauchy-Schwarz inequality to the first term in $\mathcal{\hat{V}}$. Specifically,

Written this way, our bound is itself suggestively similar to the Cauchy-Schwarz inequality. Indeed, since positive-definite kernels obey the Cauchy-Schwarz inequality, we know that this expression is positive as expected. More interestingly, since this expression does not depend on $\phi$, this bound is the same regardless of which representation we’re using.

In the (relatively common) case where $\mathcal{X}_{1}=\mathcal{X}_{2}$, we can go even further.

This can be seen from Theorem 1, since in this case $q_{\phi}\left(\omega\right)=\mathop{\mathbb{E}}_{x_{1}\sim\mathcal{X}_{1}}\left[\phi\left(x_{1},\omega\right)^{2}\right]$. By exchanging the expectations it follows that $\mathop{\mathbb{E}}_{\omega\sim\Omega}q_{\phi}\left(\omega\right)=\mathop{\mathbb{E}}_{x_{1}\sim\mathcal{X}_{1}}\left[K\left(x_{1},x_{1}\right)\right]$.

The significance of this result is that we are able to write down the exact sample variance without any reference to $\phi$. This is not just a statement about the error of a particular representation, it is a statement of equality that holds across all representations. It does not matter which $\phi$ you choose. Every random features representation has the same sample error as every other.