Variance-Reducing Couplings for Random Features

## Abstract

Abstract Random features (RFs) are a popular technique to scale up kernel methods in machine learning, replacing exact kernel evaluations with stochastic Monte Carlo estimates.
They underpin models as diverse as efficient transformers (by approximating attention) to sparse spectrum Gaussian processes (by approximating the covariance function).
Efficiency can be further improved by speeding up the convergence of these estimates: a variance reduction problem.
We tackle this through the unifying lens of optimal transport , finding couplings to improve RFs defined on both Euclidean and discrete input spaces.
They enjoy theoretical guarantees and sometimes provide strong downstream gains, including for scalable approximate inference on graphs.
We reach surprising conclusions about the benefits and limitations of variance reduction as a paradigm, showing that other properties of the coupling should be optimised for attention estimation in efficient transformers.

## 1 Introduction

Kernel methods are ubiquitous in machine learning (Canu and Smola, ; Smola and Schölkopf, ; Kontorovich et al., ; Campbell, ).
Through the kernel trick, they provide a mathematically principled and elegant way to perform nonlinear inference using linear learning algorithms.
The eponymous positive definite *kernel function* $k:\mathcal{X}\times\mathcal{X}\to\mathbb{R}$ measures the ‘similarity’ between two datapoints.
The input domain $\mathcal{X}$ may be continuous, e.g. the set of vectors in $\mathbb{R}^{d},$ or discrete, e.g. the set of graph nodes or entire graphs.

Random features for kernel approximation. Though very effective on small datasets, kernel methods suffer from poor scalability.
The need to materialise and invert the *Gram matrix* $\mathbf{K}\coloneqq[k(\bm{x}_{i},\bm{x}_{j})]_{i=1}^{N}$ leads to a time complexity cubic in the size of the dataset $N$.
Substantial research has been dedicated to improving scalability by approximating this matrix, a prominent example being *random features* (RFs) (Rahimi and Recht, ; ; Avron et al., ; Liu et al., ).
These randomised mappings $\phi:\mathbb{R}^{d}\to\mathbb{R}^{s}$ construct low-dimensional or sparse feature vectors that satisfy

$$ $k(\bm{x},\bm{y})=\mathbb{E}\left(\phi(\bm{x})^{\top}\phi(\bm{y})\right).$ (1) $$

The expectation $\mathbb{E}$ is taken over an ensemble of *random frequencies* $\{\bm{\omega}\}_{i=1}^{m}$ drawn from a distribution $\mathcal{\eta}$.
The space in which $\{\bm{\omega}\}_{i=1}^{m}$ live and manner in which they are combined to construct $\phi(\bm{x})$ depends on the particular input space $\mathcal{X}$ and kernel function $k$ being approximated.
This paper will consider several examples.
The set of RFs $\{\phi(\bm{x}_{i})\}_{i=1}^{N}$ can be used to construct a low-rank or sparse approximation of the Gram matrix, providing substantial space and time complexity savings. RFs exist for a variety of kernels, including for continuous and discrete input spaces (Dasgupta et al., ; Johnson, ; Choromanski et al., ; Goemans and Williamson, ; Rahimi and Recht, ; Choromanski, ; Tripp et al., ).

Variance reduction for RFs.
Eq. 1 can be understood as a *Monte Carlo* (MC) estimate of $k$.
In applications, it is often found that this estimate converges slowly.
This can be addressed by taking many samples $m$, but this undermines the efficiency gains of RFs.
Therefore, substantial effort has been dedicated to *reducing the variance* of the kernel estimates.
Variance reduction methods include quasi-Monte Carlo (QMC; Dick et al., ; Yang et al., ), common random numbers (CRNs; Glasserman and Yao, ), antithetic variates (Hammersley and Morton, ) and structured Monte Carlo (SMC; Yu et al., ). These techniques work by replacing i.i.d. frequencies $\{\bm{\omega}_{i}\}_{i=1}^{m}$ by a *dependent ensemble*, with the sample dependencies designed to improve RF convergence.

Limitations of previous techniques. The best choice of dependencies between $\{\bm{\omega}_{i}\}_{i=1}^{m}$ is an active research area.
Though straightforward to apply, standard QMC techniques are suboptimal.
They are based on hard-coded ‘low-discrepancy sequences’ so typically do not incorporate information about the particular kernel $k$ being approximated.
Empirical performance may be poor and theoretical guarantees lacking in the low-sample, high-dimensional regime (Rowland et al., ; Morokoff and Caflisch, ), which is precisely where RFs are most important.
On the other hand, hand-crafted SMC dependencies, which impose strict geometrical conditions like orthogonality between frequencies, tend to fare better (Yu et al., ).
But they are difficult to design, theoretical guarantees are hard-won and optimality is not guaranteed.
RFs for estimating kernels defined on discrete spaces like the nodes of a graph have only recently been developed (Choromanski, ; Tripp et al., ), so here very few effective variance reduction techniques have even been proposed.
This paper asks: *can we devise a principled framework for coupling RFs, providing variance reduction across basis functions and input domains, including with very few samples?*

Optimal transport.
To answer this, we propose to frame variance reduction as *optimal transport* (OT): an active research area of applied mathematics that studies how to move (probability) mass between distributions as efficiently as possible (Villani et al., ).
This novel perspective equips us with proof techniques and numerical tools to identify the *best possible dependencies* between samples, giving lower kernel estimator variance compared to previous approaches.
OT allows us to improve couplings for RFs in both Euclidean and discrete spaces, including with different basis functions.
To our knowledge, this has never before been achieved in the same paper.

Our contributions.
This work presents unifying strategies to reduce the variance of random features.

- 1.
We frame the problem of variance reduction of RFs as *optimal transport* (OT) (Sec. 2), and use this perspective to improve the convergence of *three* popular classes of RFs: random Fourier features, random Laplace features and graph random features.
- 2.
For random Fourier features (RFFs) and random Laplace features (RLFs), we exactly solve the OT problem for the *norms* of $m=2$ orthogonal frequencies (Sec. 3).
We introduce *pairwise norm-coupling*, which guarantees lower variance for arbitrary $m$.
- 3.
For graph random features (GRFs), we couple the *lengths* of random walks by finding a bipartite matching between the quantiles of the marginal distributions (Sec. 4).
This is the first time a coupling between random walks has been optimised using data, beating hard-coded algorithms.
- 4.
We test our algorithms on UCI datasets and real-world graphs, verifying that OT couplings substantially reduce kernel estimator variance (Secs 3 and 4).
We show that this sometimes translates to much better performance in downstream tasks, including for approximate inference with scalable graph-based Gaussian processes. However, we also reach surprising conclusions about the limitations of variance reduction for RFs, including for efficient transformers.

All proofs are saved for the Appendices, but are also sketched in the main body where space allows.

## 2 Preliminaries

From kernel estimation to optimal transport (OT).
Define the *kernel estimator* $\widehat{k}(\bm{x}_{i},\bm{x}_{j})\coloneqq\phi(\bm{x}_{i})^{\top}\phi(\bm{x}_
{j})$.
Recall that $\phi(\cdot)$ is computed using random frequencies $\{\bm{\omega}_{i}\}_{i=1}^{m}$, with the space in which they live, distribution from which they are drawn, and manner in which they are combined dependent on the particular kernel being approximated.
The estimator is unbiased provided each frequency $\bm{\omega}_{i}$ obeys some marginal distribution $\eta$.
Importantly, independence of $\{\bm{\omega}_{i}\}_{i=1}^{m}$ is *not* required: any joint distribution with marginals $\eta$ gives an unbiased estimator. We refer to the set of such joint distributions as *couplings*.

The coupling between the frequencies determines the estimator variance.
We want to solve:

$$ $\textrm{minimise }\mathcal{I}(\mu)=\mathbb{E}_{\bm{\omega}_{1:m}\sim\mu}c(\bm{ \omega}_{1:m})\quad\textrm{ for }\quad\mu\in\Lambda_{m}(\eta),$ (2) $$

where we defined the *cost function* $c(\bm{\omega}_{1:m})\coloneqq\left(\phi(\bm{x})^{\top}\phi(\bm{y})\right)^{2}$ and $\Lambda_{m}(\eta)$ denotes the set of couplings of $m$ random variables with marginal measures $\eta$.
This is precisely the *Kantorovich formulation* of a multi-marginal OT problem (see Eq. 4 of the seminal OT text of Villani ()).
We will generally consider cost functions where the minimiser exists and we want to find efficient new MC couplings, so the task is to find the *optimal coupling* $\mu^{*}=\arg\min_{\mu\in\Lambda_{m}(\chi_{d})}\left[\mathbb{E}_{\omega_{1:m}
\sim\mu}c(\omega_{1:m})\right]$ with the smallest estimator variance.
The relationship between variance reduction and OT was also noted by Rowland et al. () in a different context.

(Approximately) solving the OT problem.
The formulation of Eq. 2 depends on the particular RF mechanism and kernel being approximated.
We will show that one can solve it exactly for RFFs and RLFs (Sec. 3) and approximately for GRFs (Sec. 4), which have input domains $\mathcal{X}=\mathbb{R}^{d}$ ($d$-dimensional Euclidean space) and $\mathcal{X}=\mathcal{N}$ (the set of graph nodes) respectively.
This gives new couplings with lower RF variance than previous algorithms.

## 3 Random Fourier features and random Laplace features

RFFs and RLFs. To begin, we consider the task of approximating the popular *Gaussian kernel* $k(\bm{x}_{i},\bm{x}_{j})\coloneqq\exp(-\|\bm{x}_{i}-\bm{x}_{j}\|^{2}/2)$ with data $\{\bm{x}_{i}\}_{i=1}^{N}\subset\mathbb{R}^{d}$.
This can be achieved using Rahimi and Recht’s celebrated *random Fourier features* (RFFs) (Rahimi and Recht, ),

$$ $\phi_{\textrm{RFF}}(\bm{x})=\sqrt{\frac{1}{m}}\left(\odot_{i=1}^{m}\left[\sin( \bm{\omega}_{i}^{\top}\bm{x}),\cos(\bm{\omega}_{i}^{\top}\bm{x})\right]\right),$ (3) $$

where $\odot$ denotes concatenation.
These provide an unbiased estimate if the frequencies $\{\bm{\omega}_{i}\}_{i=1}^{m}$ are marginally Gaussian, $\bm{\omega}_{i}\sim\mathcal{N}(0,\mathbf{I}_{d})$.
RFFs are widely used for scaling kernel methods such as Gaussian processes (GPs; Williams and Rasmussen, ) and support vector machines (SVMs; Scholkopf and Smola, ).
The time complexity of computing the exact posterior of a GP is $\mathcal{O}(N^{3}),$ where $N$ is the number of datapoints.
Using RFFs, one can approximate the posterior with $m\ll N$ features, reducing this cost to $\mathcal{O}(Nm^{2})$.
Changing basis functions, $k(\bm{x}_{i},\bm{x}_{j})$ can also be approximated using *random Laplace features* (RLFs) (Yang et al., ),

$$ $\phi_{\textrm{RLF}}(\bm{x})=\sqrt{\frac{1}{m}}\exp(-\|\bm{x}\|^{2})(\odot_{i=1 }^{m}\exp(\bm{\omega}_{i}^{\top}\bm{x})),$ (4) $$

where again $\bm{\omega}_{i}\sim\mathcal{N}(0,\mathbf{I}_{d}).$
Unlike RFFs, RLFs guarantee positive kernel estimates.
This makes them better suited to approximating attention in efficient transformers (Choromanski et al., ), where negative estimates cause training instabilities.
Using $m$ RLFs to get a low-rank decomposition of attention with $N$ $d$-dimensional tokens, one can reduce the time complexity of transformers from $\mathcal{O}(N^{2}+Nd)$ to $\mathcal{O}(Nmd)$ with low performance loss.

Orthogonal random features.
A common variance reduction technique for both RFFs and RLFs is the *orthogonality trick* (Yu et al., ; Rowland et al., ; Reid et al., ; Choromanski et al., ).
Exploiting the isotropy of $\mathcal{N}(0,\mathbf{I}_{d})$, one can constrain the frequency vectors $\{\bm{\omega}_{i}\}_{i=1}^{m}$ to be exactly orthogonal whilst preserving their marginal distributions.
This is found to reduce the kernel estimator variance and improve performance in downstream tasks.
Whilst this technique couples the *directions* of the random frequencies $\{\widehat{\bm{\omega}}_{i}\}_{i=1}^{m}$, their *norms* $\{\omega_{i}\}_{i=1}^{m}$ (with $\omega_{i}\coloneqq\|\bm{\omega}_{i}\|_{2}$) are left independent so the coupling is suboptimal.
By solving an OT problem, we will show how coupling the norms can further reduce estimator variance.

### 3.1 Solving the OT problem for maximal variance reduction

Consider an ensemble of $m$ orthogonal random frequency directions $\{\widehat{\bm{\omega}}_{i}\}_{i=1}^{m}$, jointly randomly rotated so they are marginally isotropic.
Our task is to couple their norms $\{\omega_{i}\}_{i=1}^{m}$ to suppress the RFF and RLF kernel estimator variance.
The marginal distribution of each $\omega_{i}$ must be $\chi_{d}$ (a Chi distribution with $d$ degrees of freedom) to ensure that each $\bm{\omega}_{i}$ is marginally Gaussian.
We can extend recent results by Reid et al. () to compute the OT cost functions.

###### Lemma 3.1 (OT formulation for RFFs and RLFs) .

When estimating $k(\bm{x},\bm{y})$ with $m$ orthogonal RFFs and RLFs, the OT formulation of the variance reduction problem is:

$$ $\mu^{*}=\arg\min_{\mu\in\Lambda_{m}(\chi_{d})}\left[\mathbb{E}_{\omega_{1:m} \sim\mu}c(\omega_{1:m})\right],\quad\textrm{where}$ (5) $$

$$ $c_{\textrm{RFF}}(\omega_{1:m})=\sum_{i,j\neq i}^{m}\sum_{k=0}^{\infty}\frac{(- 1)^{k}z^{2k}\left(\omega_{i}^{2}+\omega_{j}^{2}\right)^{k}}{2^{2k}k!\Gamma(k+ \frac{d}{2})},\quad c_{\textrm{RLF}}(\omega_{1:m})=\sum_{i,j\neq i}^{m}\sum_{k =0}^{\infty}\frac{v^{2k}(\omega_{i}^{2}+\omega_{j}^{2})^{k}}{2^{2k}k!\Gamma(k+ \frac{d}{2})},$ (6) $$

with $z\coloneqq\|\bm{x}-\bm{y}\|_{2}$ and $v\coloneqq\|\bm{x}+\bm{y}\|_{2}$.
$\Gamma$ is the gamma function.

This is a tough multi-marginal OT problem.
However, remarkably, we can solve it *exactly*, under mild asymptotic assumptions for RFFs, when $m=2$.
The following result is novel.

###### Theorem 3.2 (Solution to OT problem when $m=2$ ) .

Denote by $F_{\chi_{d}}(\cdot)$ the cumulative distribution function (CDF) of $\chi_{d}$.
Consider $m=2$ orthogonal frequencies with norms $(\omega_{1},\omega_{2})$.
For RLFs, the OT problem in Eq. 5 is solved by the *negative monotone* coupling

$$ $F_{\chi_{d}}(\omega_{1})+F_{\chi_{d}}(\omega_{2})=1.$ (7) $$

For RFFs, Eq. 7 ensures lower cost than any other coupling, provided $z$ is sufficiently small.

*Proof sketch.* We defer a full proof of this important result to App. A.2; here is a brief sketch. OT plans satisfy a property called ‘$c$-monotonicity’, which specifies how the support of the optimal coupling depends on the cost function.
For RLFs, $c_{\textrm{RLF}}$ immediately implies negative monotonicity (Eq. 7).
For RFFs, this is only true for the first nontrivial term in $z$.
By bounding the contribution from the remaining terms, one can show that Eq. 7 still guarantees lower variance than any other coupling if $z$ is small enough.
Specifically, letting $\mu_{\textrm{NM}}$ denote the negative monotone coupling, for any other coupling $\mu^{\prime}\in\Lambda_{2}(\eta)\backslash\left\{\mu_{\textrm{NM}}\right\}$ there exists some constant $\delta(\mu^{\prime})>0$ such that $\mathcal{I}(\mu_{\textrm{NM}})<\mathcal{I}(\mu^{\prime})$ for all $z<\delta$ (Lemma A.6). ∎

Given $m=d$ orthogonal frequencies, one can partition the ensemble into $\lfloor\frac{d}{2}\rfloor$ orthogonal pairs, with one remaining frequency if $d$ is odd.
For every pair, one can impose negative monotone coupling (Eq. 7).
We refer to such ensembles as *pairwise norm-coupled* (PNC).

###### Definition 3.3 (Pairwise norm-coupled RFs) .

PNC is no more expensive than i.i.d. norms.
To reduce the variance further, one can take multiple independent PNC ensembles.
An important corollary of Thm. 3.2 is as follows.

###### Corollary 3.4 (Superiority of pairwise norm-coupled RFs) .

For any $m$, the variance of pairwise norm-coupled RFs is guaranteed to be lower than orthogonal RFs with independent norms, in full generality for RLFs and provided $z$ is small enough for RFFs.

Negative monotone coupling differs from OT plans usually seen in machine learning; it is a *space-filling* coupling that seeks long transport plans that give diverse samples.
However, it is a popular heuristic technique for variance reduction via common random numbers (CRNs) in computational statistics (Glasserman and Yao, ).
To our knowledge, this is the first result applying it to improving the convergence of orthogonal RFs, and the first corresponding guarantees for variance reduction.
We make one further theoretical contribution for RLFs.

###### Theorem 3.5 (Recovering antithetic sampling with RLFs) .

For RLFs with $m=2$ frequencies whose respective orientations $(\widehat{\bm{\omega}}_{1},\widehat{\bm{\omega}}_{2})$ are unconstrained, variance is minimised by conditioning that $\bm{\omega}_{1}=-\bm{\omega}_{2}$ almost surely (that is, opposite directions and equal norms).

This coupling is known as *antithetic sampling* (Hammersley and Morton, ).
Thm. 3.5 shows that, given a PNC ensemble $\{\bm{\omega}_{i}\}_{i=1}^{d}$, we can obtain further variance reduction by augmenting it to $\{\pm\bm{\omega}_{i}\}_{i=1}^{d}$.
Antithetic sampling is also a common (though often heuristically motivated) variance reduction strategy used e.g. when estimating attention in Performers (Choromanski et al., ).
We can reinterpret its effectiveness as an OT coupling.

###### Lemma 3.1 (OT formulation for RFFs and RLFs) .

###### Theorem 3.2 (Solution to OT problem when m = 2 𝑚 2 m=2 italic_m = 2 ) .

###### Definition 3.3 (Pairwise norm-coupled RFs) .

###### Corollary 3.4 (Superiority of pairwise norm-coupled RFs) .

###### Theorem 3.5 (Recovering antithetic sampling with RLFs) .

### 3.2 Pushing further with numerical OT solvers

Multi-marginal OT.
In Sec. 3.1 we proposed PNC RFs: a computationally efficient coupling that is guaranteed to reduce variance for any $m$.
We obtained it by solving the variance reduction OT problem exactly in $m=2$, then combining $\lfloor\frac{d}{2}\rfloor$ independent copies to get the ensemble.
Can we do better by inducing dependencies between the *all* the $m$ frequencies’ norms?
Solving this multi-marginal OT problem analytically is a tough open problem.

Copulas as numerical OT solvers.
Whilst an analytic solution to the multi-marginal OT variance reduction problem is (for now) out of reach, we can make progress using a numerical OT solver.
Our strategy is to restrict $\Lambda_{m}(\chi_{d})$, the full set of joint distributions over $m$ random variables with $\chi_{d}$ marginals, to a tractable subset amongst which we can efficiently optimise and sample.
One such subset is provided by *Gaussian copulas* (Nelsen, ; Haugh, ): joint distributions obtained by taking a multivariate Gaussian and pushing each of its coordinates forward first with the Gaussian CDF $F_{\mathcal{N}}$, and then the $\chi_{d}$ inverse CDF $F_{\chi_{d}}^{-1}$.
Fig. 1 gives a visual overview for $d=2$.
If the diagonal terms of the underlying Gaussian covariance matrix $\mathbf{\Sigma}\in\mathbb{R}^{m\times m}$ are equal to $1$ (i.e. it is a correlation matrix), this has the prescribed marginals so unbiasedness is baked in.
Meanwhile, correlations *between* the random variables are controlled by the off-diagonal entries of $\mathbf{\Sigma}$.
This parameterises a broad set of couplings, including PNC (Def. 3.3).
In App. A.5 we demonstrate that it is possible to use gradient descent with the reparameterisation trick to *learn* the optimal copula covariance matrix $\mathbf{\Sigma}$, approximately solving the multi-marginal OT problem.
We do this by minimising the kernel approximation error on training data, exploiting the fact that all operations to construct the features are differentiable.
In doing so, we optimise the RF coupling.
Remarkably, this data-dependent optimisation does *not* to find couplings much better than PNC: see the training curves in Fig. 5.
This suggests that our scheme may already be close optimal for $m\neq 2$.
Intuitively, one cannot simultaneously anticorrelate too many random variables, so strong pairwise couplings already perform very well.
Whilst copulas have previously been used as numerical OT solvers (Chi et al., ), this is (to our knowledge) their first application to learning a Monte Carlo coupling.

Figure: Figure 1: Copula schematic for $d=2$. Random variables are drawn from a Gaussian distribution with correlation matrix $\mathbf{\Sigma}$. They are pushed forward using $F_{\mathcal{N}}$ then $F_{\chi_{d}}^{-1}$ to obtain coupled variables with marginal $\chi_{d}$ distributions. $\mathbf{\Sigma}$ is learned using gradient-based optimisation, approximately solving the multi-marginal OT problem in Eq. 5.

### 3.3 Experiments for norm-coupled RFs

To test PNC RFs (Def. 3.3), we now compute kernel estimates with RFFs and RLFs for UCI datasets.
We choose the kernel lengthscale parameters based on a training set, by training a GP (RFFs) or selecting reasonable values for Performers (RLFs) (Choromanski et al., ).
We then compute the kernel approximation RMSE on a test set.
Full details are in App. B.1.

**Table 1: Performance of RFFs and RLFs on kernel estimation with UCI datasets with different coupling schemes, taking $d$ random frequencies for RFFs and $2d$ for RLFs (see main text for details). We show RMSEs to the ground truth kernel values, normalised such that the RMSE of the i.i.d. estimator is equal to one. Lower is better. Error bars are standard errors on the reported RMSEs. Our couplings consistently give the lowest variance.**
| Fourier Features | Concrete | Abalone | CPU | Power | Airfoil | Boston |
| --- | --- | --- | --- | --- | --- | --- |
| i.i.d. | $1.000$ $\pm 0.028$ | $1.000$ $\pm 0.042$ | $1.000$ $\pm 0.082$ | $1.000$ $\pm 0.037$ | $1.000$ $\pm 0.023$ | $1.000$ $\pm 0.018$ |
| Halton | $1.028$ $\pm 0.029$ | $0.991$ $\pm 0.042$ | $0.995$ $\pm 0.082$ | $0.913$ $\pm 0.033$ | $0.927$ $\pm 0.021$ | $1.176$ $\pm 0.022$ |
| Orthogonal | $0.627$ $\pm 0.019$ | $0.535$ $\pm 0.023$ | $0.617$ $\pm 0.070$ | $0.669$ $\pm 0.024$ | $0.586$ $\pm 0.013$ | $0.639$ $\pm 0.016$ |
| + PNC | $\mathbf{0.563}$ $\pm 0.019$ | $\mathbf{0.433}$ $\pm 0.019$ | $\mathbf{0.544}$ $\pm 0.071$ | $\mathbf{0.547}$ $\pm 0.020$ | $\mathbf{0.481}$ $\pm 0.011$ | $\mathbf{0.606}$ $\pm 0.018$ |
| Laplace Features | Concrete | Abalone | CPU | Power | Airfoil | Boston |
| i.i.d. | $1.000$ $\pm 0.092$ | $1.000$ $\pm 0.036$ | $1.000$ $\pm 0.086$ | $1.000$ $\pm 0.018$ | $1.000$ $\pm 0.026$ | $1.000$ $\pm 0.029$ |
| Halton | $0.721$ $\pm 0.067$ | $0.777$ $\pm 0.031$ | $0.779$ $\pm 0.084$ | $0.728$ $\pm 0.015$ | $0.721$ $\pm 0.021$ | $0.893$ $\pm 0.028$ |
| Orthogonal | $0.418$ $\pm 0.041$ | $0.546$ $\pm 0.026$ | $\mathbf{0.614}$ $\pm 0.098$ | $0.527$ $\pm 0.013$ | $0.489$ $\pm 0.016$ | $0.360$ $\pm 0.019$ |
| + PNC + antithetic | $\mathbf{0.367}$ $\pm 0.043$ | $\mathbf{0.486}$ $\pm 0.027$ | $\mathbf{0.618}$ $\pm 0.119$ | $\mathbf{0.438}$ $\pm 0.013$ | $\mathbf{0.418}$ $\pm 0.016$ | $\mathbf{0.324}$ $\pm 0.019$ |

Results for variance reduction.
Table 1 shows the results.
For RFFs, we take $m=d$ orthogonal frequencies.
For RLFs, we also include their antiparallel directions, giving $m=2d$ frequencies.
For each dataset, the RMSEs are normalised by the result with i.i.d. features.
As a further baseline, we include RFs constructed using *Halton sequences* (Dick et al., ; Yang et al., ), a fixed, off-the-shelf QMC scheme that can provide small gains but is clearly suboptimal. The third row shows orthogonal frequencies with independent norms (Yu et al., ). When we *also* couple the frequencies’ norms using our PNC scheme (plus antithetic sampling for RLFs, due to Thm 3.5), we access even lower estimator variance at no extra computational cost.
Note that the small $z$ condition for RFFs is found to be nonrestrictive in practice.

Downstream tasks.
We have achieved our objective of variance reduction, a popular and intensely studied goal in the literature (Yu et al., ; Rowland et al., ; Reid et al., ; Likhosherstov et al., ; Yang et al., ; Le et al., ; Bojarski et al., ; Choromanski et al., ; Lyu, ; Shen et al., ; Dao et al., ; Munkhoeva et al., ).
It is conventionally understood that PNC RFs should therefore improve downstream performance in applications.
Surprisingly, when we run exhaustive Gaussian process experiments in App. B.2, we do *not* observe such a gain.

The reason for this counterintuitive behaviour is as follows.
When optimising a coupling, we minimise the variance of *pointwise* kernel estimates $\{k(\bm{x}_{i},\bm{x}_{j})\}_{i,j=1}^{N}$.
However, functions like the predictive mean and KL divergence are highly nonlinear in these estimates.
For example, they may involve inverting a Gram matrix.
Downstream quantities therefore depend on the *joint* distribution of the kernel estimates, which are modified nontrivially by the coupling.
Variance reduction alone cannot guarantee an improvement.

Performers. As a concrete example, consider estimating *attention*, $a_{ij}\coloneqq k(\bm{x}_{i},\bm{x}_{j})/\sum_{l=1}^{N}k(\bm{x}_{i},\bm{x}_{l})$ using random Laplace features (Choromanski et al., ).
This normalises the kernel evaluation between the $i$ and $j$ tokens by the sum with all the other tokens.
Taylor expanding, if the kernel estimators have equal means $\mu$, the average mean square error ${\textrm{MSE}(\widehat{a}_{i})}\coloneqq\frac{1}{N}\sum_{j=1}^{N}\textrm{MSE}(
\widehat{a}_{ij})$ obeys

$$ ${\textrm{MSE}(\widehat{a}_{i})}=\frac{1}{N^{2}\mu^{2}}\left(\frac{1}{N}\sum_{j =1}^{N}\textrm{Var}(\widehat{k}(\bm{x}_{i},\bm{x}_{j}))-\frac{1}{N^{2}}\sum_{j _{1},j_{2}=1}^{N}\textrm{Cov}(\widehat{k}(\bm{x}_{i},\bm{x}_{j_{1}}),\widehat{ k}(\bm{x}_{i},\bm{x}_{j_{2}}))\right)+\mathcal{O}(\frac{1}{N^{3}}).$ (8) $$

**Table 2: Performer test accuracies on ImageNet with different coupling schemes. Counterintuitively, *maximising* the pointwise kernel estimator variance by positively correlating feature norms boosts performance – a different OT problem to the naive, obvious choice.**
| Laplace Features | Test acc. |
| --- | --- |
| Orthogonal | $0.625$ $\pm 0.003$ |
| Orthogonal + PNC | $0.620$ $\pm 0.003$ |
| Orthogonal + PM | $\mathbf{0.633}$ $\pm 0.003$ |

By coupling the frequency norms, PNC reduces $\textrm{Var}(\widehat{k}(\bm{x}_{i},\bm{x}_{j}))$ as intended.
However, it also reduces the covariance $\textrm{Cov}(\widehat{k}(\bm{x}_{i},\bm{x}_{j_{1}}),\widehat{k}(\bm{x}_{i},\bm
{x}_{j_{2}}))$, so ${\textrm{MSE}(\widehat{a}_{i})}$ does not actually substantially improve overall.
In stark contrast, if we instead take the *positive* monotone (PM) coupling where $\{\omega_{i}\}_{i=1}^{m}$ are all equal almost surely, then $\textrm{Var}(\widehat{k}(\bm{x}_{i},\bm{x}_{j}))$ is *maximised* (see App. B.3).
But these strong, positive correlations also increase $\textrm{Cov}(\widehat{k}(\bm{x}_{i},\bm{x}_{j_{1}}),\widehat{k}(\bm{x}_{i},\bm
{x}_{j_{2}}))$ by an even greater amount.
Hence, we find that ${\textrm{MSE}(\widehat{a}_{i})}$ *falls* (see Fig. 7).
This is surprising: maximising the pointwise kernel estimator variance by solving the OT problem with the ‘wrong’ sign on the cost function reduces the MSE of the attention scores after normalisation. In fact, the improvement is so big that it increases the average test accuracy of Performers trained on ImageNet (Deng et al., ) by +0.8%, whereas PNC makes no statistically significant difference.
See Table 2.
This demonstrates the limitations of simple variance reduction and invites a more careful treatment, considering the downstream quantities of interest.
App. B.3 gives further discussion and transformer training details.

## 4 Graph random features

We now shift our attention from $\mathbb{R}^{d}$ to the discrete domain.
Consider an undirected graph $\mathcal{G}(\mathcal{N},\mathcal{E})$ where $\mathcal{N}\coloneqq\{v_{1},...,v_{N}\}$ is the set of nodes and $\mathcal{E}$ is the set of edges, with $(v_{i},v_{j})\in\mathcal{E}$ if and only if there exists an edge between $v_{i}$ and $v_{j}$ in $\mathcal{G}$.
*Graph node kernels* $k:\mathcal{N}\times\mathcal{N}\to\mathbb{R}$ are positive definite, symmetric functions defined on pairs of nodes of $\mathcal{G}$, reflecting some notion of their ‘closeness’ via the graph edges and weights.
$k$ captures the structure of $\mathcal{G}$, letting practitioners repurpose popular kernelised learning algorithms to the discrete domain (Smola and Kondor, ; ).
Examples include the diffusion, regularised Laplacian, cosine and random walk kernels, all of which are typically considered as functions of the graph Laplacian matrix (Kondor and Lafferty, ).
We give a short introduction in App. C.1.

Graph random features.
As in the Euclidean setting, graph kernel methods scale poorly due to the $\mathcal{O}(N^{3})$ time complexity of inverting the Gram matrix $\mathbf{K}\coloneqq[k(v_{i},v_{j})]_{i,j=1}^{N}$.
In fact, even *computing* $\mathbf{K}$ often incurs a cubic cost since it involves multiplying large adjacency matrices.
Research has been dedicated to improving efficiency by approximating $\mathbf{K}$, including *graph random features* (GRFs) (Choromanski, ; Reid et al., ).
These are sparse vectors $\{\phi_{\textrm{GRF}}(v_{i})\}_{i=1}^{N}\subset\mathbb{R}^{N}$ that satisfy $\mathbf{K}_{ij}=\mathbb{E}\left[\phi_{\textrm{GRF}}(v_{i})^{\top}\phi_{\textrm
{GRF}}(v_{j})\right]$, so their dot product is equal to the true graph kernel in expectation.
GRFs are constructed in subquadratic time.

Coupled random walks.
For GRFs, the ‘frequencies’ $\{\bm{\omega}_{i}\}_{i=1}^{m}$ are *simple random walks*: sequences of graph nodes $(v_{i})_{i=1}^{l}$ with $(v_{i},v_{i+1})\in\mathcal{E}$.
At every timestep, the walker chooses one of its neighbours uniformly and at random.
The length of the walk $l$ is also a random variable, drawn from a geometric distribution $l\sim\textrm{G}(p),p\in(0,1)$. In other words, the walker terminates with probability $p$ at every timestep.
Random walks are usually taken to be independent, but this can lead to slow mixing times, poor efficiency and high kernel estimator variance (Alon et al., ; Zhou et al., ).
A natural question is: *can we couple graph random walks to improve the convergence of GRFs?*

Sec. 3 couples Gaussian vector norms.
A simple, analogous approach to couple random walks is via their *lengths*.
For GRFs, the constraint for unbiasedness is that the marginal distribution of each variable $l$ must remain geometric.
Reid et al. () proposed a simple, ad-hoc algorithm to achieve this called *antithetic termination*, which directly anticorrelates the walkers’ termination events at every timestep (see App. C.2).
They provide asymptotic theoretical guarantees, but only for particular choices of graph kernel.
In the next section, we will see that our novel method of *optimising* a coupling with data performs much better.

Figure: Figure 2: Schematic overview of $\sigma$-coupled GRFs with $\sigma=24351$. *Left*: permutation density with uniform marginals. *Centre*: bipartite matching between quantiles of geometric distributions over walk lengths. *Right*: ensemble of graph random walks with lengths to be coupled.

### 4.1 Approximating and solving the OT problem for GRFs

Here, we present our novel approach for formulating and approximately solving the variance reduction OT problem with GRFs.
It works by mapping to a corresponding *bipartite matching problem* which we can solve efficiently with linear programming techniques.
Fig. 2 gives a visual overview.

Constructing GRFs from random walks. To obtain $\phi_{\textrm{GRF}}(v_{i})\in\mathbb{R}^{N}$, one samples $m$ random walks $\{\bm{\omega}_{k}^{(i)}\}_{k=1}^{m}$ out of node $v_{i}\in\mathcal{N}$ and averages their ‘projections’,

$$ $\phi_{\textrm{GRF}}(v_{i})=\frac{1}{m}\sum_{k=1}^{m}\psi(\bm{\omega}_{k}^{(i)}).$ (9) $$

The projection function $\psi(\cdot):\Omega\to\mathbb{R}^{N}$ maps from the set of graph random walks $\Omega\coloneqq\left\{(v_{i})_{i=1}^{l}\,|\,v_{i}\in\mathcal{N},(v_{i},v_{i+1}
)\in\mathcal{E},l\in\mathbb{N}\right\}$ to a sparse $N$-dimensional feature vector satisfying $k(v_{i},v_{j})=\mathbb{E}_{\bm{\omega}^{(i)},\bm{\omega}^{(j)}}[\psi(\bm{
\omega}^{(i)})^{\top}\psi(\bm{\omega}^{(j)})]$.
$\psi(\cdot)$ depends on the particular kernel being approximated.
We direct the reader to the work of Reid et al. () for an introduction to GRFs, and also provide background in App. C.1.
$\psi(\cdot)$ is a complicated function; it is difficult to reason about analytically but straightforward to compute for a particular walk.
Moreover, its input walks are discrete random variables so $\psi(\cdot)$ is *not* differentiable with respect to the lengths $l_{k}^{(i)}$ (where $k=1,...,m$ and $v_{i}\in\mathcal{N}$ is the walker’s start node).
This precludes straightforward gradient-based optimisation (Sec. 3.2).

A pair of walkers. Initially consider just $m=2$ walkers.
The kernel estimator $\widehat{k}(v_{i},v_{j})$ is

$$ $\widehat{k}(v_{i},v_{j})=\phi_{\textrm{GRF}}(v_{i})^{\top}\phi_{\textrm{GRF}}( v_{j})=\frac{1}{4}\left(\psi(\bm{\omega}_{1}^{(i)})+\psi(\bm{\omega}_{2}^{(i)} )\right)^{\top}\left(\psi(\bm{\omega}_{1}^{(j)})+\psi(\bm{\omega}_{2}^{(j)}) \right),$ (10) $$

which is unbiased provided (i) the marginal distribution of each $\bm{\omega}_{1,2}^{(i,j)}$ is a simple random walk with geometrically distributed length (hereafter denoted $\eta_{\textrm{G}}$) and (ii) walks from node $v_{i}$ are independent from walks from node $v_{j}$.
The variance of the estimator depends on

$$ $\mathbb{E}\left(\widehat{k}(v_{i},v_{j})^{2}\right)=\frac{1}{16}\mathbb{E}_{ \bm{\omega}_{1,2}^{(i,j)}}\left(\left[\left(\psi(\bm{\omega}_{1}^{(i)})+\psi( \bm{\omega}_{2}^{(i)})\right)^{\top}\left(\psi(\bm{\omega}_{1}^{(j)})+\psi(\bm {\omega}_{2}^{(j)})\right)\right]^{2}\right),$ (11) $$

where the expectation is taken over both the *directions* and *lengths* of the random walks.
Suppose the directions remain independent, but the lengths $l_{1}^{(i)}$ and $l_{2}^{(i)}$ (and likewise $l_{1}^{(j)}$ and $l_{2}^{(j)}$) are to be coupled for variance reduction, analogously to the vector norms in Sec. 3.
Let $\mathbb{E}_{\textrm{dirs}}$ denote the expectation over the walkers’ directions.
We want to minimise:

$$ $\mathbb{E}_{(l_{1}^{(i)},l_{2}^{(i)})\sim\mu,(l_{1}^{(j)},l_{2}^{(j)})\sim\mu} \mathbb{E}_{\textrm{dirs}}\left(\left[\left(\psi(\bm{\omega}_{1}^{(i)})+\psi( \bm{\omega}_{2}^{(i)})\right)^{\top}\left(\psi(\bm{\omega}_{1}^{(j)})+\psi(\bm {\omega}_{2}^{(j)})\right)\right]^{2}\right)\textrm{ for }{\mu\in\Lambda_{2}( \eta_{\textrm{G}})}.$ (12) $$

OT, permutation densities and bipartite matchings.
The OT problem in Eq. 12 is analytically intractable.
To make progress, we must make approximations.
We report full details in App. C.3, limiting the main text to a high-level discussion of the main points.

First, to make the objective amenable to Monte Carlo approximation, we move $\mathbb{E}_{\textrm{dirs}}$ inside the square.
This is because, unlike the expression in Eq. 12, $\mathbb{E}_{\textrm{dirs}}(\psi(\bm{\omega}_{1,2}^{(i,j)}))$ can be efficiently estimated by simulating random walks.
Second, we must optimise amongst the class of couplings
$\Lambda_{2}(\eta_{\textrm{G}})$, joint distributions of two discrete random variables with geometrically distributed marginals.
As in Sec. 3.2, a sensible numerical approach is to limit oneself to a tractable subclass of $\Lambda_{2}(\eta_{\textrm{G}})$.
Taking inspiration from numerical OT, consider the family of measures $\pi^{(\sigma)}$ on $[0,1]^{2}$ described by the *permutation densities*
$p_{\sigma}(x,y)\coloneqq n1_{\sigma(\lceil nx\rceil)=\lceil ny\rceil}$,
with $\sigma$ a permutation of order $n$ (that is, a bijection $\sigma:[\![n]\!]\to[\![n]\!]$).
The unit square is split into a $n\times n$ grid where each row and column has a single ‘tile’ of probability density $n$ and is 0 0 otherwise (Fig. 2 left).
Both marginal distributions of $\pi^{(\sigma)}$ are uniform on $[0,1]$ and are transformed to a probability distribution $\eta$ by pushing forward with the inverse CDF, $F_{\eta}^{-1}(\cdot)\coloneqq\inf\{x\in\mathbb{R}:F_{\eta}(x)\geq\cdot\}$.
Transforming both coordinates in this way yields a joint measure $\mu^{(\sigma)}\in\Lambda_{2}(\eta)$ which will give an unbiased estimator.
$n!$ such couplings exist for a given permutation order $n$; we aim to efficiently find the one with the lowest variance.

The permutations $\sigma\in S_{n}$ can be interpreted as *matchings* between the $n$ quantiles of the geometric distributions over the lengths of a pair of walkers (Fig. 2 centre).
With the correct choice of $\sigma$, they can ensure that e.g. if one of the walk lengths is short then the other tends to be long, diversifying the ensemble (Fig. 2 right).
Optimising $\sigma$, the approximate OT problem can be written

$$ $\sigma^{*}=\arg\min_{\sigma\in S_{n}}\sum_{q_{1}\in[\![n]\!]}\sum_{q_{2}\in[\! [n]\!]}\left[\left(\widehat{\psi}(q_{1}^{(i)})+\widehat{\psi}(\sigma(q_{1})^{( i)})\right)^{\top}\left(\widehat{\psi}(q_{2}^{(j)})+\widehat{\psi}(\sigma(q_{2 })^{(j)})\right)\right]^{2}$ (13) $$

where $\widehat{\psi}(q^{(i)})\coloneqq\mathbb{E}_{u\sim\mathcal{U}((\frac{q-1}{n},
\frac{q}{n}])}\left(\mathbb{E}_{\textrm{dirs}}{\psi}(F^{-1}_{\eta_{\textrm{G}}
}(u)^{(i)})\right)$.
$\mathcal{U}((a,b])$ is the uniform distribution on the interval $(a,b]$. Eq. 13 is a *quadratic assignment problem* (Finke et al., ; Burkard et al., ).
This family is generally NP hard, but in our case the cost function has some convenient extra symmetric structure.
In fact, the following is true.

###### Theorem 4.1 (Solving Eq. 13 exactly) .

Given a set of vectors $\{\widehat{\psi}(q^{(i,j)})\}_{q=1}^{n}\subset\mathbb{R}^{N}$,
Eq. 13 can be solved with high probability in time complexity independent of $N$. Moreover, in the special case where $i=j$, it can be solved in polynomial time in $n$ (under mild technical assumptions on the set).

*Proof sketch.* Details and definitions are in App. C.3.1.
The time complexity can be made independent of the number of nodes $N$ by performing dimensionality reduction using the celebrated Johnson-Lindenstrauss transformation (Dasgupta et al., ), which preserves pairwise dot products with high probability.
In the special case $i=j$, Eq. 13 can be rewritten as finding a permutation $\sigma$ that minimises the $L_{2}$ norm of some particular $N^{2}$-dimensional vector.
In App. C.3.1 we provide a novel algorithm to achieve this efficiently by projecting onto a sequence of random Gaussian vectors, requiring only a mild geometrical condition called $\epsilon$-separation. See Lemma C.2 for full details. ∎

More pragmatically, one can set $q_{1}=q_{2}$ to simplify to a related *linear assignment problem*, which can be solved efficiently in time $\mathcal{O}(n^{3})$ using e.g. the Hungarian algorithm (Kuhn, ).
We empirically investigate how this final approximation modifies the objective in App. C.3.
Taking the optimal permutation and corresponding coupling $\mu^{(\sigma)}$, we define *$\sigma$-coupled GRFs* as follows.

###### Definition 4.2 ( $\sigma$ -coupled GRFs) .

Besides being easy to optimise and sample from, there are also more rigorous OT motivations for the choice of $\sigma$-couplings $\mu^{(\sigma)}$.
They relate to the asymptotic behaviour of $\mu^{(\sigma)}$ as the permutation order $n\to\infty$ and the *stability of OT plans* (Villani, ).
We defer this technical point to App. D.
As in Sec. 3.2, another interesting question is whether one could couple the lengths of $m>2$ walkers.
This is challenging and has received little attention in the literature.
One possibility would be to combine $m-1$ permutations, finding a minimum-weights $m$-partite matching with all subgraphs constrained to be complete $K_{m}$.
Another approach would be to approximately solve this multi-marginal OT problem using the Sinkhorn-Knopp algorithm (Sinkhorn and Knopp, ; Cuturi, ).

Broader applicability. As a final remark, the utility of our algorithm extends to graph-based estimators beyond GRFs.
For example, it can be used to improve estimates of the *PageRank vector*, a popular measure of the importance of the nodes of a graph proposed by Page et al. () to rank websites in search engine results.
$\sigma$-couplings consistently match or beat the previous best algorithm for coupling walk lengths (Reid et al., ).
See Fig. 9 in App. E for the full results.

###### Theorem 4.1 (Solving Eq. 13 exactly) .

###### Definition 4.2 ( σ 𝜎 \sigma italic_σ -coupled GRFs) .

### 4.2 Experiments with σ 𝜎 \sigma italic_σ -coupled GRFs

We now empirically evaluate $\sigma$-coupled GRFs for variance reduction of graph node kernel estimates.
For real-world graphs, we show that lower variance unlocks better approximate inference with scalable graph-based Gaussian processes (GPs), a novel application of GRFs.

Gram matrix approximation. GRFs take the *termination probability* $p_{\textrm{halt}}$ as a hyperparameter, determining the rate of decay of the geometric distribution over walk length.
A smaller value of $p_{\textrm{halt}}$ samples longer walks and gives more accurate kernel estimates, but takes longer to run.
The optimal coupling changes depending on $p_{\textrm{halt}}$.
We consider the values $p_{\textrm{halt}}\in\{0.1,0.2,0.3,0.4,0.5\}$, finding the optimal permutation $\sigma$ in each case.
To find the optimal $\sigma$, we solve the matching problem (App. C.3)

Figure: Figure 3: Kernel estimator error vs termination probability. Insets show permutations for $p_{\textrm{halt}}\in\{0.1,0.2,0.3,0.4,0.5\}$.
Refer to caption: x1.png

for a random Erdős-Rényi graph with $N=100$ nodes, taking a permutation order $n=30$ and choosing the $2$-regularised Laplacian kernel as our target.
We use the Hungarian algorithm, averaging the cost matrix over every possible node pair $(v_{i},v_{j})\in\mathcal{N}^{2}$. Having computed couplings $\mu^{(\sigma)}$ for each value of $p_{\textrm{halt}}$, we then test the corresponding $\sigma$-coupled GRFs on a variety of real-world graphs.
Fig. 3 shows the results for cora ($N=2708$), with the rest left to App. F.1.
We plot the relative Frobenius norm error of the Gram matrix approximation $\|\mathbf{K}-\widehat{\mathbf{K}}\|_{\textrm{F}}/\|\mathbf{K}\|_{\textrm{F}}$ with walkers that are i.i.d., antithetic (Reid et al., ) or $\sigma$-coupled.
For each $p_{\textrm{halt}}$, $\sigma$-coupled GRFs give equally good or smaller kernel estimator errors.
Our OT approach often substantially outperforms antithetic termination: a data-independent, hard-coded algorithm designed specifically to improve GRFs (Reid et al., ).
We include visualisations of the optimal permutations for different values of $p_{\textrm{halt}}$ in the inset, verifying that the $\sigma$-coupling adapts to different hyperparamaters.

Novel application: $\sigma$-coupled GRFs for scalable graph-based GPs.
We now apply $\sigma$-coupled GRFs to scalable graph-based Gaussian processes (GPs), where improved estimation of the covariance function permits better approximate inference.
Scalable GPs are a novel application of GRFs that may be of independent interest (Borovitskiy et al., ; Mostowsky et al., ).

Consider the task of *probabilistic graph interpolation*.
This aims to predict unknown graph function values, along with principled uncertainty estimates, from an observed set (Pfaff et al., ).
Take mesh graphs $\mathcal{G}$ where every node $v_{i}\in\mathcal{N}$ has a normal vector $\bm{n}_{i}\in\mathbb{R}^{3}$ (Dawson-Haggerty, ).
Our task is to predict the $z$-components of a masked set, $\{(\bm{n}_{i})_{z}\}_{i=1}^{N_{\textrm{test}}}$ with $N_{\textrm{test}}=\lfloor 0.05N\rfloor$.
To achieve this, we use a graph-based GP with a heat kernel covariance function.
We compute a sparse, unbiased approximation of this kernel using GRFs with $\{16,32,64\}$ walkers that are i.i.d., antithetic (Reid et al., ) or $\sigma$-coupled.
Details of GP hyperparameter optimisation are given in App. F.2.
Fig. 4 shows the results.
For mesh graphs of different sizes (the largest as big as $8700$ nodes), we plot the relative Frobenius norm error of the Gram matrix approximation, the test root mean square error (RMSE), and the KL divergence to the true posterior.
Our variance reduction method unlocks more accurate predictions and better uncertainty quantification, sometimes by a factor of $>2$.

Figure: Figure 4: Graph GP regression. Rows show the kernel estimation error, test set RMSE, and KL-divergence between the true and approximate posterior. Lower is better. $\sigma$-coupled GRFs give better predictions and uncertainty estimates, sometimes by a factor of $>2$. Standard errors are shaded.
Refer to caption: x2.png

Probabilistic interpolation of traffic data.
We also train a scalable graph-based GP on a traffic flow dataset of the highways of San Jose, California, curated by Borovitskiy et al. () using data from Chen et al. () and OpenStreetMap.
The graph has $1016$ nodes, with speed only known at $325$.
We use $250$ randomly-chosen nodes as training data and the remainder as test data.
With this sparse, noisy dataset, $\sigma$-coupled GRFs again give substantially better predictions and uncertainty estimates.
Fig. 11 in App. F.3 shows the full results.

## 5 Discussion and outlook

OT provides a powerful, unifying paradigm for variance reduction with random features.
It offers perspectives, proof techniques and numerical algorithms for finding novel RF couplings on continuous and discrete input domains, substantially beating previous algorithms in both settings and with disparate basis functions.

Variance reduction is *not* all you need. Whilst the presence of variance reduction is unambiguous, downstream benefits tell a more nuanced story.
With GRFs for scalable GPs, variance reduction permits much better approximate inference (Sec. 4.2).
With RFFs and RLFs, this is not the case.
For instance, when approximating attention in Performers (Choromanski et al., ), *maximising* the pointwise kernel estimator variance – the ‘wrong’ OT problem – turns out to improve predictive performance after row normalisation.
This shows that, though popular, naive variance reduction is not always the right goal.
We believe this to be underappreciated in the literature.

Right framing, wrong cost function.
Therefore, we posit that OT provides the *right framing* for the problem of coupling RFs, but sometimes pointwise kernel variance is the *wrong cost function*.
This choice may not fully capture how the *joint* distribution over kernel estimates determines downstream performance.
Coupling to optimise e.g. the spectral properties of $\widehat{\mathbf{K}}$ (Choromanski et al., ; Avron et al., ) or the variance of row-normalised attention scores may prove better.
These objectives are rarely considered in the literature. Fortunately, OT provides a suite of theoretical and numerical tools achieve this; one simply modifies the cost function in Eq. 2, optimising a different characteristic of the coupling.
We hope this research will spur future work in this exciting direction.

## 6 Contributions and acknowledgements

Relative contributions.
IR conceptualised the project, proposed the coupling mechanisms in Defs 3.3 and 4.2, proved the major theoretical contributions, ran the GRF experiments and wrote the manuscript.
SM designed and ran the RFF and RLF GP experiments (Sec. 3.3), helped shape the project’s direction and made core contributions to the text.
KC acted as the senior lead, providing technical guidance and developing the algorithms for the matching problem in App. C.3.1.
RET met frequently throughout the project, giving important advice and support.
AW provided helpful discussion, supervision and feedback on the manuscript.

Acknowledgements and funding.
IR acknowledges support from a Trinity College External Studentship.
Part of the work was completed as a student researcher at Google.
SM acknowledges funding from the Vice Chancellor’s and the George and Marie Vergottis scholarship of the Cambridge Trust, and the Qualcomm Innovation Fellowship.
RET is supported by Google, Amazon, ARM, Improbable and an EPSRC grant EP/T005386/1.
AW acknowledges support from a Turing AI fellowship under grant EP/V025279/1 and the Leverhulme Trust via CFI.

We thank Bruno Mlodozeniec for his suggestion to use copulas in Sec. 3.2 and Mark Rowland for insightful discussions about multi-marginal optimal transport and the limitations of pointwise variance reduction.
Viacheslav Borovitskiy helped guide our discussion of scalable graph-based GPs and, together with Iskander Azangulov, kindly provided updated code for loading the traffic data graph in Sec. 4.2 and App. F.3.
We thank Matt Ashman and Arijit Sehanobish for their thoughtful feedback on the text, and Jihao Andreas Lin for interesting suggestions about possible GP applications.

## 7 Ethics and reproducibility

Ethics statement: Our work is foundational with no immediate ethical concerns apparent to us. However, increases in scalability provided by improvements to MC algorithms could exacerbate existing and incipient risks of machine learning, from bad actors or as unintended consequences.

Reproducibility statement: Every effort has been made to ensure the work’s reproducibility.
The core algorithms are presented in Defs 3.3 and 4.2, with exhaustive details and discussion in the Appendices.
Theoretical results are proved with full assumptions in Apps A and C.3.1, with proof sketches included in the main text for clarity.
We will make source code available if accepted.
All datasets are available online.
We give links to suitable repositories in every instance.
Where possible, results are reported with uncertainties to facilitate comparison.

## Appendix A Solving the OT problem for RFFs and RLFs

In this appendix, we provide proofs for the theoretical results in Sec. 3.1 and supplement the discussion of copula-based numerical OT solvers in Sec. 3.2.

### A.1 Proof of Lemma 3.1

We begin by proving Lemma 3.1, which formulates variance reduction for RFFs and RLFs as an optimal transport problem.
We will first reason about the simpler case of RLFs, then consider RFFs.

*Proof of Lemma 3.1*. Consider the following recently-derived result by Reid et al. [].

###### Lemma A.1 (Kernel estimator MSE for RLFs [ Reid et al. , 2023 ] ) .

When estimating the Gaussian kernel $k(\bm{x},\bm{y})\coloneqq\exp(-\|\bm{x}-\bm{y}\|_{2}^{2}/2)$ for datapoints $\{\bm{x},\bm{y}\}\subset\mathbb{R}^{d}$ using random Laplace features (synonymously, positive random features), the mean square error of the kernel estimate $\widehat{k}(\bm{x},\bm{y})$ is given by:

$$ $\textrm{MSE}(\widehat{k})=\frac{e^{-2x^{2}-2y^{2}}}{m}\left((e^{2v^{2}}-e^{v^{ 2}})+(m-1)(\rho(\bm{x},\bm{y})-e^{v^{2}})\right)$ (14) $$

where $m$ is the number of sampled random frequencies, $v\coloneqq\|\bm{x}_{i}+\bm{x}_{j}\|_{2}$ is a data-dependent scalar and $\rho(\bm{x}_{i},\bm{x}_{j})$ is the *RF-conformity*,

$$ $\rho(\bm{x},\bm{y})\coloneqq\frac{\Gamma(\frac{d}{2})}{m(m-1)}\sum_{i,j\neq i} \mathbb{E}_{\omega_{ij}}\left(\sum_{k=0}^{\infty}\frac{v^{2k}\omega_{ij}^{2k}} {2^{2k}k!\Gamma(k+\frac{d}{2})}\right).$ (15) $$

Here, $\omega_{ij}\coloneqq\|\bm{\omega}_{i}+\bm{\omega}_{j}\|_{2}$ is the norm of the resultant of a pair of distinct frequencies and $\Gamma$ is the Gamma function.

*Proof.* Reid et al. []. ∎

The simple derivation, reported in full by Reid et al. [], is based on rewriting the angular integral as a Hankel transform, yielding a Bessel function of the first kind with a known Taylor expansion.

Supposing that $\bm{\omega}_{i}\perp\bm{\omega}_{j}$, we have that $\omega_{ij}=\sqrt{\omega_{i}^{2}+\omega_{j}^{2}}$, where $\omega_{i}\coloneqq\|\bm{\omega}_{i}\|_{2}$ denotes the $L_{2}$-norm of the frequency $\bm{\omega}_{i}$.
Note that, for $\bm{\omega}_{i}$ to be marginally Gaussian, we require that the marginal distribution over its *norm* is $\chi_{d}$, a Chi distribution with $d$ degrees of freedom.
Substituting this into Eq. 14 and dropping terms unmodified by the coupling and irrelevant multiplicative factors, it is straightforward to arrive at the expression for the cost function $c_{\textrm{RLF}}$ in Eq. 6.

Finding the cost function for RFFs is only slightly more difficult.
We begin by citing another recent result by Reid et al. [].

###### Lemma A.2 (Kernel estimator MSE for RFFs [ Reid et al. , 2023 ] ) .

When estimating the Gaussian kernel $k(\bm{x},\bm{y})\coloneqq\exp(-\|\bm{x}-\bm{y}\|_{2}^{2}/2)$ for datapoints $\{\bm{x},\bm{y}\}\subset\mathbb{R}^{d}$ using random Fourier features, the mean square error of the kernel estimate $\widehat{k}(\bm{x},\bm{y})$ is given by:

$$ $\text{MSE}(\widehat{k})=\frac{1}{m}\left(\frac{(1-e^{-z^{2}})^{2}}{2}+(m-1)( \zeta(\bm{x},\bm{y})-e^{-z^{2}})\right)$ (16) $$

where $m$ is the number of samples, $\bm{z}\coloneqq\bm{x}-\bm{y}$ and $\zeta(\bm{x},\bm{y})$ is defined by

$$ $\zeta(\bm{x},\bm{y})\coloneqq\frac{1}{m(m-1)}\sum_{i,j\neq i}\mathbb{E}_{\bm{ \omega}_{i},\bm{\omega}_{j}}\left[\cos(\bm{\omega}_{i}^{\top}\bm{z})\cos(\bm{ \omega}_{j}^{\top}\bm{z})\right].$ (17) $$

*Proof.* Reid et al. []. ∎

Note the close resemblance to Eq. 15. The only term that depends on couplings between the random frequencies $\{\bm{\omega}_{i}\}_{i=1}^{m}$ is $\zeta(\bm{x},\bm{y})$, which we seek to suppress with carefully engineered correlations. From elementary trigonometry, $\cos\bm{\omega}_{i}^{\top}\bm{z}\cos\bm{\omega}_{j}^{\top}\bm{z}=\frac{1}{2}(
\cos\left((\bm{\omega}_{i}+\bm{\omega}_{j})^{\top}\bm{z}\right)+\cos\left((\bm
{\omega}_{i}-\bm{\omega}_{j})^{\top}\bm{z}\right)$, and for any coupling scheme $\bm{\omega}_{i}\pm\bm{\omega}_{j}$ is isotropic. Defining the random variables $\omega^{(+)}\coloneqq\|\bm{\omega}_{i}+\bm{\omega}_{j}\|_{2}$ and $\omega^{(-)}\coloneqq\|\bm{\omega}_{i}-\bm{\omega}_{j}\|_{2}$ and integrating out the angular part,

$$ $\zeta(\bm{x}_{i},\bm{x}_{j})=\frac{1}{m(m-1)}\sum_{i,j\neq i}\Gamma(d/2)2^{ \frac{d}{2}-2}\mathbb{E}_{\omega^{(\pm)}}\left[(\omega^{(+)}z)^{1-\frac{d}{2}} J_{\frac{d}{2}-1}(\omega^{(+)}z)+(\omega^{(+)}z)^{1-\frac{d}{2}}J_{\frac{d}{2} -1}(\omega^{(-)}z)\right]$ (18) $$

where $J_{\alpha}(x)$ is a Bessel function of the first kind, order $\alpha$.
If $\bm{\omega}_{i}\perp\bm{\omega}_{j}$, $\omega^{(+)}=\omega^{(-)}$ so their distributions are identical.
It follows that we can write

$$ $\zeta(\bm{x},\bm{y})=\frac{\Gamma(\frac{d}{2})}{m(m-1)}\sum_{i,j\neq i}\mathbb {E}_{\omega_{i},\omega_{j}}\left[\sum_{k=0}^{\infty}\frac{(-1)^{k}z^{2k}\left( \omega_{i}^{2}+\omega_{j}^{2}\right)^{k}}{2^{2k}k!\Gamma(k+\frac{d}{2})}\right].$ (19) $$

Again dropping multiplicative factors and terms unmodified by the coupling, we arrive at the RFF OT cost function $c_{\textrm{RFF}}$ specified in Eq. 6.
This completes the derivation. ∎

###### Lemma A.1 (Kernel estimator MSE for RLFs [ Reid et al. , 2023 ] ) .

###### Lemma A.2 (Kernel estimator MSE for RFFs [ Reid et al. , 2023 ] ) .

### A.2 Proof of Thm. 3.2

We now solve the OT problem formulated in Lemma 3.1 *exactly* in the special case that $m=2$ (with mild asymptotic assumptions for RFFs).
For the reader’s convenience, we copy it from the main text below.

###### Theorem A.3 (Solution to OT problem when $m=2$ ) .

Denote by $F_{\chi_{d}}(\cdot)$ the cumulative distribution function (CDF) of $\chi_{d}$.
Consider $m=2$ orthogonal frequencies with norms $(\omega_{1},\omega_{2})$.
For RLFs, the OT problem in Eq. 5 is solved by the *negative monotone* coupling

$$ $F_{\chi_{d}}(\omega_{1})+F_{\chi_{d}}(\omega_{2})=1.$ (20) $$

For RFFs, Eq. 7 ensures lower cost than any other coupling, provided $z$ is sufficiently small.

The proof of Thm. 3.2 uses ideas from optimal transport theory. In particular, it modifies arguments made for a related problem by (among others) Thorpe [], to which we direct the interested reader for further context and discussion. Before giving the proof, we establish some basic definitions and results.

###### Definition A.4 ( $c$ -monotone sets) .

Given a cost function $c:\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$, we refer to a set $\Gamma\in\mathbb{R}^{2}$ as $c$-*monotone* if for all pairs $(x_{1},y_{1}),(x_{2},y_{2})\in\Gamma$ we have that

$$ $c(x_{1},y_{1})+c(x_{2},y_{2})\leq c(x_{1},y_{2})+c(x_{2},y_{1}).$ (21) $$

It is intuitive that $c$-monotonicity should be a property of the support of Kantorovich optimal transport plan: if we could have accessed lower cost by sending $x_{1}\to y_{2}$ and $x_{2}\to y_{1}$ instead of $x_{1}\to y_{1}$ and $x_{2}\to y_{2}$, the plan would have done this instead. This is formalised as follows.

###### Lemma A.5 (Support of optimal transport plan is $c$ -monotone [ Thorpe , 2019 ] ) .

Consider $\eta\in\mathcal{P}(\mathbb{R})$, and assume that $\mu^{*}\in\Lambda(\eta)$ is the Kantorovich optimal transport plan for a continuous cost function $c(x,y)$. Then for all $(x_{1},y_{1}),(x_{2},y_{2})\in\emph{{supp}}(\mu^{*})$ we have that

$$ $c(x_{1},y_{1})+c(x_{2},y_{2})\leq c(x_{1},y_{2})+c(x_{2},y_{1}).$ (22) $$

*Proof.* Thorpe []. ∎

We are now ready to provide our proof of Thm. 3.2.

*Proof of Thm. 3.2*.
Inspecting the cost functions in Eq. 6, it is clear that in the special case $m=2$ we have that

$$ $c_{\textrm{RFF}}(\omega_{1},\omega_{2})=\sum_{k=0}^{\infty}\frac{(-1)^{k}z^{2k }\left(\omega_{1}^{2}+\omega_{2}^{2}\right)^{k}}{2^{2k}k!\Gamma(k+\frac{d}{2}) },\quad c_{\textrm{RLF}}(\omega_{1},\omega_{2})=\sum_{k=0}^{\infty}\frac{v^{2k }(\omega_{1}^{2}+\omega_{2}^{2})^{k}}{2^{2k}k!\Gamma(k+\frac{d}{2})},$ (23) $$

with $z\coloneqq\|\bm{x}-\bm{y}\|_{2}$ and $v\coloneqq\|\bm{x}+\bm{y}\|_{2}$ as usual.

First consider $c_{\textrm{RLF}}$.
We claim that for any $x_{1},x_{2},y_{1},y_{2}\geq 0$ satisfying Eq. 22 for this $c_{\textrm{RLF}}$, $x_{1}<x_{2}$ implies $y_{1}\geq y_{2}$. This is seen by observing that

$$ $\begin{multlined}c_{\textrm{RLF}}(x_{1},y_{1})+c_{\textrm{RLF}}(x_{2},y_{2})-c _{\textrm{RLF}}(x_{1},y_{2})-c_{\textrm{RLF}}(x_{2},y_{1})\\ =\sum_{k=0}^{\infty}\frac{v^{2k}}{2^{2k}k!\Gamma(k+\frac{d}{2})}\left[(x_{1}^{ 2}+y_{1}^{2})^{k}+(x_{2}^{2}+y_{2}^{2})^{k}-(x_{1}^{2}+y_{2}^{2})^{k}-(x_{2}^{ 2}+y_{1}^{2})^{k}\right]\\ =\sum_{k=0}^{\infty}\frac{v^{2k}}{2^{2k}k!\Gamma(k+\frac{d}{2})}\sum_{i=0}^{k} {k\choose i}\left[(y_{2}^{2})^{k-i}-(y_{1}^{2})^{k-i}\right]\left[(x_{2}^{2})^ {i}-(x_{1}^{2})^{i}\right].\end{multlined}c_{\textrm{RLF}}(x_{1},y_{1})+c_{ \textrm{RLF}}(x_{2},y_{2})-c_{\textrm{RLF}}(x_{1},y_{2})-c_{\textrm{RLF}}(x_{2 },y_{1})\\ =\sum_{k=0}^{\infty}\frac{v^{2k}}{2^{2k}k!\Gamma(k+\frac{d}{2})}\left[(x_{1}^{ 2}+y_{1}^{2})^{k}+(x_{2}^{2}+y_{2}^{2})^{k}-(x_{1}^{2}+y_{2}^{2})^{k}-(x_{2}^{ 2}+y_{1}^{2})^{k}\right]\\ =\sum_{k=0}^{\infty}\frac{v^{2k}}{2^{2k}k!\Gamma(k+\frac{d}{2})}\sum_{i=0}^{k} {k\choose i}\left[(y_{2}^{2})^{k-i}-(y_{1}^{2})^{k-i}\right]\left[(x_{2}^{2})^ {i}-(x_{1}^{2})^{i}\right].$ (24) $$

Supposing that $x_{1}<x_{2}$, Eq. 22 is satisfied if and only if $y_{1}\geq y_{2}$, verifying the statement.

Denote by $\Gamma_{\textrm{RLF}}\coloneqq\textrm{supp}(\mu^{*})$ the support of the optimal transport plan for the cost function $c_{\textrm{RLF}}$, and consider some point $(x_{0},y_{0})\in\Gamma$.
An immediate implication of the statement above is that

$$ $\Gamma_{\textrm{RLF}}\subset\{(x,y):x\leq x_{0},y\geq y_{0}\}\cup\{(x,y):x\geq x _{0},y\leq y_{0}\}.$ (25) $$

Let $A=[0,x_{0}]\times[y_{0},\infty)$, $B=[0,x_{0})\times[0,y_{0})$, $C=[x_{0},\infty)\times[0,y_{0}]$ and $D=(x_{0},\infty]\times(y_{0},\infty]$.
Note that $A\cup B\cup C\cup D=\mathbb{R}^{+}\times\mathbb{R}^{+}$ so $\mu^{*}(A\cup B\cup C\cup D)=1$ since the measure is normalised.
The subsets are also disjoint apart from $A\cap C=(x_{0},y_{0})$, a singleton of zero measure.(^1^11e.g. since $(x_{0},y_{0})\subset\{x_{0}\}\times\mathbb{R}^{+}$ and $\mu^{+}(\{x_{0}\}\times\mathbb{R}^{+})=p_{\chi}(x=x_{0})=0$ since the marginal measure $\chi$ is nonatomic.)
Eq. 25 implies that $\mu^{*}(B)=0=\mu^{*}(D)$, whereupon $1=\mu^{*}(A\cup C)=\mu^{*}(A)+\mu^{*}(C)=\mu^{*}(A\cup B)+\mu^{*}(C\cup B)$. Now $A\cup B=([0,x_{0}]\times\mathbb{R}^{+})\backslash(\{x_{0}\}\times[0,y_{0}))$ and the set $\{x_{0}\}\times[0,y_{0})$ is zero measure.
Therefore, $\mu^{*}(A\cup B)=\mu^{*}([0,x_{0}]\times\mathbb{R}^{+})=F_{\chi_{d}}(x_{0})$. Likewise, $\mu^{*}(C\cup B)=\mu^{*}(\mathbb{R}^{+}\times[0,y_{0}])=F_{\chi_{d}}(y_{0})$. Relabelling $(x_{0},y_{0})\in\Gamma$ by $(\omega_{1},\omega_{2})$, Eq. 20 immediately follows.
This completes the proof that negative monotone coupling minimises the kernel estimator variance for orthogonal RLFs.

Let us now turn to the case of RFFs.
The optimal transport plan $\mu^{*}$ will instead be $c_{\textrm{RFF}}$-monotone.
Unfortunately, Eq. 24 does *not* hold in general for $c_{\textrm{RFF}}$, so we cannot immediately use the same arguments to conclude that the OT plan is negative monotone.
However, it *does* hold if we just consider the first few terms of its Taylor expansion in $z$.
The following is true.

###### Lemma A.6 (Negative monotone coupling for RFFs) .

Denote by $\mu_{\textrm{NM}}$ the negative monotone coupling.
Consider a coupling $\mu^{\prime}\in\Lambda_{2}(\eta)\backslash\left\{\mu_{\textrm{NM}}\right\}$, i.e. any other feasible transport map which is not negative monotone.
There exists some constant $\delta(\mu^{\prime})>0$ such that $\mathcal{I}(\mu_{\textrm{NM}})<\mathcal{I}(\mu^{\prime})$ for all $z<\delta$ (where $\mathcal{I}(\mu)$ denotes the expectation of the cost function under $\mu$).

*Proof*. Recall that, for RFFs, the cost function is of the form

$$ $c_{\textrm{RFF}}(\omega_{1},\omega_{2})=\sum_{k=0}^{\infty}\frac{(-1)^{k}z^{2k }\left(\omega_{1}^{2}+\omega_{2}^{2}\right)^{k}}{2^{2k}k!\Gamma(k+\frac{d}{2})},$ (26) $$

and that we would like to solve $\mu^{*}=\textrm{arg min}_{\mu\in\Lambda_{2}(\eta)}\left[\mathbb{E}_{{{\omega}_
{1,2}}\sim\mu}c_{\textrm{RFF}}(\omega_{1},\omega_{2})\right].$
In general, $\mu^{*}=\mu^{*}(z,d)$.
This is a very challenging OT problem; we do not believe that a simple closed-form solution exists.
However, we can make progress expanding in $k$.
The $k=0$ and $k=1$ terms are trivial since they do not contain $\omega_{1,2}$ cross terms.
We have

$$ $\mu^{*}=\textrm{arg min}_{\mu\in\Lambda_{2}(\eta)}\left[\frac{z^{4}}{32\Gamma( 2+\frac{d}{2})}\mathbb{E}\left((\omega_{1}^{2}+\omega_{2}^{2})^{2}\right)+\sum _{k=3}^{\infty}\frac{(-1)^{k}z^{2k}}{2^{2k}k!\Gamma(k+\frac{d}{2})}\mathbb{E} \left(\left(\omega_{1}^{2}+\omega_{2}^{2}\right)^{k}\right)\right].$ (27) $$

We can solve the OT problem exactly for the first term, recovering negative monotone coupling.
That is,

$$ $\textrm{arg min}_{\mu\in\Lambda_{2}(\eta)}\left[\frac{z^{4}}{32\Gamma(2+\frac{ d}{2})}\mathbb{E}\left((\omega_{1}^{2}+\omega_{2}^{2})^{2}\right)\right]=\mu_{ \textrm{NM}}$ (28) $$

for any $z>0$ and $d$ finite.
Note also that

$$ $\textrm{max}_{\mu\in\Lambda_{2}(\eta)}\mathbb{E}\left(\left(\omega_{1}^{2}+ \omega_{2}^{2}\right)^{k}\right)=2^{k}\mathbb{E}\left(\omega^{2k}\right)$ (29) $$

because the coupling that *maximises* the expectation is positive monotone (this follows from our previous arguments almost automatically).
This is evaluated as the $k$th moment of a $\chi^{2}_{d}$ distribution, which is nothing other than $2^{k}\Gamma(\frac{k}{2}+d)/\Gamma(\frac{k}{2})$.
This means we can upper bound the magnitude of the $k$th term in the expansion by: $z^{2k}/\left(k!\Gamma(\frac{k}{2})\right)$.
So for *any* coupling (feasible transport plan) we can upper bound the magnitude of the sum on the right by: $g(z)\coloneqq\sum_{k=3}^{\infty}z^{2k}/\left(k!\Gamma(\frac{k}{2})\right)$.
This goes to 0 0 as $z\to 0$ and increases monotonically.

Consider some coupling $\mu^{\prime}\in\Lambda_{2}(\eta)\backslash\left\{\mu_{\textrm{NM}}\right\}$.
Since the minimiser is unique (if $\mu^{\prime}$ is not negative monotone, it is not a minimiser), there is a positive constant $c$ such that

$$ $c\coloneqq\mathbb{E}_{\mu^{\prime}}\left(\omega_{1}^{2}+\omega_{2}^{2}\right)^ {2}-\mathbb{E}_{\mu_{\textrm{NM}}}\left(\omega_{1}^{2}+\omega_{2}^{2}\right)^{ 2}>0.$ (30) $$

But we also have that

$$ $\left|\mathbb{E}_{\mu_{\textrm{NM}}}\sum_{k=3}^{\infty}\frac{(-1)^{k}z^{2k} \left(\omega_{1}^{2}+\omega_{2}^{2}\right)^{k}}{2^{2k}k!\Gamma(k+\frac{d}{2})} -\mathbb{E}_{\mu^{*}}\sum_{k=3}^{\infty}\frac{(-1)^{k}z^{2k}\left(\omega_{1}^{ 2}+\omega_{2}^{2}\right)^{k}}{2^{2k}k!\Gamma(k+\frac{d}{2})}\right|<2g(z).$ (31) $$

Therefore, to guarantee that $\mathcal{I}(\mu_{\textrm{NM}})<\mathcal{I}(\mu^{\prime})$, it is sufficient that

$$ $\frac{z^{4}}{32\Gamma(2+\frac{d}{2})}c>2g(z).$ (32) $$

Since $g(z)/z^{4}$ is also monotonically increasing and evaluates to 0 0 and $z=0$ (by inspecting its Taylor expansion), it is always possible to solve this to find $z=\delta(\mu^{\prime})$ such that, for $z<\delta(\mu^{\prime})$, $\mathcal{I}(\mu_{\textrm{NM}})<\mathcal{I}(\mu^{\prime})$ is guaranteed.
This holds for any dimensionality $d$ and any measure $\mu^{\prime}$ that is not negative monotone. ∎

This shows that the negative monotone coupling gives lower expected cost with RFFs than any other coupling if $z$ is sufficiently small, concluding our proof of Thm. 3.2. ∎

Comments on small $z$ for RFFs.
We have seen that for RFFs it is only possible to prove optimality of $\mu_{\textrm{NM}}$ when $z$ is small enough.
Here, we comment on why this result is nonetheless interesting and useful.
Firstly, note that in practice pairwise norm coupling still substantially suppresses kernel estimator variance, even at data lengthscales chosen by training an exact GP independent of the RF construction (Table 1).
Even if at bigger $z$ the method no longer provides the smallest possible estimator variance, it can still substantially reduce it compared to an i.i.d. coupling.
Second, from a more theoretical perspective, the large $d$, small sample regime is exactly where standard QMC methods often fail.
It is interesting that OT-driven methods can still provide theoretical guarantees in this low-sample, high-dimensionality setting.
Last, we note that, for RFF variance reduction schemes, it is very common to only guarantee gains in the asymptotic limit.
This is also the case e.g. for orthogonality [Reid et al., , Yu et al., ]: a well-established and widely-used algorithm.

###### Theorem A.3 (Solution to OT problem when m = 2 𝑚 2 m=2 italic_m = 2 ) .

###### Definition A.4 ( c 𝑐 c italic_c -monotone sets) .

###### Lemma A.5 (Support of optimal transport plan is c 𝑐 c italic_c -monotone [ Thorpe , 2019 ] ) .

###### Lemma A.6 (Negative monotone coupling for RFFs) .

### A.3 Proof of Corollary 3.4

We now prove that pairwise norm-coupled RFs (Def. 3.3) provide strictly lower kernel estimator variance than i.i.d. RFs.

*Proof of Corollary 3.4*.
Supposing that we have $m=d$ frequencies, the sum $\sum_{i,j\neq i}^{d}$ has $d(d-1)$ terms in total.
Of these, $2\lfloor\frac{d}{2}\rfloor$ correspond are negative monotone norm couplings, and the remainder are independent.
The independent terms are the same in the pairwise norm-coupled and fully i.i.d. configurations, so can be ignored.
By Thm. 3.2, we have seen that negative monotone coupling exactly solves the variance reduction OT problem for RLFs, so these variance contributions will be strictly smaller in the norm-coupled case.
It immediately follows that pairwise norm-coupled RLFs give strictly lower kernel estimator variance than orthogonal independent-norm RLFs.
For RFFs, Thm. A.6 shows that negative monotone coupling is better than i.i.d. if $z$ is small enough, so the result again follows.
∎

### A.4 Proof of Thm. 3.5

We now drop the restriction that $\widehat{\bm{\omega}}_{1}\perp\widehat{\bm{\omega}}_{2}$ and consider the variance reduction problem for $m=2$ frequencies whose respective direction is unconstrained.
We will prove Thm. 3.5, which asserts that in this case the best possible coupling is antithetic, $\bm{\omega}_{1}=-\bm{\omega}_{2}$.

*Proof of Thm. 3.5*.
Recalling the expression for RLF variance in Lemma A.1, the more general OT problem under consideration is

$$ $\mu^{*}=\arg\min_{\mu\in\Lambda_{2}(\mathcal{N})}\left[\mathbb{E}_{\bm{\omega} _{1},\bm{\omega}_{2}\sim\mu}\sum_{k=0}^{\infty}\frac{v^{2k}\|\bm{\omega}_{1}+ \bm{\omega}_{2}\|_{2}^{k}}{2^{2k}k!\Gamma(k+\frac{d}{2})}\right].$ (33) $$

The term in square parentheses is an expectation of an infinite sum, every term of which is greater than or equal to 0 0.
The sum is manifestly minimised if $\|\bm{\omega}_{1}+\bm{\omega}_{2}\|_{2}=0$, which sets every term (apart from the first) to 0 0.
This is achieved if and only if $\bm{\omega}_{1}=-\bm{\omega}_{2}$: a valid coupling called ‘antithetic sampling’.
Any other joint distribution assigns nonzero probability to the event $\|\bm{\omega}_{1}+\bm{\omega}_{2}\|_{2}>0$, so this optimal coupling is unique. ∎

### A.5 Copulas as numerical OT solvers

In the main text, we noted that finding an analytic solution to the multi-marginal OT problem for RFFs and RLFs (Eq. 5) is an open problem.
In Sec. 3.2, we briefly presented an alternative numerical approach using *copulas*.
Here, we discuss this in greater detail.

Copulas.
A copula is a multivariate cumulative distribution function (CDF) whose marginals are uniformly distributed on $[0,1]$.
By Sklar’s theorem [Sklar, ], its joint distribution can be arbitrary.
Given a copula, we can easily enforce the constraint that its marginals are $\chi_{d}$ by pushing each coordinate forward with the inverse CDF $F_{\chi_{d}}^{-1}(\cdot)$, whilst retaining necessary flexibility in the joint to reduce estimator variance.
Copulas can be used to model dependencies between random variables and are popular tools in quantitative finance [Haugh, ].

Gaussian copulas.
The general family of copulas is still intractable to optimise and sample from so we constrain ourselves to *Gaussian copulas*.
These are distributions with uniform marginals whose joint distributions are determined by multivariate Gaussians, defined below.

###### Definition A.7 (Gaussian copula) .

Let $\{g_{i}\}_{i=1}^{m}\sim\mathcal{N}(0,\bm{\Sigma})$ where $\bm{\Sigma}\in\mathbb{R}^{m\times m}$ is a correlation matrix, i.e. a positive definite matrix with unit diagonals, and let $F_{\mathcal{N}}$ be the CDF of the standard univariate Gaussian.
We say $\{u_{i}\}_{i=1}^{m}$, where $u_{i}\coloneqq F_{\mathcal{N}}(g_{i})$, is distributed according to a Gaussian copula with covariance $\bm{\Sigma}$. We use the notation $\{u_{i}\}_{i=1}^{m}\sim\text{GC}(\bm{\Sigma})$ to denote this.

Parameterising correlation matrices.
Gaussian copulas are easy to sample from since they involve sampling a multivariate Gaussian and applying the univariate Gaussian CDF.
We are therefore left with the task of finding an appropriate correlation matrix $\bm{\Sigma},$ for which we turn to numerical optimisation.
The family of $m\times m$ correlation matrices can be parameterised by a vector $\bm{\theta}\in\mathbb{R}^{m(m-1)/2}.$
In fact, there exist tractable bijections between unconstrained vectors of real numbers $\bm{\theta}\in\mathbb{R}^{m(m-1)/2}$ and lower triangular Cholesky factors $\mathbf{L}_{\bm{\theta}}$ such that $\bm{\Sigma}=\mathbf{L}_{\bm{\theta}}\mathbf{L}_{\bm{\theta}}^{\top}$ is a valid correlation matrix [Bhat and Mondal, ].
In particular, suppose that for each $i=1,\dots,N,$ and $j=1,\dots,i,$ we have $\theta_{ij}\in\mathbb{R}^{+},$ where $\theta_{ii}=1.$
Then the parameterisation we use is

$$ $L_{ij}=\begin{cases}\frac{\theta_{ij}}{s_{i}}&\text{ for }i\leq j,\\ 0&\text{ otherwise},\end{cases}$ (34) $$

where $s_{i}=\sqrt{\sum_{j=1}^{i}{\theta_{ij}^{2}}}.$
Note that, since we are directly parameterising the Cholesky factor, we can sample from the associated Gaussian copula with $\mathcal{O}(m^{2})$ computational cost.

Optimising correlation matrices.
In order to pick an appropriate correlation matrix $\bm{\Sigma},$ we optimise it directly to minimise the root mean squared error (RMSE) loss

$$ $\mathcal{L}(\bm{\theta})=\mathbb{E}_{\{u_{i}\}}\left[\sqrt{\frac{1}{N^{2}}\sum _{i,j=1}^{N}(\phi_{\text{RF}}(\bm{x}_{i})^{\top}\phi_{\text{RF}}(\bm{x}_{j})-k (\bm{x}_{i},\bm{x}_{j}))^{2}}~{}\right],$ (35) $$

where $\{u_{i}\}\sim\text{GC}(\mathbf{L}_{\bm{\theta}}\mathbf{L}_{\bm{\theta}}^{\top}).$
Note that $\phi_{\text{RF}}$ here depends on $u_{i}$ by pushing forward with $F_{\chi_{d}}^{-1}$ and using the result as random frequency norms, though we have suppressed this dependence for notational simplicity.
Assuming that $\phi_{\text{RF}}$ is differentiable with respect to $\omega_{i},$ which is the case in RFFs and RLFs, we can optimise the copula parameters $\bm{\theta}$ by estimating the loss in Eq. 35, computing its gradients with respect to $\bm{\theta},$ and updating its values accordingly.

Training curves. Fig.
5 shows an example of training curves for RFFs and RLFs using a numerically optimised Gaussian copula to learn an appropriate norm-coupling, here on the Boston dataset.
We observe that the numerically optimised copula recovers the performance of the pairwise norm coupling scheme we proposed.
This suggests that the proposed scheme may in fact be (close to) optimal.
Rigorous analytical investigation is an important avenue for future work.

Figure: Figure 5: Training RMSE losses on a split of the Boston dataset. For the ‘orthogonal + copula (learned)’ configuration, we display the raw training loss (light blue) as well as an exponential moving average (dark blue). The performance of the coupling scheme returned by optimising the copula matches the performance of the ‘orthogonal + PNC’ scheme that we proved is optimal for the $m=2$ case. The optimisation is very noisy but this can be mitigated by increasing the number of samples used in the reparameterisation trick.
Refer to caption: x3.png

###### Definition A.7 (Gaussian copula) .

## Appendix B RFF and RLF experimental details

In this appendix, we supplement the discussion in Sec. 3.3, providing more details of our experimental setup for Gram matrix estimation.
We also apply norm-coupled RFs to sparse spectrum Gaussian processes [Lázaro-Gredilla et al., ], showing that in this case variance reduction does not help downstream performance.
We provide experimental details for the Performer [Choromanski et al., ] results in Table 2.

### B.1 Details for RFF and RLF experiments

Overview.
In both our RFF and RLF experiments, we compare different coupling schemes for approximating the Gaussian kernel.
The Gaussian kernel, including a lengthscale parameter $\ell,$ an output scale variable $\sigma_{v}$ and a noise scale parameter $\sigma_{n}$, takes the form

$$ $k(x_{i},x_{j})=\sigma_{v}^{2}\exp\left(-\frac{1}{2\ell^{2}}||\bm{x}_{i}-\bm{x} _{j}||_{2}^{2}\right).$ $$

Our baselines include standard methods for sampling random frequency vectors for use within RFFs and RLFs: i.i.d. sampling and Halton sequences [Halton, ].
In addition, for both settings, we consider ensembles of frequency vectors that are coupled to have orthogonal directions but i.i.d. lengths.
For a dataset of dimension $d,$ for RFFs we use ensembles of $d$ orthogonal vectors.
For RLFs we use ensembles of $2d$ vectors, including $d$ orthogonal basis vectors and their $d$ antiparallel vectors.

Selecting kernel hyperparameters.
We want to compare our coupling schemes using realistic kernel hyperparameter values, which we determine as follows.
A realistic application setting for RFFs is within GPs for probabilistic regression.
Therefore, we first fit a GP on a tractable subset of the data, specifically a maximum of 256 randomly chosen datapoints, to select appropriate parameters $\ell,\sigma_{v}$ and $\sigma_{n}.$
We optimise the exact GP marginal likelihood with respect to these hyperparameters, and subsequently fix them.
On the other hand, it is well-documented that RLFs suffer from poor estimator concentration when the data norm becomes large because of the exponential function in the feature map (Eq. 4); see e.g. Thm. 4 and App. F6 of Choromanski et al. [] or Thm. 4.3 of Likhosherstov et al. [], where the authors bound the $L_{2}$-norm of queries and keys.
This is anecdotally responsible for the deterioration in performance of Performers when networks become very deep.
To reflect this fact and choose a data regime where vanilla RLFs can perform reasonably well (so we can assess any gains from our coupling), we set the lengthscale $\ell$ to two twice the average summed norm of the data, namely

$$ $\ell=\frac{2}{N_{\textrm{train}}^{2}}\sum_{i,j=1}^{N_{\textrm{train}}}||\bm{x} _{i}+\bm{x}_{j}||_{2},$ (36) $$

over the training set.
We train the rest of the kernel parameters ($\sigma_{v}$ and $\sigma_{n}$) to maximise the marginal likelihood of the data under the exact GP.

Splitting procedure.
To obtain mean evaluation metrics and standard errors, we evaluate the methods on multiple random splits as follows.
For each dataset, we conduct cross validation with 20 splits, splitting each dataset into a training and a test set.
Because we train an exact GP to determine kernel hyperparameters and evaluate its predictive NLL, we need to limit the number of datapoints used in both the training and the test set.
We set them to a maximum of 256 points each by sub-sampling at random without replacement.
After training the GP, we evaluate the metrics on the test set, and repeat this procedure for all 20 splits.

Optimisation details.
We train the exact GP using the Adam optimiser [Kingma and Ba, ], using a learning rate of $10^{-2}.$
The exact GP optimisation stage converges around $1000$ steps, and we run it up to $5000$ steps.

### B.2 PNC RFFs for Gaussian processes

Kernel and posterior approximations.
Suppose that we have drawn an ensemble of frequency vectors from which we construct random features $\{\phi_{\text{RF}}(\bm{x}_{i})_{i=1}^{N}\}$.
Group these in a large design matrix $\bm{\Phi}$,

$$ $\bm{\Phi}\coloneqq[\phi_{\text{RF}}(\bm{x}_{i})]_{i=1}^{N},$ (37) $$

where $\phi_{\text{RF}}(\bm{x}_{i})$ of course depends on the ensemble frequencies $\{\bm{\omega}_{i}\}_{i=1}^{m}$ (suppressed for notational compactness).
For RFFs $\bm{\Phi}\in\mathbb{R}^{N\times 2m}$ whereas for RLFs $\bm{\Phi}\in\mathbb{R}^{N\times m}$.
We estimate the Gram matrix by $\widehat{\mathbf{K}}\coloneqq\bm{\Phi}\bm{\Phi}^{\top}$.

As noted by Lázaro-Gredilla et al. [], this RF kernel approximation is exactly equivalent to a linear model, namely

$$ $\bm{y}=\bm{w}\bm{\Phi}+\bm{\epsilon},$ (38) $$

where $\bm{w}\sim\mathcal{N}(\bm{0},\mathbf{I})$ and $\bm{\epsilon}\sim\mathcal{N}(\bm{0},\sigma_{n}^{2}\mathbf{I}).$
The prior covariance of this linear model is $\bm{\Phi}^{\top}\bm{\Phi}+\sigma_{n}^{2}\mathbf{I},$ which is, by construction, equal in expectation to the exact covariance produced by the kernel, namely $\mathbf{K}+\sigma_{n}^{2}\mathbf{I}$.
The predictive means of the approximate linear model and corresponding exact model are

$$ $\displaystyle\bm{\mu}_{\text{approx}}$ $\displaystyle=\bm{\Phi}_{p}\left(\frac{1}{\sigma_{n}^{2}}\bm{\Phi}_{d}\bm{\Phi }_{d}^{\top}+\bm{I}\right)^{-1}\bm{\Phi}_{d}\bm{y},$ (39) $\displaystyle\bm{\mu}_{\text{exact}}$ $\displaystyle=\mathbf{K}_{pd}(\mathbf{K}_{dd}+\sigma_{n}^{2}\mathbf{I})^{-1} \bm{y},$ (40) $$

whereas the predictive covariances are

$$ $\displaystyle\mathbf{C}_{\text{approx}}$ $\displaystyle=\bm{\Phi}_{p}^{\top}\left(\frac{1}{\sigma_{n}^{2}}\bm{\Phi}_{d} \bm{\Phi}_{d}^{\top}+\bm{I}\right)^{-1}\bm{\Phi}_{p}+\sigma_{n}^{2}\mathbf{I},$ (41) $\displaystyle\mathbf{C}_{\text{exact}}$ $\displaystyle=\mathbf{K}_{pp}-\mathbf{K}_{pd}(\mathbf{K}_{dd}+\sigma_{n}^{2} \mathbf{I})^{-1}\mathbf{K}_{dp}.$ (42) $$

Here, $\bm{\Phi}_{d}$ and $\bm{\Phi}_{p}$ are the design matrices corresponding to the training inputs and prediction outputs respectively, $\mathbf{K}_{dd}$ is the covariance matrix corresponding the training inputs, $\mathbf{K}_{pp}$ is the covariance matrix corresponding to the prediction inputs and $\mathbf{K}_{pd}$ and $\mathbf{K}_{dp}$ are the cross-covariance matrices between the training and prediction datapoints.
These models become exactly equivalent in the limit of an infinite number of features $m$ since the kernel approximation becomes exact.

We have seen that pairwise norm coupling improves the approximation of the Gram matrices $\mathbf{K}$.
In particular, we are able to suppress the variance of each pointwise kernel estimate (Thm. 3.2 and Corr. 3.4), and therefore the relative Frobenius norm error between the true and approximate Gram matrices (Table 1).
In light of the discussion above, it would be natural to assume that this would result in more accurate approximations of the predictive mean and covariance.
However, in the following section we will see that surprisingly this is *not* the case.

Evaluating posterior approximation quality.
Table 3 takes the RFs from Sec. 3.3, where we found that our coupling can substantially improve the quality of kernel estimation.
It then reports the KL divergence between the exact predictive posterior and the approximate predictive posteriors computed with RFFs and RLFs, respectively.
To be clear, Eqs 39 and 41 can be rewritten

$$ $\bm{\mu}_{\text{approx}}=\widehat{\mathbf{K}}_{pd}(\widehat{\mathbf{K}}_{dd}+ \sigma_{n}^{2}\bm{I})^{-1}\bm{y},~{}~{}~{}\mathbf{C}_{\text{approx}}=\widehat{ \mathbf{K}}_{pp}-\widehat{\mathbf{K}}_{pd}(\widehat{\mathbf{K}}_{dd}+\sigma_{n }^{2}\mathbf{I})^{-1}\widehat{\mathbf{K}}_{dp}$ (43) $$

where $\widehat{\mathbf{K}}_{dd}={\mathbf{\Phi}}_{d}^{\top}{\mathbf{\Phi}}_{d}$,
$\widehat{\mathbf{K}}_{pd}={\mathbf{\Phi}}_{d}^{\top}{\mathbf{\Phi}}_{p}$,
$\widehat{\mathbf{K}}_{pd}={\mathbf{\Phi}}_{p}^{\top}{\mathbf{\Phi}}_{d}$ and $\widehat{\mathbf{K}}_{pp}={\bm{\Phi}}_{p}^{\top}{\bm{\Phi}}_{p}.$
It is then straightforward to compute the KL divergence between Gaussian distributions with means $\bm{\mu}_{\text{approx}}$, $\bm{\mu}_{\text{exact}}$ and covariances $\mathbf{C}_{\text{approx}}$, $\mathbf{C}_{\text{exact}}$.

**Table 3: KL divergences (nats per datapoint) between exact and approximate GP predictive posterior on held out test sets. Note that we use $d$ features for RFFs and $2d$ features for RLFs. Reported errors are equal to two standard errors, i.e. $98\%$ confidence intervals, computed by averaging across splits. Note that differences *between* data splits are responsible for large errors; *within* each split we take enough trials that the standard errors become small. For every individual data split, orthogonal is statistically significantly better than i.i.d. but there is no additional benefit from PW norm-coupled. Results differ between splits because we consider a different subset of the dataset.**
| Fourier Features | Concrete | Abalone | CPU | Power | Airfoil | Boston |
| --- | --- | --- | --- | --- | --- | --- |
| i.i.d. | $1.491$ $\pm 0.106$ | $0.013$ $\pm 0.001$ | $1.570$ $\pm 0.417$ | $0.150$ $\pm 0.007$ | $0.357$ $\pm 0.025$ | $2.128$ $\pm 0.197$ |
| Halton | $1.548$ $\pm 0.104$ | $0.014$ $\pm 0.001$ | $1.596$ $\pm 0.419$ | $0.138$ $\pm 0.006$ | $0.356$ $\pm 0.024$ | $2.263$ $\pm 0.204$ |
| Orthogonal | $1.166$ $\pm 0.113$ | $0.004$ $\pm 0.000$ | $1.635$ $\pm 0.423$ | $0.029$ $\pm 0.002$ | $0.235$ $\pm 0.017$ | $1.990$ $\pm 0.191$ |
| + PNC | $1.168$ $\pm 0.113$ | $0.004$ $\pm 0.000$ | $1.589$ $\pm 0.416$ | $0.029$ $\pm 0.002$ | $0.235$ $\pm 0.017$ | $1.985$ $\pm 0.190$ |
| Laplace Features | Concrete | Abalone | CPU | Power | Airfoil | Boston |
| i.i.d. | $0.552$ $\pm 0.064$ | $0.028$ $\pm 0.003$ | $5.925$ $\pm 1.961$ | $0.199$ $\pm 0.011$ | $0.230$ $\pm 0.024$ | $0.759$ $\pm 0.068$ |
| Halton | $0.574$ $\pm 0.064$ | $0.024$ $\pm 0.003$ | $5.811$ $\pm 1.897$ | $0.146$ $\pm 0.009$ | $0.213$ $\pm 0.022$ | $0.834$ $\pm 0.074$ |
| Orthogonal | $0.486$ $\pm 0.059$ | $0.014$ $\pm 0.002$ | $5.494$ $\pm 1.774$ | $0.059$ $\pm 0.005$ | $0.165$ $\pm 0.018$ | $0.679$ $\pm 0.058$ |
| + PNC + antithetic | $0.482$ $\pm 0.058$ | $0.014$ $\pm 0.002$ | $5.468$ $\pm 1.780$ | $0.050$ $\pm 0.006$ | $0.165$ $\pm 0.018$ | $0.673$ $\pm 0.058$ |

Surprisingly, we find that, even though our couplings improve the accuracy of kernel approximation, the approximate mean and covariance and hence the KL divergence to the true posterior do not necessarily improve.
In Table 1 we routinely see variance reductions of $10$-$20$%, but even with very many trials this is not reflected in the data splits used for Table 3.

The reason for this experimental finding is that, as is clear in Eq. 43, the posterior is highly nonlinear in $\widehat{\mathbf{K}}$.
This is on account of the presence of matrix multiplications and inversions.
These nonlinear operations on the Gram matrix entries mean that, despite our *pointwise* estimates being unbiased, $\bm{\mu}_{\text{approx}}$ and $\mathbf{C}_{\text{approx}}$ are in fact biased.
We have achieved our objective of variance reduction of $\widehat{\mathbf{K}}$ – with pairwise norm coupling, it is both theoretically mandated and empirically observed.
But clearly this does not necessitate variance (or bias) reduction for $\bm{\mu}_{\text{approx}}$ and $\mathbf{C}_{\text{approx}}$.
The relationship between the distribution of $\widehat{\mathbf{K}}$ and the distribution of the approximate posterior is more complex.

Variance reduction does not always help predictive performance.
To sharpen this point, we now take a *single* data split of the power dataset and plot the the kernel approximation RMSE (i.e. kernel estimator variance), as well as various quantities of predictive interest, against the number of features $m$.
We exclude the Halton coupling (which is consistently worse than ‘orthogonal’) for clarity of presentation.
Fig. 6 shows the results.

The left hand panel confirms that we have achieved our stated objective of variance reduction.
Moreover, for all couplings the quality of approximation improves as we introduce more ensembles of size $d$.
Reading left to right, the other three panels show: (i) the KL divergence between the exact and approximate GP *predictive posteriors* (as in Table 3), (ii) the predictive RMSE on a held out test set, and (iii) the KL divergence between the exact and approximate GP *priors*.
In every instance, orthogonality provides a substantial gain but, despite the encouraging results for kernel RMSE, there is no additional benefit from PNC.
However, it would be wrong to draw the simplistic conclusion that PNC does not give large enough variance savings to see downstream gains: comparable reductions from orthogonality yield substantial improvements.
The problem is more fundamental, relating to how the joint distribution of kernel estimates – beyond just the second moment of its pointwise entries – interacts with nonlinear operations like matrix multiplication and inversion.
Table 4 is a companion to Table 3, reporting predictive RMSEs as in Fig. 6 for the other datasets.

**Table 4: Predictive RMSEs (normalised to RMSE of i.i.d.) of the approximate GP predictive posteriors on held out test sets. Reported errors are equal to two standard errors, i.e. $98\%$ confidence intervals, computed by averaging across splits. Substantially lower kernel estimator variance does *not* guarantee better predictive performance, though performance tends to be no worse.**
| Fourier Features | Concrete | Abalone | CPU | Power | Airfoil | Boston |
| --- | --- | --- | --- | --- | --- | --- |
| i.i.d. | $1.000$ $\pm 0.027$ | $1.000$ $\pm 0.043$ | $1.000$ $\pm 0.179$ | $1.000$ $\pm 0.016$ | $1.000$ $\pm 0.025$ | $1.000$ $\pm 0.067$ |
| Halton | $1.013$ $\pm 0.027$ | $1.001$ $\pm 0.043$ | $1.005$ $\pm 0.179$ | $0.991$ $\pm 0.016$ | $0.999$ $\pm 0.025$ | $1.021$ $\pm 0.068$ |
| Sobol | $1.010$ $\pm 0.025$ | $1.000$ $\pm 0.045$ | $1.004$ $\pm 0.138$ | $0.991$ $\pm 0.022$ | $0.997$ $\pm 0.029$ | $1.027$ $\pm 0.058$ |
| Orthogonal | $0.917$ $\pm 0.026$ | $0.994$ $\pm 0.044$ | $1.019$ $\pm 0.182$ | $0.915$ $\pm 0.019$ | $0.930$ $\pm 0.025$ | $0.974$ $\pm 0.066$ |
| + PNC | $0.917$ $\pm 0.026$ | $0.994$ $\pm 0.044$ | $1.033$ $\pm 0.180$ | $0.915$ $\pm 0.019$ | $0.930$ $\pm 0.025$ | $0.975$ $\pm 0.066$ |
| Laplace Features | Concrete | Abalone | CPU | Power | Airfoil | Boston |
| i.i.d. | $1.000$ $\pm 0.028$ | $1.000$ $\pm 0.043$ | $1.000$ $\pm 0.192$ | $1.000$ $\pm 0.017$ | $1.000$ $\pm 0.029$ | $1.000$ $\pm 0.073$ |
| Halton | $1.007$ $\pm 0.028$ | $0.997$ $\pm 0.044$ | $1.003$ $\pm 0.190$ | $0.964$ $\pm 0.017$ | $0.989$ $\pm 0.029$ | $1.019$ $\pm 0.075$ |
| Sobol | $1.007$ $\pm 0.025$ | $0.997$ $\pm 0.045$ | $0.998$ $\pm 0.176$ | $0.964$ $\pm 0.022$ | $0.990$ $\pm 0.039$ | $1.017$ $\pm 0.071$ |
| Orthogonal | $0.979$ $\pm 0.027$ | $0.990$ $\pm 0.044$ | $1.003$ $\pm 0.187$ | $0.899$ $\pm 0.018$ | $0.958$ $\pm 0.028$ | $0.979$ $\pm 0.071$ |
| + PNC + antithetic | $0.981$ $\pm 0.027$ | $0.991$ $\pm 0.044$ | $1.006$ $\pm 0.187$ | $0.909$ $\pm 0.018$ | $0.958$ $\pm 0.028$ | $0.982$ $\pm 0.071$ |

Figure: Figure 6: Results on a single train-test split of the power dataset, showing kernel approximation RMSE, KL divergence to the true predictive posterior, test RMSE of predictions and KL divergence to the true prior. We have successfully reduced the variance of the kernel approximation, but this may not help downstream metrics. Standard errors are shaded but too small to easily see.
Refer to caption: x4.png

### B.3 Norm-coupled RLFs in Performers

In Sec. 3.3, we used random Laplace features for attention approximation in Performers [Choromanski et al., ], a type of efficient transformer that estimates softmax using a low rank decomposition.
Here, we discuss the results in greater detail.

Derivation of Eq. 8.
Let $X_{ij}$ denote the random variable $\widehat{k}(\bm{x}_{i},\bm{x}_{j})$, which is an unbiased estimate of $\exp(\bm{x}_{i}^{\top}\bm{x}_{j})$ constructed using features $\{\exp({\|\bm{x}_{i}\|^{2}}/{2})\phi_{\textrm{RLF}}(\bm{x}_{i})\}_{i=1}^{N}$.
Unbiasedness follows trivially from the fact that RLFs give an unbiased estimate of the Gaussian kernel.
Normalising by the sum of attention scores, let us define $\widehat{a}_{ij}\coloneqq X_{ij}/\sum_{j}X_{ij}$.
Let $X_{ij}=\mu_{ij}+\delta_{ij}$ with $\mu_{ij}\coloneqq\mathbb{E}(X_{ij})$ and $\mathbb{E}(\delta_{ij})=0$.
Then

$$ $\begin{multlined}\widehat{a}_{ij}^{2}=(\mu_{ij}^{2}+2\mu_{ij}\delta_{ij}+ \delta_{ij}^{2})(\sum_{k}\mu_{ik}+\sum_{k}\delta_{ik})^{-2}\\ =\frac{\mu_{ij}^{2}+2\mu_{ij}\delta_{ij}+\delta_{ij}^{2}}{N^{2}\bar{\mu}_{i}^{ 2}}\left(1-2\frac{\sum_{k}\delta_{ik}}{N\bar{\mu}_{i}}+3\frac{\sum_{kl}\delta_ {ik}\delta_{il}}{N^{2}\bar{\mu}_{i}^{2}}+\mathcal{O}(\frac{1}{N^{3}})\right), \end{multlined}\widehat{a}_{ij}^{2}=(\mu_{ij}^{2}+2\mu_{ij}\delta_{ij}+\delta_ {ij}^{2})(\sum_{k}\mu_{ik}+\sum_{k}\delta_{ik})^{-2}\\ =\frac{\mu_{ij}^{2}+2\mu_{ij}\delta_{ij}+\delta_{ij}^{2}}{N^{2}\bar{\mu}_{i}^{ 2}}\left(1-2\frac{\sum_{k}\delta_{ik}}{N\bar{\mu}_{i}}+3\frac{\sum_{kl}\delta_ {ik}\delta_{il}}{N^{2}\bar{\mu}_{i}^{2}}+\mathcal{O}(\frac{1}{N^{3}})\right),$ (44) $$

where we defined $\bar{\mu}_{i}\coloneqq\frac{1}{N}\sum_{j}\mu_{ij}$, the average groundtruth attention score across tokens.
Since we are targeting $\frac{\mu_{ij}}{N\bar{\mu_{i}}}$,

$$ $\textrm{MSE}(\widehat{a}_{ij})=\frac{1}{N^{2}\bar{\mu_{i}}^{2}}\left(\delta_{ ij}^{2}-\frac{4\mu_{ij}\delta_{ij}}{N\bar{\mu}_{i}}\sum_{k}\delta_{ik}+\frac{3 \mu_{ij}^{2}}{N^{2}\bar{\mu}_{i}^{2}}\sum_{kl}\delta_{ik}\delta_{il}\right)+ \mathcal{O}\left(\frac{1}{N^{3}}\right).$ (45) $$

As in the main text, denote $\textrm{MSE}(\widehat{a}_{i})\coloneqq\frac{1}{N}\sum_{j}\textrm{MSE}(\widehat
{a}_{ij})$, the average mean squared error over the tokens to which $i$ attends.
To better see the intuitive behaviour, suppose also that $\mu_{ij}=\bar{\mu}_{i}\hskip 1.99997pt\forall\hskip 1.99997ptj$. Then

$$ $\begin{multlined}\textrm{MSE}(\widehat{a}_{i})=\frac{1}{N^{2}\bar{\mu_{i}}^{2} }\left(\frac{1}{N}\sum_{j}\delta_{ij}^{2}-\frac{1}{N^{2}}\sum_{jk}\delta_{ij} \delta_{ik}\right)+\mathcal{O}\left(\frac{1}{N^{3}}\right)\\ =\frac{1}{N^{2}\bar{\mu}_{i}^{2}}\left(\frac{1}{N}\sum_{j}\textrm{Var}( \widehat{k}(\bm{x}_{i},\bm{x}_{j}))-\frac{1}{N^{2}}\sum_{j,k}\textrm{Cov}( \widehat{k}(\bm{x}_{i},\bm{x}_{j}),\widehat{k}(\bm{x}_{i},\bm{x}_{k}))\right)+ \mathcal{O}\left(\frac{1}{N^{3}}\right),\end{multlined}\textrm{MSE}(\widehat{a }_{i})=\frac{1}{N^{2}\bar{\mu_{i}}^{2}}\left(\frac{1}{N}\sum_{j}\delta_{ij}^{2 }-\frac{1}{N^{2}}\sum_{jk}\delta_{ij}\delta_{ik}\right)+\mathcal{O}\left(\frac {1}{N^{3}}\right)\\ =\frac{1}{N^{2}\bar{\mu}_{i}^{2}}\left(\frac{1}{N}\sum_{j}\textrm{Var}( \widehat{k}(\bm{x}_{i},\bm{x}_{j}))-\frac{1}{N^{2}}\sum_{j,k}\textrm{Cov}( \widehat{k}(\bm{x}_{i},\bm{x}_{j}),\widehat{k}(\bm{x}_{i},\bm{x}_{k}))\right)+ \mathcal{O}\left(\frac{1}{N^{3}}\right),$ (46) $$

as reported in Eq. 8.

Positive monotone coupling.
In contrast to Eq. 7, we say that random variables $(\omega_{1},\omega_{2})$ are *postive monotone-coupled* if $\omega_{1}=\omega_{2}$ almost surely.
This is a valid transport plan, the identity.
Clearly, all $m$ frequencies in an ensemble can be simultaneously positive monotone-coupled by making them equal.
It is intuitive that this should maximise the kernel estimator variance since we sample only one frequency lengthscale.
The proof is a trivial extension of the arguments made for Thm. 3.2 in App. A.2, simply taking $c_{\textrm{RLF}}\to-c_{\textrm{RLF}}$ so that the support of the OT plan swaps.
We omit it for brevity.
It is also intuitive that this will increase the covariance between kernel estimates.

Variance, covariance, and attention approximation error.
Modifying the norm coupling between frequencies changes both $\textrm{Var}(\widehat{k}(\bm{x}_{i},\bm{x}_{j}))$ and $\textrm{Cov}(\widehat{k}(\bm{x}_{i},\bm{x}_{j}))$, so $\textrm{MSE}(\widehat{a}_{i})$ can change unpredictably.
Fig. 7 shows the results for randomly $N=16$ synthetic, normally distributed $16$-dimensional keys, $\bm{x}\sim\mathcal{N}(0,\frac{1}{\sqrt{d}}\mathbf{I}_{d}$).
Strikingly, maximising the kernel estimator variance with a positive monotone (PM) coupling ends up *improving* the quality of attention estimation since it also also increases the covariance between the unnormalised scores.
This counterintuitive finding shows highlights the limitations of variance reduction as a paradigm.

Figure: Figure 7: Kernel variance $\textrm{Var}(\widehat{k}(\bm{x}_{i},\bm{x}_{j}))$, kernel covariance $\textrm{Cov}(\widehat{k}(\bm{x}_{i},\bm{x}_{j}))$ and attention mean squared error $\textrm{MSE}(a_{i})$ plotted against number of frequencies for different coupling schemes. As intended, PNC decreases the kernel estimator variance, but it also decreases the kernel estimator covariance so the attention MSE remains essentially unchanged. On the other hand, positive monotone coupling (PM) drastically increases both kernel variance (which it maximises) and kernel covariance. The latter offsets the former (Eq. 46), so that attention MSE actually drops: a remarkable, counterintuitive finding.
Refer to caption: x5.png

Performer experimental details.
To obtain the results in Table 2, we train a Performer-ViT [Dosovitskiy et al., , Choromanski et al., ] on the ImageNet (1M) dataset [Deng et al., ].
We estimate attention using $m=128$ random features that are orthogonal with norms that are: (i) i.i.d. (ii) PNC (Def. 3.3), or (iii) positive monotone-coupled (i.e. equal among blocks of size $d$).
We use a transformer with $12$ layers and heads, with hidden size $768$ and MLP dimension $3072$.
We take $16\times 16$ patches, and train with the Adam optimiser for $90$ epochs with a compound learning rate ($10^{4}$ steps linear warmup, constant, then cosine decay, with base LR $3\times 10^{-3}$ and final LR $1\times 10^{-5}$).
The batch size is $4096$.
To get the standard errors, we average between $34$ and $48$ seeds per coupling.
We see that the lower attention MSE unlocked by PM coupling improves predictive performance over the baseline.

## Appendix C Graph random features

In this appendix, we provide a self-contained introduction to graph random features (GRFs) and previously proposed techniques to reduce their variance.
We also explain the motivations behind the series of approximations used in Sec. 4.1, and where possible provide empirical evidence that these approximations for tractability and efficiency do not substantially degrade the performance of the learned $\sigma$-coupling.

### C.1 Constructing graph random features

For the reader’s convenience, we begin by providing a brief introduction to GRFs.

Graph node kernels.
Recall that *graph node kernels* $k:\mathcal{N}\times\mathcal{N}\to\mathbb{R}$ are positive definite, symmetric functions defined on pairs of nodes of $\mathcal{G}$. Note that one can also define kernels that take pairs of graphs as inputs, but these are not the object of our study.

Many of the most popular graph node kernels in the literature are functions of weighted adjacency matrices [Smola and Kondor, , Chapelle et al., ]. In particular, they are often functions of the *graph Laplacian matrix*,

$$ $\mathbf{L}\coloneqq\mathbf{D}-\mathbf{W},$ (47) $$

where $\mathbf{W}$ is a (weighted) adjacency matrix and $\mathbf{D}$ is a diagonal degree matrix (satisfying $\mathbf{D}_{ii}=\sum_{j}\mathbf{W}_{ij}$).
It is also common to consider its normalised variant,

$$ $\widetilde{\mathbf{L}}\coloneqq\mathbf{D}^{-\frac{1}{2}}\mathbf{L}\mathbf{D}^{ -\frac{1}{2}},$ (48) $$

whose spectrum is contained in $[0,2]$ [Chung, ]. We provide examples in Table 5.

**Table 5: Examples of graph node kernels. The $\mathrm{exp}$ and $\mathrm{cos}$ mappings are defined via Taylor series expansions rather than element-wise, e.g. $\exp(\mathbf{M})\coloneqq\lim_{n\to\infty}(\mathbf{I}_{N}+\mathbf{M}/n)^{n}$ and $\cos(\mathbf{M})\coloneqq\textrm{Re}(\exp(i\mathbf{M}))$. $\sigma$ and $\alpha$ are regularisers.**
| Name | Form |
| --- | --- |
| $d$-regularised Laplacian | $(\mathbf{I}_{N}+\sigma^{2}\mathbf{L})^{-d}$ |
| $p$-step random walk | $(\alpha\mathbf{I}_{N}-\mathbf{L})^{p},\alpha\geq 2$ |
| Diffusion | $\exp(-\sigma^{2}\mathbf{L}/2)$ |
| Inverse Cosine | $\cos\left(\mathbf{L}\pi/4\right)$ |

Formula for graph random features. Computing graph node kernels exactly is computationally expensive because of the $\mathcal{O}(N^{3})$ cost of e.g. matrix inversion or exponentiation.
Inspired by the success of random Fourier features in the Euclidean domain [Rahimi and Recht, ], the recently-introduced class of *graph random features* (GRFs) permits unbiased approximation of $k$ in subquadratic time [Choromanski, , Reid et al., ].
Intuitively, GRFs interpret the powers of weighted adjacency matrices in the Taylor expansions of the expressions in Table 5 as weighted sums over walks on a graph.
Instead of computing this infinite sum of walks exactly, we can sample $m$ finite walks (using importance sampling) to construct an unbiased estimate of $k$.

Concretely, GRFs compute a set of estimators $\{\phi_{\textrm{GRF}}(v_{i})\}_{i\in\mathcal{N}}$ such that $k(v_{i},v_{j})=\mathbb{E}[\phi_{\textrm{GRF}}(v_{i})^{\top}\phi_{\textrm{GRF}}
(v_{j})]$ by taking:(^2^22Strictly, for unbiased estimation of diagonal kernel entries $\{k(v_{i},v_{i})\}_{v_{i}\in\mathcal{N}}$ we should construct *two* independent sets of features, but this is found to make little practical difference [Reid et al., ] so we omit further discussion in this manuscript.)

$$ $\phi_{\textrm{GRF}}(v_{i})=\frac{1}{m}\sum_{k=1}^{m}\psi(\bm{\omega}_{k}^{(i)}),$ (49) $$

where $\bm{\omega}_{k}^{(i)}$ is the $k$th walk (out of a total of $m$) simulated out of the starting node $v_{i}$. We remind the reader that by a ‘walk’ we mean a sequence of neighbouring graph nodes.
The projection function $\psi(\cdot):\Omega\to\mathbb{R}^{N}$ maps from the set of graph random walks $\Omega\coloneqq\left\{(v_{i})_{i=1}^{l}\,|\,v_{i}\in\mathcal{N},(v_{i},v_{i+1}
)\in\mathcal{E},l\in\mathbb{N}\right\}$ to a sparse $N$-dimensional feature vector.
It is computed as follows:

$$ $\psi(\bm{\omega}^{(i)})_{q}\coloneqq\sum_{\omega_{iq}\in\Omega_{iq}}\frac{ \widetilde{\omega}(\omega_{iq})f(\textrm{len}(\omega_{iq}))}{p(\omega_{iq})} \mathbb{I}(\omega_{iq}\in\bm{\omega}^{(i)}),\quad q=1,...,N.$ (50) $$

In this expression,

- •
$\psi(\bm{\omega}^{(i)})_{q}$ is the $q$th component of the feature projection of the random walk $\bm{\omega}^{(i)}$ (itself a sequence of neighbouring graph nodes beginning with $v_{i}$).
- •
$\Omega_{iq}$ is the set of *all* graph walks between nodes $v_{i}$ and $v_{q}$, of which each $\omega_{iq}$ is a member.
- •
$\mathbb{I}(\cdot)$ is the indicator function that evaluates to $1$ if its argument is true and 0 0 otherwise.
- •
$\omega_{iq}\in\bm{\omega}^{(i)}$ is the condition that $\omega_{iq}$, a particular walk between nodes $v_{i}$ and $v_{q}$, is a *prefix subwalk* of $\bm{\omega}^{(i)}$, the random walk we actually sample.(^3^33Meaning that the walk $\bm{\omega}^{(i)}$ from node $v_{i}$ initially follows $\omega_{iq}$, then optionally continues to visit further nodes. Note the subtle difference in usage of the symbol $\in$ compared to in the expression $\omega_{iq}\in\Omega_{iq}$, where it means ‘a member of the set of walks’ rather than ‘a prefix subwalk’.)
- •
$f:\mathbb{N}\to\mathbb{R}$ is the *modulation function*, which controls the particular graph node kernel we choose to approximate.
- •
$\textrm{len}(\omega_{iq})$ is the length of graph walk $\omega_{iq}$.
- •
$p(\omega_{iq})$ is the *marginal probability* of sampling a random walk $\omega_{iq}$, which is known from the sampling strategy.
- •
$\widetilde{\omega}(\omega_{iq})$ is a function that returns the *products of weights* of the edges traversed by $\omega_{iq}$.

Intuitively, to construct $\psi(\bm{\omega}^{(i)})$ using a random walk $\bm{\omega}^{(i)}$, we look at the node that it visits at each of the $\textrm{len}(\bm{\omega}^{(i)})+1$ timesteps before termination.
At every node, we add a contribution to the feature at the corresponding coordinate that depends on (i) the product of edge weights traversed so far, (ii) the known marginal probability of sampling the walk so far, and (iii) a ‘modulation function’ $f$ applied to the number of steps taken so far.
We refer the reader to the original works of Choromanski [] and Reid et al. [] for further technical details, experimental results and a proof of unbiasedness.

### C.2 Antithetic termination

In this section, we briefly describe *antithetic termination* [Reid et al., ]: a previously-proposed variance reduction algorithm that couples the lengths of pairs of random walks.

Consider a random walker that terminates with probability $p$ at every timestep.
In practice, this can be achieved by drawing a ‘termination random variable’ $t$ from a uniform distribution on $[0,1]$, $t\sim\mathcal{U}([0,1])$, and ending the walk if $t<p$.
Now consider two such walkers with corresponding termination random variables $t_{1},t_{2}$.
If they are independent, $t_{1}$ and $t_{2}$ are drawn independently.
Antithetic termination instead proposes to induce a nontrivial joint distribution by first drawing $t_{1}\sim\mathcal{U}([0,1])$ as usual, and then setting $t_{2}=\mod_{1}(t_{1}+\frac{1}{2})$.
$t_{2}$ still follows a uniform marginal distribution so we preserve unbiasedness, but this coupling forces walkers to terminate at different timesteps (if $p_{\textrm{halt}}<0.5$) which the authors prove reduces the estimator variance under certain assumptions.
This is one example of a possible hand-crafted coupling between random walk lengths that improves estimator concentration.

It is also possible to improve the accuracy of estimators using graph random walks by coupling walker *directions* [Reid et al., , Luo, ].
This approach is complementary to our own and can be combined with it.
Formulating coupled walker directions in terms of OT is an interesting and involved technical challenge; we defer its study to future work.

### C.3 Approximating the OT problem: details and discussion for Sec. 4.1

In this appendix, we discuss the steps to formulate the OT matching problem in Sec. 4.1.
We will especially focus on the approximations needed to make the objective in Eq. 12 tractable, describing their motivation and where possible investigating how they modify the original objective.

To begin, we remind the reader of the OT formulation of the variance reduction problem for GRFs:

$$ $\begin{multlined}\textrm{minimise }\mathbb{E}_{(l_{1}^{(i)},l_{2}^{(i)})\sim \mu,(l_{1}^{(j)},l_{2}^{(j)})\sim\mu}\left(\mathbb{E}_{\textrm{dirs}}\left[ \left(\psi(\bm{\omega}_{1}^{(i)})+\psi(\bm{\omega}_{2}^{(i)})\right)^{\top} \left(\psi(\bm{\omega}_{1}^{(j)})+\psi(\bm{\omega}_{2}^{(j)})\right)\right]^{2 }\right)\\ \textrm{for }\mu\in\Lambda_{2}(\eta_{\textrm{G}}),\end{multlined}\textrm{ minimise }\mathbb{E}_{(l_{1}^{(i)},l_{2}^{(i)})\sim\mu,(l_{1}^{(j)},l_{2}^{(j) })\sim\mu}\left(\mathbb{E}_{\textrm{dirs}}\left[\left(\psi(\bm{\omega}_{1}^{(i )})+\psi(\bm{\omega}_{2}^{(i)})\right)^{\top}\left(\psi(\bm{\omega}_{1}^{(j)}) +\psi(\bm{\omega}_{2}^{(j)})\right)\right]^{2}\right)\\ \textrm{for }\mu\in\Lambda_{2}(\eta_{\textrm{G}}),$ (51) $$

where $\Lambda_{2}(\eta_{\textrm{G}})$ is the set of joint distributions on $\mathbb{N}^{2}$ with geometrically distributed marginals, $l_{1}^{(i)}$ is the length of walk $\bm{\omega}_{1}^{(i)}$ out of node $v_{i}$, $\mathbb{E}_{\textrm{dirs}}$ averages the *directions* of walkers of a particular length, and $\psi(\bm{\omega}_{k}^{(i)})$ is the *projection function* that maps from random walks to feature vectors (see Eq. 9).

Averaging walker trajectories.
As noted in the main text, the cost function in Eq. 51 is not analytically tractable.
We can approximate it using MC sampling of graph random walks.
To do so, it will be convenient to introduce the *direction-averaged feature vector* $\widetilde{\psi}(\cdot):\mathbb{N}\to\mathbb{R}^{N}$ satisfying

$$ $\widetilde{\psi}(l^{(i)})\coloneqq\mathbb{E}_{\textrm{dirs}}[\psi(\bm{\omega}^ {(i)})].$ (52) $$

This computes the *average* feature projection over all walks of a given length $l^{(i)}$ out of node $v_{i}$.
It is straightforward to estimate by simulating a finite number of random walks.
If we move the expectation over walker directions *inside* the square bracket, the (now approximate) OT problem becomes

$$ $\textrm{minimise }\mathbb{E}_{(l_{1}^{(i)},l_{2}^{(i)})\sim\mu,(l_{1}^{(j)},l_ {2}^{(j)})\sim\mu}\left[\left(\widetilde{\psi}(l_{1}^{(i)})+\widetilde{\psi}(l _{2}^{(i)})\right)^{\top}\left(\widetilde{\psi}(l_{1}^{(j)})+\widetilde{\psi}( l_{2}^{(j)})\right)\right]^{2}\textrm{ for }\mu\in\Lambda_{2}(\eta_{\textrm{G} }).$ (53) $$

The price we pay is that we have ignored some higher-order corrections from the variance of the directions of the walks, but as we have seen in Sec. 4.2 this does not prevent us from obtaining an effective, computationally cheap coupling.

Using $\sigma$-couplings.
Unfortunately, Eq. 53 remains intractable.
To make progress, we must restrict our search-space of joint distributions from $\Lambda_{2}(\eta_{\textrm{G}})$ – the set of *all* joint distributions with geometrically-distributed marginals – to a tractable subclass.
Taking inspiration from approaches in numerical OT, we have considered the class of *permutation densities* defined by $p_{\sigma}(x,y)\coloneqq n1_{\sigma(\lceil nx\rceil)=\lceil ny\rceil}$,
with $\sigma$ a permutation of order $n$ (that is, a bijection $\sigma:[\![n]\!]\to[\![n]\!]$), transformed coordinate-wise by pushing forward with the appropriate inverse CDF.
This may initially seem an arbitrary choice, but it is motivated by three important observations.
First, we have seen that this class of joint distributions frames the minimisation as a *matching problem* which, under further simplifying assumptions, we can efficiently solve using techniques from linear programming.
Second, it is straightforward to sample from so the learned coupling will be practically useful.
Third, it is well-motivated by results in discrete OT; see App. D for concrete theoretical guarantees.
Once again, we have seen that in practice, even at modest permutation order $n$, this class contains couplings that can substantially outperform both i.i.d. walkers and antithetic termination [Reid et al., ].

We now replace $\Lambda_{2}(\eta_{\textrm{G}})$ in Eq. 53 by the set of measures $\{\mu^{(\sigma)}\}_{\sigma\in S_{n}}$ obtained by pushing permutation densities through the required inverse CDF.
$S_{n}$ is the set of permutations $[\![n]\!]\to[\![n]\!]$.
Denote by $\widehat{\psi}(\cdot):[\![n]\!]\to\mathbb{R}^{N}$ the function

$$ $\widehat{\psi}(q^{(i)})\coloneqq\mathbb{E}_{u\sim\mathcal{U}((\frac{q-1}{n}, \frac{q}{n}])}\left(\widetilde{\psi}(F^{-1}_{\eta_{\textrm{G}}}(u)^{(i)}) \right).$ (54) $$

This takes the direction-averaged feature vector for node $v_{i}$ (see Eq. 52) and evaluates a *further* average over walk lengths corresponding to the $q$th quantile of the geometric distribution $\eta_{\textrm{G}}$.
Again, this is straightforward to estimate by simulating walks.
With a further approximation of the objective, we can consider

$$ $\sigma^{*}=\arg\min_{\sigma\in S_{n}}\sum_{q_{1}\in[\![n]\!]}\sum_{q_{2}\in[\! [n]\!]}\left[\left(\widehat{\psi}(q_{1}^{(i)})+\widehat{\psi}(\sigma(q_{1})^{( i)})\right)^{\top}\left(\widehat{\psi}(q_{2}^{(j)})+\widehat{\psi}(\sigma(q_{2 })^{(j)})\right)\right]^{2}$ (55) $$

which is a matching problem.

Solving the matching problem.
Though we can now efficiently estimate the cost function with MC, it is is *still* difficult to solve the matching problem in Eq. 55 efficiently because it is quadratic [Finke et al., , Burkard et al., ].
Minimum weights quadratic bipartite matching problems are generally NP-hard, but here the problem has extra symmetry that makes it simpler.

In App. C.3.1, we discuss possible approaches to directly solve the matching problem in Eq. 55.
In particular, we prove that it can be solved with high probability at a time complexity independent of the number of nodes $N$, and give a polynomial time algorithm for the special case of diagonal kernel entries $k(v_{i},v_{i})$.
Notwithstanding this progress, as a simpler first recourse we can just make a further approximation by restricting the sum to $q_{1}=q_{2}$ diagonal terms.
This is the approximation used to get the results in the main text.
Doing so, we arrive at

$$ $\sigma^{*}=\arg\min_{\sigma\in S_{n}}\sum_{q\in[\![n]\!]}\left[\left(\widehat{ \psi}(q^{(i)})+\widehat{\psi}(\sigma(q)^{(i)})\right)^{\top}\left(\widehat{ \psi}(q^{(j)})+\widehat{\psi}(\sigma(q)^{(j)})\right)\right]^{2}.$ (56) $$

This is now a minimum-weights bipartite matching problem with *linear* weights.
Even though the search space of possible permutations of order $n$ is of size $n!$, it is well known that it can be solved efficiently and exactly using algorithms from linear programming (e.g. the Hungarian algorithm [Kuhn, ]) in time $\mathcal{O}(n^{3})$.
As stated in Sec. 4.1 of the main text, we can use this optimal permutation to construct *$\sigma$-coupled GRFs*.

This approximation might seem unreasonable since it involves discarding all the off-diagonal $q_{1}\neq q_{2}$ terms from the objective, but for modest permutation order $n$ we can empirically test the effect this has on the quality of the obtained coupling.
We achieve this by comparing to the *best* possible permutation obtained by exhaustively searching the $n!$ possibilities.

Fig. 8 plots the results for GRF estimates of the $2$-regularised Laplacian kernel on the eurosis graph, comparing the approximate and exact objectives.
The former can be computed efficiently but the latter is evaluated at a cost exponential in $n$, so we limit to $n=10$ and smaller.
We plot the kernel approximation error (normalised by the i.i.d. result) for different permutation orders.
We also plot the *Cayley distance* between the permutations, which measures the difference between them.
As $n$ grows, the permutations become more different so we are clearly not finding the same coupling.
Nonetheless, we do not see a statistically significant difference in the amount of variance reduction.
This suggests that our linear proxy objective is reasonable.
We defer more detailed analytic study of this curious phenomenon to future work.

Figure: Figure 8: Relative Frobenius norm error (normalised by i.i.d. result) with permutations of order $n$ learned using exact and approximate objectives. ‘Cayley distance’ shows the difference between the learned permutations. Although dropping off-diagonal terms from the objective in Eq. 55 changes the $\sigma$-coupling, it does not significantly detriment variance reduction.
Refer to caption: x6.png

How important are the approximations?
This section has provided comprehensive details about the series of approximations required to make the GRF variance reduction OT problem tractable and solve it efficiently.
We close by emphasising that, even if these approximations mean that we are not *exactly* solving the original OT problem, we clearly nonetheless obtain a computationally efficient coupling that offers much greater variance reduction compared to heuristically-motivated techniques.
Analytic intractability is common in OT (Sec. 3 being an obvious exception), but fortunately this framework also equips us with a powerful arsenal of numerical tools to achieve our goal of finding effective, cheap sample couplings.

#### C.3.1 Proof of Thm. 4.1 : towards an analytic solution to Eq. 55

To assuage any possible dissatisfaction with the $q_{1}=q_{2}$ approximation in Eq. 56, in this section we discuss progress towards solving the optimisation problem in Eq. 55 *exactly*. In particular, we prove Thm. 4.1:
that it can be solved with high probability and

- 1.
at a time complexity independent from the dimensionality of vectors $\{\widehat{\psi}(q^{(i,j)})\}_{q=1}^{n}$ (at the cost of a one-time pre-processing), or
- 2.
in polynomial time in the special case $i=j$ (under mild assumptions about the distribution of the averaged projection vectors $\{\widehat{\psi}(q^{(i,i)})\}_{q=1}^{n}$).

This is possible because of the extra structure in the quadratic matching problem.
A general efficient solution remains (for now) out of reach, but we hope this progress will spur further research in this direction.
Note especially that independence of the time complexity from the dimensionality of $\{\widehat{\psi}(q^{(i,i)})\}_{q=1}^{n}$ means that the difficulty of the optimisation does not depend on the number of nodes $N$: a very
convenient property for large graphs.

*Proof.*
First, for clarity of exposition, we will introduce more convenient notation and important definitions.
For a given $\epsilon>0$ and a vector $\mathbf{v}\in\mathbb{R}^{d}$, define the set $N_{\epsilon}(\mathbf{v})=\{\mathbf{w}\in\mathbb{R}^{d}:|\alpha_{\mathbf{v},
\mathbf{w}}-\frac{\pi}{2}|\leq\epsilon\}$, where $\alpha_{\mathbf{v},\mathbf{w}}$ is an angle between $\mathbf{v}$ and $\mathbf{w}$.
Then define the property of *$\epsilon$-separation* as follows.

###### Definition C.1 .

For a given $\epsilon>0$, a set of vectors $\mathcal{V}$ and a vector $\mathbf{v}\in\mathcal{V}$, we say that $\mathbf{v}$ is *$\epsilon$-separated* in $\mathcal{V}$ if

$$ $N_{\epsilon}(\mathbf{v})\cap\bigcup_{\mathbf{x}\in\mathcal{V}\backslash\{ \mathbf{v}\}}N_{\epsilon}(\mathbf{x})=\emptyset.$ (57) $$

Simplifying notation, we can rewrite Eq. 55 as the following optimisation problem:

$$ $\sigma^{*}=\textrm{argmin}_{\sigma}\sum_{q_{1},q_{2}\in\{1,...,n\}}[(\mathbf{v }_{q_{1}}^{(i)}+\mathbf{v}_{\sigma(q_{1})}^{(i)})^{\top}(\mathbf{v}_{q_{2}}^{( j)}+\mathbf{v}_{\sigma(q_{2})}^{(j)})]^{2},$ (58) $$

where $\sigma$ is the permutation of the sequence $(1,2,...,n)$
and $\mathbf{v}_{q}^{(i,j)}\coloneqq\widehat{\psi}(q^{(i,j)})$.
With every permutation $\sigma$ of the sequence $(1,...,n)$, we will associate the perfect matching $M(\sigma)$ in its corresponding bipartite graph with monochromatic classes of size $n$ each.
Denote

$$ $\mathbf{f}^{(i)}_{M(\sigma)}\coloneqq\sum_{e\in M(\sigma)}\mathbf{u}_{e}^{(i)},$ (59) $$

where $\mathbf{u}_{e}$ is given by

$$ $\mathbf{u}_{(k,\sigma(k))}^{(i)}=\textrm{vec}\left[(\mathbf{v}_{k}^{(i)}+ \mathbf{v}^{(i)}_{\sigma(k)})\otimes(\mathbf{v}_{k}^{(i)}+\mathbf{v}_{\sigma(k )}^{(i)})\right].$ (60) $$

Here, $\otimes$ represents the outer product and vec is the ‘vectorising’ operation that flattens its input to a vector.
Note that that $\mathbf{u}_{(k,\sigma(k))}\in\mathbb{R}^{N^{2}}$.

It is straightforward to see that, in this notation, the optimisation problem becomes

$$ $\sigma^{*}=\textrm{argmin}_{\sigma}{\mathbf{f}^{(i)}_{M(\sigma)}}^{\top} \mathbf{f}^{(j)}_{M(\sigma)}.$ (61) $$

We will now discuss progress towards solving this efficiently.

Making the problem independent of $N$.
Let us first show how we can make the optimisation problem independent of the number of nodes $N$ with high probability, at the cost of a one-time pre-processing.
Our basic approach is to conduct *dimensionality reduction* of the vectors $\mathbf{u}^{(i,j)}_{e}$ via the Johnson-Lindenstrauss Transform (JLT) [Dasgupta and Gupta, ].
Specifically, we compute

$$ $\widehat{\mathbf{u}}^{(i,j)}_{(k,\sigma(k))}=\frac{1}{\sqrt{r}}\mathbf{G} \mathbf{u}^{(i,j)}_{(k,\sigma(k))},$ (62) $$

where $r\in\mathbb{N}$ denotes the number of random projections and
$\mathbf{G}\in\mathbb{R}^{r\times N^{2}}$ is a Gaussian matrix with entries taken independently at random from $\mathcal{N}(0,1)$.
From the well-known properties of JLT [Dasgupta et al., ], we have that

$$ $\mathbb{E}[(\widehat{\mathbf{u}}^{(i)}_{(k,\sigma(k))})^{\top}\widehat{\mathbf {u}}^{(j)}_{(k,\sigma(k))}]=(\mathbf{u}^{(i)}_{(k,\sigma(k))})^{\top}\mathbf{u }^{(j)}_{(k,\sigma(k))}$ (63) $$

since it is unbiased.
We also have that

$$ $1-\epsilon\leq\frac{\|\widehat{\mathbf{u}}^{(i)}_{(k,\sigma(k))}\pm\widehat{ \mathbf{u}}^{(j)}_{(k,\sigma(k))}\|^{2}_{2}}{|\mathbf{u}^{(i)}_{(k,\sigma(k))} \pm\mathbf{u}^{(j)}_{(k,\sigma(k))}|^{2}_{2}}\leq 1+\epsilon,$ (64) $$

with $r=\mathcal{O}(\frac{\log(n)}{\epsilon^{2}})$ random projections, with probability $p=1-\textrm{neg}(n)$.
Here, neg denotes the negligible function of $n$.

Eq. 64 describes the ability of the JLT to (approximately) preserve $L_{2}$-norms.
The analogous property for preservation of dot-products also holds, which of interest to us given Eq. 61.
It follows because for any $\mathbf{x},\mathbf{y}\in\mathbb{R}^{N^{2}}$,

$$ $\mathbf{x}^{\top}\mathbf{y}=\frac{\|\mathbf{x}+\mathbf{y}\|^{2}-\|\mathbf{x}- \mathbf{y}\|^{2}}{4}.$ (65) $$

The properties above lead us directly to an algorithm to decouple the time complexity of the optimisation from $N$.
One simply replaces every $\mathbf{u}^{(i,j)}_{(k,\sigma(k))}$ with its corresponding dimensionality-reduced $\widehat{\mathbf{u}}^{(i,j)}_{(k,\sigma(k))}$
for $r=O(\frac{\log(n)}{\epsilon^{2}})$ and $\epsilon>0$ small enough.
Then, by the concentration result listed above and the union bound over all pairs $(i,j)$, we conclude that the optimal
$\sigma$ (potentially not unique) will remain optimal after applying the JLT (since all dot-products are approximately preserved).
This completes the analysis.
We stress that $r$ does *not* depend directly on the dimensionality of $\mathbf{u}^{(i,j)}_{(k,\sigma(k))}$, and therefore on $N$.

Polynomial in $n$ algorithm for $i=j$.
We can make further progress in the special case
that $i=j$, i.e. if we are minimising the variance of diagonal kernel entries $k(v_{i},v_{i})$.
Denote by $T(n)$ time complexity of an efficient polynomial-time algorithm for solving the minimum weights matching problem on bipartite graphs with monochromatic parts of size $n$ each.
The following is true.

###### Lemma C.2 .

If Eq. 58 has a unique solution $\sigma^{*}$ and furthermore $\mathbf{f}_{M(\sigma^{*})}$ is $\epsilon$-separated in $\mathcal{F}=\{\mathbf{f}_{M(\sigma)}:\sigma\in S_{n}\}$ for some $0<\epsilon<\frac{\pi}{2}$ (where $S_{n}$ denotes the set of all permutations of the sequence $(1,...,n)$), then, for any $k>0$, $\sigma^{*}$ can be found in time $k\cdot T(n)$ with probability $1-(1-p_{\epsilon})^{k}$.
Here, $p_{\epsilon}$ is given by the following formula:

$$ $p_{\epsilon}=\mathbb{P}\left[\left|\arccos\left(\frac{g_{1}}{\sqrt{g_{1}^{2}+. ..+g_{N^{2}}^{2}}}\right)-\frac{\pi}{2}\right|\leq\epsilon\right],$ (66) $$

with $\mathbf{g}\in\mathbb{R}^{N^{2}}$ a random isotropic vector.

*Proof of Lemma C.2*.
If $i=j$, the expression in Eq. 61 is simply $\|\mathbf{f}_{M(\sigma)}\|_{2}^{2}$.
Thus, the goal is to find a permutation $\sigma$ that minimises the $L_{2}$-norm of the vector $\mathbf{f}_{M(\sigma)}$.
We propose the following algorithm.
At every iteration, replace each vector $\mathbf{u}_{e}$ with the projection $u_{e}=\mathbf{u}_{e}^{\top}\mathbf{g}$ for $\mathbf{g}\in\mathcal{N}(0,\mathbf{I}_{N^{2}})$. Note that, from $\epsilon$-separability, we have that if $\mathbf{g}\in N_{\epsilon}(\mathbf{v})$, then the permutation $\sigma$ corresponding to the minimum weight matching in the bipartite graph with weights given by $u_{e}$ is also equal to $\sigma^{*}$.
This captures the intuition that the shortest vector will have the smallest projection on a random vector that is almost perpendicular to it, provided none of the other vectors are *also* almost perpendicular to this projection direction.
From the fact that $\mathbf{g}$ is taken from the isotropic distribution, we know that the probability that $\mathbf{g}\in N_{\epsilon}(\mathbf{v})$ is exactly $p_{\epsilon}$, where $p_{\epsilon}$ is a geometrical constant that depends only on $N$.
Our algorithm solves the projected minimum weight matching problem (in time $T(n)$) at every iteration and stores the solution $\sigma$.
Vectors $\mathbf{g}$ at different iterations are chosen independently.
After $k$ iterations, the algorithm computes the original objective for every previously-obtained solution and selects the one with the smallest value.
From the above, this returns the optimal permutation $\sigma^{*}$ with probability $1-(1-p_{\epsilon})^{k}$, which completes the proof. ∎

Having considered both claims, the proof of Thm. 4.1 is complete. ∎

###### Definition C.1 .

###### Lemma C.2 .

## Appendix D Asymptotic optimality of permutation densities for variance reduction

In Sec. 4.1 of the main text, we alluded to the choice of permutation densities $p_{\sigma}(x,y)\coloneqq n1_{\sigma(\lceil nx\rceil)=\lceil ny\rceil}$ depicted in Fig. 2 being not only convenient (in terms of ease of sampling and optimising), but also well-motivated by results from OT theory.
Here, we make this statement concrete, showing that in the limit of infinite permutation order $n$ this class contains the solution to a broad range of OT problems.

###### Theorem D.1 (Asymptotic optimality of permutation densities) .

Consider the class of joint measures $\Lambda^{(\sigma)}\coloneqq\{\mu^{(\sigma)}\}$ specified by permutations $\sigma$ of order $n$, given by permutation densities $p_{\sigma}$ pushed forward coordinate-wise using the (left-continuous) inverse CDF $F_{\eta}^{-1}$.
Suppose also that the marginal measure $\eta$ is absolutely continuous with respect to Lebesgue measure.
Consider also a *continuous* function $f:\mathbb{R}\to\mathbb{R}$ whose expectation is to be estimated using $m=2$ coupled samples.
In the limit $n\to\infty$, the class $\Lambda^{(\sigma)}$ contains measures that converge weakly to the optimal transport plan – that is, the sample coupling which provides the smallest possible estimator variance.

*Proof*. Consider the set $\{u_{i}\coloneqq\frac{1}{n}\left(i-\frac{1}{2}\right)\}_{i=1}^{n}$ for some $n\in\mathbb{N}$.
Consider also a random variable $X$ with measure $\eta\in\mathcal{P}(\mathbb{R})$ and corresponding distribution function $F$.
Define its left-continuous inverse $F^{-1}(y)\coloneqq\inf\{x\in\mathbb{R}:F(x)\geq y\}$.
Given some permutation order $n\in\mathbb{N}$, consider the set

$$ $\left\{x_{i}\coloneqq F^{-1}\left(u_{i}\right)=F^{-1}\left(\frac{i-\frac{1}{2} }{n}\right)\right\}_{i=1}^{n}$ (67) $$

which we will refer to as the *quantile midpoints*. Then consider the discrete measure

$$ $\eta_{n}\coloneqq\frac{1}{n}\sum_{i=1}^{n}\delta_{x_{i}}$ (68) $$

where $\delta_{x_{i}}(A)=1$ if $x_{i}\in A$ and 0 0 otherwise (with $A\subset\mathbb{R}$).
Denote by $F_{n}$ the distribution associated with $\eta_{n}$.
Clearly, $\eta_{n}\to\eta$ weakly as $n\to\infty$ since $F_{n}(x)\to F(x)$ for any continuity point $x$ of $F$ [Billingsley, ].
In this sense, the measure $\eta_{n}$ is a discrete approximation to $\eta$.

Now consider the variance reduction OT problem for estimating $\mathbb{E}_{\omega\sim\eta}[f(\omega)]$ with $m=2$ samples,

$$ $\mu^{*}=\textrm{arg min}_{\mu\in\Lambda_{2}(\eta)}\left[\mathbb{E}_{(\omega_{1 },\omega_{2})\sim\mu}f(\omega_{1})f(\omega_{2})\right].$ (69) $$

Solving this analytically for arbitrary $f$ is not in general tractable.
On the other hand, if we instead take the discretised marginals $\eta_{n}$, we must solve

$$ $\mu_{n}^{*}=\textrm{arg min}_{\mu\in\Lambda_{2}(\eta_{n})}\left[\mathbb{E}_{( \omega_{1},\omega_{2})\sim\mu}f(\omega_{1})f(\omega_{2})\right].$ (70) $$

This special case of discrete marginal measures where all points have equal mass can be solved *exactly*.
Our discussion is based on the proof outlined in the introduction to Villani [], to which the interested reader is directed for full details.

OT with equal-mass discrete marginal measures.
Any measure in $\Lambda_{2}(\eta_{n})$ can be represented as a bistochastic $n\times n$ matrix $\mathbf{B}=[B_{ij}]_{i,j=1}^{N}\in\mathcal{B}_{n}$, meaning that all $B_{ij}$ are nonnegative and satisfy

$$ $\sum_{i}B_{ij}=1\;\forall\;j;\quad\sum_{j}B_{ij}=1\;\forall\;i.$ (71) $$

Therefore, the Kantorovich problem reduces to

$$ $\inf_{B\in\mathcal{B}_{n}}\left[\frac{1}{n}\sum_{ij}B_{ij}C_{ij}\right]$ (72) $$

where we have encoded the transport costs in the cost matrix $\mathbf{C}=[f(x_{i})f(x_{j})]_{i,j=1}^{n}$.
$B_{ij}$ is interpreted as the mass of the singleton $(x_{i},x_{j})$.
This is a linear minimisation problem on a convex bounded set.
It is well known that a solution always exists and corresponds to the convex hull of optimal permutation matrices.
In more detail: by Choquet’s theorem the problem admits solutions which are at the extremal points of $\mathcal{B}_{n}$ (points that cannot be written as a nontrivial linear combination of other points in $\mathcal{B}_{n}$).
By Birkhoff’s theorem these are exactly the permutation matrices $\mathbf{B}^{(\sigma)}\coloneqq[\delta_{\sigma(i)j}]_{i,j=1}^{n}$ with $\sigma\in S_{n}$.
So the optimal discrete coupling is the convex hull of the discrete joint measures corresponding to the set of optimal permutations.
See the work of Villani [] for further details.
Choosing just one of these optimal permutations $\sigma$ (since any convex combination of the corresponding measures will give the same amount of variance reduction), we have that

$$ $\mu_{n}^{*}=\mu_{n}^{(\sigma)}\coloneqq\sum_{i=1}^{n}\frac{1}{n}\delta_{(x_{i} ,x_{\sigma(i)})}.$ (73) $$

Stability of OT plans. Another important result in optimal transport is the *stability of transference plans*: namely, that if $\eta_{n}\to\eta$ weakly then the optimal coupling $\pi_{n}\in\Lambda_{2}(\eta_{n})$ also converges to the optimal $\pi\in\Lambda_{2}(\eta)$ weakly provided the cost function $c(X_{i},X_{j})$ (in our case $f(X_{i})f(X_{j})$) is continuous and

$$ $\limsup_{n\to\infty}\int c(x_{1},x_{2})\textrm{d}\pi_{n}(x,y)<\infty.$ (74) $$

This important observation underpins the effectiveness of numerical approximations to optimal transport. We direct the reader to Thm. 1.7.2 by Panaretos and Zemel [] for a proof sketch and Thm. 5.20 by Villani et al. [] for more detailed exposition. It follows that $\mu_{n}^{*}$ converges weakly to the true optimal coupling $\mu^{*}$.

The OT plan is (asymptotically) in our search class.
Lastly, we have that our search class of couplings $\Lambda^{(\sigma)}$ (corresponding to permutation densities pushed forward by $F_{\eta}^{-1}$) contains measures that converge in distribution to $\mu_{n}^{(\sigma)}$ when $n\to\infty$.
In this mathematical sense, the asymptotic limit the class of couplings amongst which we optimise includes measures that give the greatest possible variance reduction: roughly speaking, our method is ‘asymptotically optimal’.
This is intuitive because as the order of the permutation grows each ‘tile’ of nonzero density narrows and can be increasingly well-approximated by a delta function.

To see this, consider the measure on $[0,1]^{2}$ described by the permutation density $p_{\sigma}(x,y)=n1_{\sigma(\lceil nx\rceil)=\lceil ny\rceil}$.
Consider also the discrete measure on $[0,1]^{2}$ given by
$\frac{1}{n}\sum_{i=1}^{n}\delta_{u_{i}}$ with the set $\{u_{i}\}_{i=1}^{n}$ the quantile midpoints defined previously.
These measures converge in distribution (their corresponding joint CDFs can differ by at most $\frac{1}{n}$ at any point, which goes to 0 0 as $n\to\infty$).
They will also converge in distribution when pushed forward coordinate-wise by $F_{\eta}^{-1}$ by continuity of $\eta$.
It follows that at asymptotic $n$ $\mu_{n}^{(\sigma)}$ and $\mu^{(\sigma)}$ converge in distribution, which completes the proof. ∎

We remark that not all the assumptions in Thm. D.1 hold in the GRF setting.
In particular, neither the marginal measure $\eta$ nor the function $f$ is continuous.
The intention of including Thm. D.1 is to build intuition for the reader and provide some motivation for our choice of permutation densities.
We defer a rigorous investigation into relaxing these restrictive assumptions – a tough measure-theoretic problem – to future work.

###### Theorem D.1 (Asymptotic optimality of permutation densities) .

## Appendix E Estimating PageRank

In this appendix, we demonstrate a further use of the variance reduction OT techniques developed for graph random walks in Sec. 4: estimating the *PageRank* vector [Page et al., ].
This popular measure of the relative importance of the nodes $\mathcal{N}$ of a graph is expensive to compute exactly so is often estimated by sampling random walks [Fogaras et al., ].
We will show that OT can be used to find a coupling between the walks to reduce the estimator variance, demonstrating the versatility of our approach beyond improving RFs.

### E.1 Setup: estimating PageRank with random walks

The PageRank vector is the stationary distribution of Markov chain whose state space is the set of all graph nodes $\mathcal{N}$, with a transition matrix

$$ $\widetilde{\mathbf{P}}\coloneqq(1-p_{\textrm{halt}})\mathbf{P}+\frac{p_{ \textrm{halt}}}{N}\mathbf{E}.$ (75) $$

Here, $p_{\textrm{halt}}\in(0,1)$ is a scalar, $N$ is the number of nodes and $\mathbf{E}=[1]_{i,j\in\mathcal{N}}$ is a matrix whose entries are all ones. $\mathbf{P}$ is the transiton matrix of a simple random walk,

$$ $P_{ij}=\begin{cases}\frac{1}{d_{i}}&$if $(i,j)\in\mathcal{E}\\ 0&$otherwise$\end{cases}$ (76) $$

with $d_{i}$ the degree of the $i$th node. Since $\smash{\widetilde{\mathbf{P}}}$ is stochastic, aperiodic and irreducible, we define the unique PageRank vector $\bm{\rho}\in\mathbb{R}^{N}$:

$$ $\bm{\rho}^{\top}\widetilde{\mathbf{P}}=\bm{\rho}^{\top},\hskip 8.53581pt\bm{ \rho}^{\top}\bm{1}=1,$ (77) $$

where we normalised the sum of vector entries to $1$. Computing $\bm{\rho}$ analytically is expensive. Rearranging and Taylor expanding $(1-(1-p_{\textrm{halt}})\mathbf{P})^{-1}$, it is straightforward to see that

$$ $\bm{\rho}_{i}=\frac{p_{\textrm{halt}}}{N}\sum_{j\in\mathcal{N}}\sum_{k=0}^{ \infty}(1-p_{\textrm{halt}})^{k}\mathbf{P}^{k}_{ji}.$ (78) $$

This is simply a sum over all walks from each of the graph nodes $v_{j}$ to node $v_{i}$, weighted by their respective probabilities – that is, the expected number of random walkers ending at node $v_{i}$ if they terminate with probability $p_{\textrm{halt}}$ at every timestep – which invites the estimator proposed by Fogaras et al. [],

$$ $\widehat{\bm{\rho}}_{i}=\frac{1}{Nm}\sum_{v_{j}\in\mathcal{N}}\sum_{n=1}^{m} \mathbb{I}[n\textrm{th walk from node }v_{j}\textrm{ terminates at node}\hskip 2 .84526ptv_{i}].$ (79) $$

An interesting and practically important question is whether the variance of the estimator $\widehat{\bm{\rho}}$ can be improved by coupling random walks.
As with GRFs, this can be achieved using antithetic termination [Reid et al., ] (see Sec. C.2).
However, we will see that our OT length coupling approach does equally well or better.

### E.2 OT formulation of PageRank variance reduction

Given $m=2$ walks, the variance of the PageRank estimator $\textrm{Var}(\widehat{\bm{\rho}}_{i})$ depends on the quantity

$$ $\begin{multlined}\mathbb{E}(\widehat{\bm{\rho}}_{i}^{2})=\mathbb{E}\bigg{[} \frac{1}{N^{2}m^{2}}\sum_{v_{j},v_{k}\in\mathcal{N}}\sum_{n_{1},n_{2}=1}^{2} \mathbb{I}[n_{1}\textrm{th walk from node }v_{j}\textrm{ terminates at node} \hskip 2.84526ptv_{i}]\bigg{.}\\ \bigg{.}\cdot\mathbb{I}[n_{2}\textrm{th walk from node }v_{k}\textrm{ terminates at node}\hskip 2.84526ptv_{i}]\bigg{]}.\end{multlined}\mathbb{E}( \widehat{\bm{\rho}}_{i}^{2})=\mathbb{E}\bigg{[}\frac{1}{N^{2}m^{2}}\sum_{v_{j} ,v_{k}\in\mathcal{N}}\sum_{n_{1},n_{2}=1}^{2}\mathbb{I}[n_{1}\textrm{th walk from node }v_{j}\textrm{ terminates at node}\hskip 2.84526ptv_{i}]\bigg{.}\\ \bigg{.}\cdot\mathbb{I}[n_{2}\textrm{th walk from node }v_{k}\textrm{ terminates at node}\hskip 2.84526ptv_{i}]\bigg{]}.$ (80) $$

After a few lines of algebra, the OT problem to minimise this quantity is

$$ $\begin{multlined}\textrm{minimise }\mathbb{E}_{(l_{1},l_{2})\sim\mu}\sum_{v_{j }\in\mathcal{N}}\mathbb{P}[\textrm{RW of length $l_{1}$ from node $v_{j}$ terminates at node $v_{i}$}]\\ \cdot\mathbb{P}[\textrm{RW of length $l_{2}$ from node $v_{j}$ terminates at node $v_{i}$}]\textrm{ for }{\mu\in\Lambda_{2}(\eta_{\textrm{G}})}\end{ multlined}\textrm{minimise }\mathbb{E}_{(l_{1},l_{2})\sim\mu}\sum_{v_{j}\in \mathcal{N}}\mathbb{P}[\textrm{RW of length $l_{1}$ from node $v_{j}$ terminates at node $v_{i}$}]\\ \cdot\mathbb{P}[\textrm{RW of length $l_{2}$ from node $v_{j}$ terminates at node $v_{i}$}]\textrm{ for }{\mu\in\Lambda_{2}(\eta_{\textrm{G}})}$ (81) $$

without making any approximations.
$\mathbb{P}(\cdot)$ denotes the probability of the event in brackets.
As with GRFs, for tractability we restrict our search space of joint distributions to permutation densities pushed forward coordinate-wise with the geometric distribution inverse CDF, and this becomes

$$ $\begin{multlined}\sigma^{*}=\arg\min_{\sigma}\sum_{v_{j}\in\mathcal{N}}\sum_{q \in[\![n]\!]}\mathbb{P}[\textrm{RW with length in $q$th quadrant from node $v_ {j}$ terminates at node $v_{i}$}]\\ \cdot\mathbb{P}[\textrm{RW with length in $\sigma(q)$th quadrant from node $v_ {j}$ terminates at node $v_{i}$}].\end{multlined}\sigma^{*}=\arg\min_{\sigma} \sum_{v_{j}\in\mathcal{N}}\sum_{q\in[\![n]\!]}\mathbb{P}[\textrm{RW with length in $q$th quadrant from node $v_{j}$ terminates at node $v_{i}$}]\\ \cdot\mathbb{P}[\textrm{RW with length in $\sigma(q)$th quadrant from node $v_ {j}$ terminates at node $v_{i}$}].$ (82) $$

The probabilities can be efficiently estimated by simulating random walks on the graph and recording where they terminate.
This leaves us with a minimum-weights bipartite matching problem which can as usual be solved efficiently with the Hungarian algorithm [Kuhn, ].

Note that we made fewer simplifications than for GRFs (Sec. C.1).
They only requirements for tractability are restricting the class of considered joints to $\sigma$-couplings and computing MC estimates of the terms in the OT cost matrix.
We do not need to e.g. move $\mathbb{E}_{\textrm{dirs}}$ inside a square or approximate a quadratic matching problem by a linear-weights counterpart.

### E.3 Empirical results for PageRank

Figure: Figure 9: PageRank estimator errors $\|\bm{\rho}-\widehat{\bm{\rho}}\|_{2}$ for a range of real-world graphs when $2$ random walkers are taken to be i) i.i.d., ii) antithetic or iii) $\sigma$-coupled. Lower is better. The error is always at least as small and often substantially better when the walkers are coupled using our OT coupling. One standard error is shaded.
Refer to caption: x7.png

We now provide experiments using our OT coupling to improve the convergence of PageRank estimates.
Fig. 9 shows the results.
For termination probabilities $p_{\textrm{halt}}\in\{0.1,0.2,0.3,0.4,0.5,0.6,0.7,0.8,0.9\}$, we compute the optimal permutation $\sigma$ by solving the matching problem in Eq. 82 on an Erdős-Rènyi graph with $N=100$ nodes, taking permutation order $n=10$ as a hyperparameter.
We then test these couplings on a range of real-world graphs [Ivashkin, ] of different sizes, plotting the estimator error $\|\bm{\rho}-\widehat{\bm{\rho}}\|_{2}$ (where $\bm{\rho}$ is the true PageRank vector and $\widehat{\bm{\rho}}$ is its MC estimate) against the termination probability $p_{\textrm{halt}}$.
We include walkers that are i.i.d., antithetic [Reid et al., ] and $\sigma$-coupled.
At every value of $p_{\textrm{halt}}$, our method performs *at least as well* as the baselines, and in some cases (e.g. on the very large cora graph) it gives much greater variance reduction.

Note that we only consider connected subgraphs and take $1000$ repeats for standard errors.
Reading the plots from left to right for the top row and then the bottom, the number of nodes are: $\{34,62,105,115,398,986,1133,1272,2120,2708\}$.
The shape of the curve and size of the gain from coupling depends on the particular graph being considered.
For some graphs it is not possible to obtain a big reduction in PageRank estimator variance by coupling walkers (e.g. football and newsgroup), so neither antithetic termination nor our $\sigma$-coupling can provide a big gain.
A detailed investigation of how these observations relate to the mathematical graph structure is deferred to future work.

## Appendix F GRFs for GPs: additional details and experimental results

In this appendix, we supplement Sec. 4.2 by providing further experimental results for GRFs, including Gram matrix estimation on a wider range of graphs and a scalable graph-based GP on a different real-world dataset.
We also provide technical background and tedious details considered too long for inclusion in the main text.

### F.1 More results for Gram matrix estimation

First, we run the experiment in the first part of Sec. 4.2 for more real-world graphs.
We approximate the $2$-regularised Laplacian kernel $\smash{(\mathbf{I}-\sigma^{2}\widetilde{\mathbf{L}})^{-2}}$ with a regularisation parameter $\sigma=1$ and $m=2$ walkers.
Fig. 10 shows the results, with the kernel approximation error normalised by the i.i.d. result (unlike in Fig. 3) for clarity of comparison.
The $\sigma$-coupling consistently reduces the estimator variance, providing equally good or better results compared to the hand-crafted antithetic termination algorithm.
The size of the gain depends on the graph structure.

Figure: Figure 10: Relative Frobenius norm error $\|\mathbf{K}-\widehat{\mathbf{K}}\|_{\textrm{F}}/\|\mathbf{K}\|_{\textrm{F}}$ between true and approximated Gram matrices for different termination probabilities when walkers are i) i.i.d., ii) antithetic [Reid et al., ] or iii) $\sigma$-coupled using our OT approach. All results are normalised by the i.i.d. variance for easy comparison. Lower is better. The $\sigma$-coupling is consistently as good or better across values of $p_{\textrm{halt}}$ and graph topologies and sizes. One standard error is shaded.
Refer to caption: x8.png

### F.2 Scalable GPs on graphs

A very short introduction to graph-based GPs.
In certain applications of Gaussian processes, kernels based on Euclidean distance are unsuitable – e.g. when predicting traffic congestion since locations that are spatially close may not be connected by roads.
Kernels defined on the nodes of a graph $\mathcal{G}$ may be more appropriate, giving structured covariance matrices that encode the dependence between pairs of vertices.
We can then perform inference and make predictions using *GPs on graphs* [Borovitskiy et al., , Zhi et al., ].

Like their Euclidean counterparts, exact GPs on graphs suffer $\mathcal{O}(N^{3})$ time complexity, making them impractical for very large structures.
Techniques to improve their scalability include graph Fourier feature approximations, which approximate the kernel matrix using a truncated eigenvalue expansion computed using the Lanczos algorithm.
The price of increased speed is that our kernel estimate is biased. Another approach is to restrict to kernels whose inverses are sparse, e.g. a Matérn kernel with a small integer $\nu$-parameter, at the expense of lost flexibility [Borovitskiy et al., ].
In this paper, we have proposed to instead use GRFs, which give an unbiased estimate to the kernel matrix in subquadratic time. GRFs are sparse: intuitively, $\phi_{\textrm{GRF}}(v_{i})$ is constructed by simulating $m$ random walks out of $v_{i}$ and adding weighted contributions only at the coordinates corresponding to visited nodes.
For walks with geometrically-distributed lengths this is typically a small subset of $\mathcal{N}$. Therefore, the matrix $\smash{\widehat{\mathbf{K}}\coloneqq[\phi(v_{i})^{\top}\phi(v_{j})]_{i,j=1}^{N}}$ is a sparse, unbiased estimate to *any* kernel, not just specific families.(^4^44Strictly: any graph node kernel that is expressible as a function of a weighted adjacency matrix.) We can use established numerical routines for sparse linear algebra to speed up training, providing computational savings on big graphs.
This is analogous to the sparse spectrum GPs discussed in App. B [Lázaro-Gredilla et al., ]; the interested reader might benefit from reviewing this section first.

Our contribution to scalable graph-based GPs.
Using GRFs for scalable approximate inference on graphs is itself a novel contribution that deserves detailed exploration in future work.
Since our central goal is to use perspectives from OT to improve the convergence of GRFs, we confine our focus to the question: *can coupled random walks give more accurate estimates of graph GP posteriors?*
We defer a detailed comparison of our method to other scalable graph-based GP techniques to a future self-contained paper.

Training GRF-GPs: technical details.
Here, provide details to supplement Sec. 4.2.
We use mesh graphs made available by Dawson-Haggerty [], taking the *faces* as our graph nodes.
Reading left to right, the number of nodes are: $894,999,1280,1572,3838,8700$.
For probabilistic mesh interpolation, we use the diffusion kernel $\mathbf{K}=\kappa\exp(-\gamma^{2}\widetilde{\mathbf{L}})$.
Here, $\widetilde{\mathbf{L}}$ is the normalised graph Laplacian (Eq. 48) and $\kappa,\gamma$ are learnable parameters corresponding to the signal variance and kernel lengthscale respectively.
The observation noise $\sigma$ is also learnable. For every graph, we construct a kernel estimate by sampling $m=8$ i.i.d. random walks out of every node.
We then train the corresponding vanilla GP by optimising the log marginal likelihood on observed nodes, using the Adam optimiser [Kingma and Ba, ] for $1000$ epochs (after which both the objective value and the hyperparameters are empirically found to have converged).
We freeze $\kappa,\gamma,\sigma$ and compute the corresponding exact kernel.
We approximate this exact kernel using $\{16,32,64\}$ random walks that are either i.i.d., antithetic or $\sigma$-coupled. We perform approximate Bayesian inference in each case, computing the accuracy of the kernel estimate, the average test RMSE and the KL divergence to the true posterior distribution (unnormalised by the number of test points).
We use permutations of order $n=50$ and a termination probability $p_{\textrm{halt}}=0.4$.
As reported in the main body, *coupling the lengths of graph random walkers permits better approximate Bayesian inference with GPs on graphs*.

### F.3 Probabilistic traffic interpolation

In this section, we provide further details about the final experiment of Sec. 4.2, which uses a scalable graph-based GP to predict traffic flow speeds on the highways of San Jose, California.

Dataset. Readers are directed to the original paper [Borovitskiy et al., ] for exhaustive details about the dataset curation and graph computation.
Here, we simply remark that this graph is sufficiently large that exact GP methods become slow ($N=1016$), and observations are only made at a small subset of the nodes ($325$) so good uncertainty quantification is essential.

Results for GRFs.
We randomly divide the nodes into training and test datasets of sizes $N_{\textrm{train}}=250$ and $N_{\textrm{test}}=75$ respectively.
As described in App. F.2, we train a scalable GRF-GP with $m=8$ i.i.d. walkers set up to approximate the diffusion kernel, optimising the training data negative log marginal likelihood.
Next, we freeze the kernel and noise parameters, and compare the performance of GRFs with $m\in\{4,8,16\}$ i.i.d., antithetic [Reid et al., ] and $\sigma$-coupled walkers.
As before, we consider the quality of Gram matrix approximation, the accuracy of predictions (test RMSE) and the quality of the predictive uncertainties (KL divergence to true posterior).
Fig. 11 shows the results.
$\sigma$-GRFs consistently do best on all metrics, even with a severely restricted sampling budget.

Figure: Figure 11: Results for probabilistic traffic interpolation experiment. Plots show the kernel approximation accuracy, test RMSE and KL divergence to the true posterior for graph kernels estimated using i.i.d., antithetic and $\sigma$-coupled GRFs. Lower is better. $\sigma$-coupling gives the best results for approximate inference. One standard error over random draws of walkers is shaded.
Refer to caption: x9.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: My main concern is about the justification and also computational aspects of the proposed framework. As the authors explain, the main limitation of this work is that variance is reduced at a specific pair of points, which also leads to some limitations in the performance of downstream tasks. The condition of fixed marginals in the OT framework actually guarantees an unbiased kernel as a process, i.e. for every pair of points. Is not it more reasonable to minimize the variation of the entire process rather than a specific point?

The pointwise view also leads to a potential computational limitation. If I understood correctly, in practice one needs to generate an independent set of coupled features for each pair of points, which requires to run the variance reduction method once for each pair. This leads to a huge computational burden in a large-scale problem.

More specific questions regarding the presentation:

1- I am not familiar to the notion of negative monotone coupling and for me (7) does not make sense. Can you provide more explanation?
2- If I am correct, theorem 3.2. only works for sufficiently small z, this should be explained in the introduction before this equation to avoid confusion.
3-In line 290, what is the definition of \hat{a}_i?
4-In line 387, does "walker's direction" mean the random walk conditioned on its length?
5-In Theorem 4.1, I do not understand the definition of q^{i,j}. Also q^{i} is not exactly defined.
soundness: 3
strengths: I find the general idea of the paper fascinating, namely formulating the variance reduction task as optimal transport. This approach can also be applied to similar Monte Carlo methods in future studies. It also allows for the application of various techniques in OT for studying and improving kernel estimation methods.
summary: This paper addresses a well-known Monte Carlo method for estimating the value of kernel functions using random feature embeddings. In particular, it proposes a method to reduce the variance of the estimate by coupling individual random features. This task is framed as a general optimal transport problem and is further discussed in two different scenarios: first random Fourier and Laplace features, and second features obtained by a random walk on graphs, corresponding to graph kernels. In the first case, exact results for two features are obtained and some extensions to more features is discussed. Moreover a search algorithm over a parameterized family of couplings is proposed. In the second case, a simplified coupling of the length of the random walks is discussed.

Numerical experiments are presented for each case that show the effectiveness of these algorithms.
weaknesses: Although the paper is well-written and easy to understand, the presentation of the paper can be improved. Much of the technical details that are summarized in the main body cannot be easily understood without referring to the appendix. There are confusing notations and definitions, especially for the case of graph kernels. I will mention some of them in the next part. 

I also have a more specific questions/concerns about the justification of this framework which is explained in the next part.

I am willing to raise my score if my concerns are addressed.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 4
questions: Regarding the choice of metric, what are the benefits of using KL over predictive NLL (or NLPD)? I noticed that on page 22, it is mentioned that the predictive NLL is computed, so perhaps I missed where these results are presented.

Are there applications where the kernel reconstruction is the primary focus, rather than the derived confidence intervals? If so it might be beneficial to emphasise these to the reader.

One further small suggestion - at present the authors describe as "surprising" the result that kernel accuracy improvement often does not lead to downstream test improvements. On the contrary, I think there have been several reports of this phenomenon in the literature. For example, from Liu et al 2004.11154 "We find that ORF and SSF yield the best approximation quality. Despite that most algorithms achieve different approximation errors, there is no significant difference in the test accuracy."
soundness: 3
strengths: The manuscript is well presented, structured, and clearly written.

There is a high degree of novelty in the application of optimal transport in determining the feature couplings. 

The proposed method is tested on a good selection of real world datasets.
summary: The paper presents a novel approach for improving kernel estimation methods via random features. By utilising optimal transport, they tackle the specific cases of random fourier features, random laplace features, and graph random features. They demonstrate this technique on a range of real world datasets. And while some practical advantages are found, it is also noted that reducing the kernel variance does not necessarily improve the quality of the forecasts.
weaknesses: I have some concerns regarding some of the experimental results presented in the manuscript:

Regarding Table 1 - Running only d features for RFF strikes me as a rather low number to choose, especially for these modestly sized datasets. It would be helpful to clarify how the results impacted when using a larger numbers of features. It would be great to see something analogous to the Figure 1 of the Yu et al 'Orthogonal Random Features' paper which depicts 1 < d/D < 10. 

While I can appreciate that normalisation of RMSE's to 1.0 makes for a clean reference point, this also obfuscates the absolute performance from the reader (which is significant for the purposes of reproducibility, and comparisons with other works).

Regarding Table 3 - In addition to, or instead of, providing the KL results, I think the predictive performance (NLL or NLPD) would be preferable. I believe this would provide a better proxy for the model evidence. My concern with the KL divergence is that it does not heavily penalise models which are overconfident.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. What are the computational implications of implementing these optimal transport techniques at scale?
2. Can the theoretical findings on variance reduction be generalized to other kernel methods or models?
soundness: 3
strengths: 1. **Innovative Use of Optimal Transport**: The paper employs OT as a novel framework to address the variance reduction problem in RFs, which is a creative approach that ties together theoretical insights with practical application.
2. **Comprehensive Coverage**: It addresses both Euclidean and discrete input spaces, providing a unified strategy applicable across different domains.
3. **Theoretical Guarantees and Empirical Validation**: This paper offers theoretical guarantees and empirical validations, enhancing the understanding of variance reduction in RF performance.
summary: The paper explores variance reduction techniques for random features (RFs) in kernel methods to improve efficiency in large-scale machine learning tasks. By framing variance reduction as an optimal transport (OT) problem, the authors propose couplings that enhance the convergence of RFs on both Euclidean and discrete input spaces. They introduce novel coupling techniques for various RF classes, including Random Fourier Features (RFFs), Random Laplace Features (RLFs), and Graph Random Features (GRFs). The paper highlights the effectiveness of the proposed methods in reducing variance and improving downstream performance in certain cases, although it also discusses limitations.
weaknesses: 1. The definition and relationship between variance reduction in kernel estimation and optimal transport should be more detailed and introduced for non-specialists to understand directly.
2. The computation and theorem (Theorem 3.2) is only for $m=2$, and the authors apply the copula tool as numerical OT solvers for multi-marginal OT.
3. Despite variance reduction, the downstream impact on model performance can be inconsistent, especially for transformers.
4. More comparisons between this paper and [1] and [2] should be clarified, since some of the idea and technical proof skill is from these paper.


[1]. *General Graph Random Features*. Nips 2023.
[2]. *Simplex Random Features*. ICML 2023.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Here below some comments and observations:

- After Eq 1, it is not clear if *each* $\phi$ has to be intended depending on $m$ random frequencies (also, $\\{\boldsymbol\omega\\}$ should be replaced with $\\{\boldsymbol\omega_i\\}$: it would be sufficient, for example, to write down $\phi(\boldsymbol x;\boldsymbol\omega_{1:m})$ in place of $\phi(\boldsymbol x)$. At line 048 I am not sure if Eq 1 is a "MonteCarlo estimate" of $k$, as in Eq 1 there is no empirical sampling. 
- I might being missing something trivial, but why $\mathcal I(\mu)=\mathbb E_\mu[(\phi^\top\phi)^2]$ and not $\mathcal I(\mu)=\mathbb E_\mu[(\phi^\top\phi)^2]-\mathbb E_\mu[\phi^\top\phi]^2$? I suppose that, given the constraint on the marginals, the second term is irrelevant if $\phi(\boldsymbol x;\boldsymbol\omega_{1:m})^\top\phi(\boldsymbol y;\boldsymbol\omega_{1:m})=\sum_{i=1}^m\psi_i(\boldsymbol x;\boldsymbol\omega_i)^\top\psi_i(\boldsymbol y;\boldsymbol\omega_i)$ for some $\psi_i$, but this might be not true in general. Unless I am wrong, would the more general case lead to a much more complicated treatment? The treatment of RFF and RLF seems to be strongly case-specific, although relevant.
- The discussion of the application to GRFs leads to encouraging results. Do the authors have an intuition on why this is the case? Are such positive performances retained independently from the level of the sparsity of the graph? In general, I was wondering how much the "sparsity" of the underlying space where the kernel lives affects the effectiveness of the method.
soundness: 3
strengths: The paper aims to solve the problem of minimizing the variance of kernel estimators by random features by mapping the problem in an optimal transportation (OT) problem. By using OT concepts, (sub-optimal but interesting) coupling strategies are identified in both the Euclidean case and in a discrete case, namely graph kernels. The paper links therefore two ideas (kernel random features and optimal transport) that are very popular in the machine learning community, with an interesting digression on graph random features and graph kernels. Moreover, I appreciate the rich numerical exploration of the results, that highlighted the benefits and limitations of the method.
summary: The paper considers the problem of finding a suitable joint measure for random features in a kernel evaluation by applying a variance reduction principle. The application of such a principle in general results in a coupling of the features, expressed, for example, by a constraint on their modulus in the case of features in the Euclidean case. The principle is given in terms of an optimal transport formulation and its application is exemplified by a number of experiments on both Euclidean kernels and graph kernels. The authors show that if on the one hand, a smaller kernel variance is obtained as expected, other quantities might benefit from different strategies, encouraging a rethinking of the relevance of kernel variance reduction.
weaknesses: My main concern is that the proposed OT formulation does not obviously lead to a simple, applicable rule, as shown by the authors themselves. Even in the random Fourier features case, under the simplifying pairwise coupling assumption, guarantees can be proven only for close enough arguments of the kernel function to estimate. The examples discussed and the appendices show that the application of the OT ideas might be cumbersome, not straightforward, and require a certain ingenuity to perform a number of simplifications and approximations (whose effectiveness is only shown *a posteriori*): this comes with no surprise as solving an OT problem is generically hard. More importantly, if the quantities of interest are nonlinear functions of the kernel, counterintuitive results might be found, e.g., *maximizing* the variance might be beneficial. It appears then unclear how to decide *a priori* what the best coupling strategy is. Even if, in principle, one might think to generalize the minimization problem in Eq 2 (so that the functional $\mathcal I(\mu)$ depends on the variance of other target quantities), as the authors themselves explain, the possible complication of such generalization, and its feasibility, seem hard to anticipate.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. In Lemma 3.1, the forms of c_{RFF} and c_{RLF} are very complicated and the solution \mu^* seem to depend on x and y. In Theorem 3.2 with m=2, the solution seems to be independent of x and y. Is it the case for m=2? Or the solution will always be irrelevant to x and y no matter how large m is? In addition, it will be better to give more explanation about the statement ``provided z is sufficiently small". For example, how small z is w.r.t. the order? How to check if z is small enough in practice? 

2. In Section 3.3, what is the definition of RMSE? Is it the point-wise RMSE or the Frobenius norm error for the gram matrix approximation?  What is the choice of d and how RMSEs change with d empirically? 

3. Though the paper provides that Theorem that the pairwise norm-coupled RFs is guaranteed to be lower than orthogonal RFs with independent norms, it will be better to provide an explicit error upper bound for the variance, which will be easier to see the dependence of the error on different parameters. 

4. It is interesting to see that the variance reduction alone cannot guarantee an improvement in the downstream tasks. I am wondering is it the case for most datasets considered in Table 1? If so, the application of this RFs may be somehow limited. 

5. In page 8, the paper proposes to ``move Edirs inside the square" to make the computation feasible. It will be better to discuss the effect of such adjustment.  By moving the expectation, the objective function indeed changes and solving the adjusted objective function does not necessarily reduce the variance? 

6. In the numerical experiments w.r.t. RGFs, the proposed method performs the best also in the downstream tasks for different datasets, which shows a different performance from RFFs and RLFs.  Is there any intuition behind this?
soundness: 3
strengths: 1. The connection between variance reduction of RFs and OT is clearly presented.
2. RFs for different settings including RFFs, RLFs and GRFs are deeply studied. 
3. A through simulation study has been conducted and some interesting phenomenons show up in the numerical experiments.
summary: This paper proposes an efficient construction of random features for kernel methods. Motivated by the optimal transport, it finds couplings to reduce the variance of random features defined on both Euclidean and discrete input spaces. Numerical experiments are conducted and theoretical properties are analyzed.
weaknesses: Please see following questions.

Meta-review: This paper provides a novel approach to reducing the variance of Randon Fourier feature maps via a connection to optimal transport. All reviewers felt that the connection was interesting and new, and will likely be of interest to many in the ICLR community. There were some concerns regarding ultimate value to downstream applications, as reducing embedding variance did not always lead to better downstream performance. I think the authors make a convincing argument that highlighting this point can be viewed a feature instead of a bug, as many prior works focus on variance reduction. Overall, a nice that that we would be happy to see at ICLR.

Decision: Accept (Poster)