Variance-Reducing Couplings for Random Features
Authors: Isaac Reid
Stratis Markou
Krzysztof Marcin Choromanski
Richard E. Turner
Adrian Weller
Venue: ICLR 2025 Poster (2025)
Keywords: Monte Carlo
variance reduction
quasi Monte Carlo
transformers
performers
optimal transport
random Fourier features
graphs
Gaussian processes
kernels
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: A unifying framework for variance reduction of random features in Euclidean and discrete space, with applications to geometric GPs and efficient transformers

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

Published as a conference paper at ICLR 2025
VARIANCE-REDUCING COUPLINGS FOR RANDOM FEATURES
Isaac Reid', Stratis Markou', Krzyszzof Choromanski RRichard E. Turner', Adrian Weller University of Cambridge, ?Google DeepMind, "Columbia, 4Alan Turing Institute
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 approx- imating attention) sparse spectrum Gaussian processes (by approximating the covariance function). Efficiency can be further improved by speeding vergence of these estimates: variance reduction problem. We tackle this through the unifying lens of optimal transport, finding couplings to improve RFs defined on both Euclidean and discrete imput spaces. They enjoy theoretical guarantees and sometimes provide strong downstream gains, including for scalable inference on graphs. We reach surprising conclusions about the benefits and limitations of variance reduction paradigm, showing that other properties of the coupling should be optimised for attention estimation in efficient transformers.
the con-
INTRODUCTION
Kernel methods are ubiquitous in machine learning (Canu and Smola, 2006; Kontorovich et al., 2008; Campbell, 2002). Through the kernel trick, they provide mathematically principled and elegant way to perform nonlinear inference using linear learning algorithms. The eponymous kernel function k X' x X + measures the "similarity" between two datapoints. The input domain X' may be continuous, c.g. the set of vectors in Rd, discrete, c.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 K == [k(2,, x, 2 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, prominent example being random features (RFs) (Rahimi and Recht, 2007: 2008; Avron et al., 2017b; Liu et al., 2022). These randomised mappings á Rdd R construct low-dimensional sparse feature vectors that satisfy
k(x, 3) = E (Φ(x; {w.)),)) Toy: (w:)=))-
(1)
The expectation E is taken over an ensemble of random frequencies {wi) drawn from distri- bution 17. The space in which {w) live and manner in which they are combined to construct p(æ; { depends on the particular input space X' and kernel function k being approximated.
This paper will consider several examples. Hereafter, the dependence on (w:)m will be suppressed reduce notational clutter. The set of RFs (o(xx,) ) can be used to construct 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 al., 2010; Johnson, 1984; Choromanski et al., 2020; Rahimi and Recht, 2007; Tripp et al., 2024).
Variance reduction for RFs. Replacing by the mean over random samples of [wr] Eq. can be understood as 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- Carlo (QMC; Dick et al., 2013; Yang et al., 2014a), common random numbers (CRNs; Glasserman and Yao, 1992), antithetic variates (Hammersley and Morton, 1956) and structured Monte Carlo (SMC; Yu et al., 2016). These techniques work by replacing i.i.d. frequencies (w:)) by dependent ensemble, with the sample dependencies designed to improve RF convergence.
Senior lead.

---

Published as a conference paper at ICLR 2025
Limitations of previous techniques. The best choice of dependencies between [,,];" is active research area. Though straightforward to apply, standard QMC techniques are suboptimal. They are based on hard-coded fow-discrepancy sequences' so typically do not incorporate information about the particular kernel k heing approximated. Empirical performance may be poor and theoretical guarantees lacking in the low-sample, high-dimensional regime (Rowland al., 2018; Morokoff and Caflisch, 1995), 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., 2016). 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 graph have only recently been developed (Choromanski, 2023; Tripp et al., 2024), 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 busis 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 efficiently as possible (Villani et al., 2009). 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 previous approaches. OT allows us to improve couplings for RFs in both Fuclidean 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.
For random Fourier features (RFFs) and random Laplace features (RLFs), we exactly solve the OT problem for the norms of norm-coupling, which guarantees lower variance for arbitrary m.
2 orthogonal frequencies (Sec. 3). We introduce pairwise
3. For graph random features (GRFs), 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 (Sees 3 and 4). We show that this sometimes translates to much better performance in downstream tasks, including for 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 k(x,, 22)) == o(x;)" Φ(æ,). Recall that ø(-) is computed using random frequencies {w, )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 w; obeys some marginal distribution 1). Importantly, independence of {w; ) is nor required: any joint distribution with marginals 7) gives an unbiased estimator. We refer the set of such joint distributions as couplings.
The coupling between the frequencies determines the estimator variance. We want solve:
minimise I(p) = chw for fa E An))))
(2)
where we defined the cost function c(wi-)) = (d(x)) TØ()))² and (n) denotes the set of couplings of m random variables with marginal measures 77. This is precisely the Kantorovich formulation of multi-marginal OT problem (see Eq. 4 of the seminal OT text of Villani (2021)). 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 μ² arg min A... (xa) [E m (( with the smallest estimator variance. The relationship between variance reduction and OT was also noted by Rowland et al. (2018) in different context.
2

---

Published as a conference paper at ICLR 2025
(Approximately) solving the OT problem. The formulation of Eq. 2 depends 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 X' Rd (d-dimensional Fuclidean space) and X' = (the set of graph nodes) respectively. This gives new couplings with lower RF variance than previous algorithms.
RANDOM FOURIER FEATURES AND RANDOM LAPLACE FEATURES
RFFss and RLFs, To begin, we consider the task of approximating the popular Gaussian kernel k(z; my) == cxp(-IIaz x;,1² /2) with data { C R This can be achieved using Rahimi and Recht's 's celebrated random Fourier features (RFFs) (Rahimi and Recht, 2007),
V 1
(RFF(x))
m (
[sin( x)) [os[(w])))))
(3)
where denotes concatenation. These provide an unbiased estimate if the frequencies {w; m are marginally Gaussian, W, ~ (0, Ia). RFFs are widely used for scaling kernel methods such as Gaussian processes (GPs; Williams and Rasmussen, 2006) and support vector machines (SVMs; Scholkopf and Smola, 2018). The time complexity of computing the exact posterior of GP is ((N"), where N is the number of datapoints. Using RFFs, one can approximate the posterior with m & N features, reducing this cost ONNm²). Changing basis functions, k( 2.)) can also be approximated using random Laplace features (RLFs) (Yang et al., 2014b),
ØRLF(W)
exp(
æ (o exp(w, x)),
where again w; ~ (O Ia). Unlike RFFs, RLFs guarantee positive kernel estimates. This makes them better suited to approximating attention in efficient transformers (Choromanski et al., 2020), where negative estimates cause training instabilities. Using RLFs to get low-rank decomposition of attention with N d-dimensional tokens, one can reduce the time complexity of transformers from O(N2 Nd) to O(Nmd) with low performance loss.
Orthogonal random features. common variance reduction technique for both RFFs and RLFs is the orthogonality trick (Yu et al., 2016; Rowland al., 2018; Reid al., 2023; Choromanski et al., 2018). Exploiting the isotropy of /O Ia), one can constrain the frequency vectors (w.) to be exactly orthogonal whilst preserving their marginal distributions. This is found reduce the kernel estimator variance and improve performance in downstream tasks. Whilst this technique couples the directions of the random frequencies (w) N= their norms {w:] (with wi = l|w)ll2) 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 (w, 1m jointly randomly rotated so they are marginally isotropic. Our task is to couple their norms {wr to suppress the RFF and RLF kernel estimator variance. The marginal distribution of cach w, must be Xa (a Chi distribution with d degrees of freedom) to ensure that each w; is marginally Gaussian. We can extend recent results by Reid et al. (2023) to compute the OT cost functions.
Lemma 3.1 (OT formulation for RFFs and RLFs). When estimating k(æ, y) with m orthogonal RFFs and RLFs, the OT formulation of the variance reduction problem is:
14" = arg min (EAm((X)
[[E [Eurmre((wim)
where
(5)
2 22 (w) w3)
m
2k (w? w/)²
(RFF(W1sm)
CREF(Wi:))
(6)
22k IT (k: + 2)) ( k=0
22k R + 2))
1.j K=
with = ||
4112 and
x + Ul12. I is the gamma function.
This is a tough multi malti-marginal OT problem. However, remarkably, we can solve it exactly, under mild asymptotic assumptions for RFFs, when m 2. The following result is novel.
3

---

Published as a conference paper at ICLR 2025
Theorem 3.2 (Solution to OT problem when m 2). Denote by Exa()) the cumulative distribution function (CDF) of Xa. Consider m orthogonal frequencies fieunncce with norms (2i1,W2). For RLFs, the OT problem in Eq. 5 is solved by the negative monotone coupling
Fxa(wr) + (xx( =.
For RFFs, Eq. 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 brief sketch. OT plans satisfy a property called 'c-monotonicityy, which specifies how the support of the optimal coupling depends on the cost function. For RLFs, CRLF 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, a is small enough. Specifically, letting UNM denote the negative monotone coupling, for any other coupling yl E (n)) {WNM) there exists some constant ((n)) such that I(NNM) I(n') for all 2 < 8 (Lemma A.6).
can show that Eq. 7 still guarantees lower variance than any other coupling
Given d orthogonal frequencies, can partition the ensemble into 1 orthogonal pairs, with one remaining frequency if d is odd. For every pair, one can impose negative monotone coupling (Eq. 7). We refer such ensembles pairwise norm-coupled (PNC).
Definition 3.3 (Pairwise norm-coupled REs). RFs are pairwise norm-coupled (PNC) if d orthogonal frequencies [wi} are arranged in [ pairs, each of which is negalive montone-coupled so that Fx (W1) + Fxa(ws)
Different pairs are independent.
PNC is no more expensive than 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; is a space- filling coupling that seeks long transport plans that give diverse samples. However, it is popular heuristic technique for variance reduction via common random numbers (CRNs) computational statistics (Glasserman and Yao, 1992). 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 (i1, (22) are unconstrained, variance minimised by conditioning that W1 W2 almost surely (that is, opposite directions and equal norms).
This coupling is known as antithetic sampling (Hammersley and Morton, 1956). Thm. 3.5 shows that, given PNC ensemble {w; )d (#w, Antithetic sampling is also a common (though often heuristically motivated) variance reduction strategy used e.g. when estimating attention in Performers (Choromanski al., 2020). We can reinterpret its effectiveness as an OT coupling.
obtain further variance reduction by augmenting it to
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 m 2, then combining [ independent copies to get the ensemble. Can do better by inducing dependencies between the all the 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 Am (xa), the full set of joint distributions over m random variables with Xa marginals, to a tractable subset amongst which we can efficiently optimise and sample. One such subset is provided by Gaussian copulas (Nelsen, 2006; Haugh, 2016): joint distributions obtained by
4

---

Published as a conference paper at ICLR 2025
Kernel RMSE
Predictive RMSE
Prodictive NLL
Predictive KLD
Prior KLD
evaluation 0.
4.00
2300
0.02
4.33
2.89
0.03
0.02
0.00
10
(/d)
L.dd.
orthogonal
PNC
kernel
linear regression
Figure 1: Downstream performance on a split of the POWER dataset. Kernel estimator RMSE, test predictive RMSE, test negative log likelihoods, KL divergence to the true predictive posterior, and KL divergence to the true prior, estimated with RFFs, We have successfully reduced the variance of the kernel approximation, but this may not help all downstream metrics. We also include test predictive RMSEs with the exact kernel (to which we converge) and linear regression for comparison.
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
L.D.
1,000 #0.028
1.000
=0.042
1.000
=0.052
1.000
=0.057
1.000
=0.025
1.000 =0.015
HALTON
1,028 +0.020
0.991
+0.042
0,995
=0.082
0.913
=0.05s
0.927
1.176 +0.022
ORTHOGONAL
0.627
0.535
20.023
0.617
0.669
20.024
0.586
20.013
0.639 20.016
PNC
0.563 100001
0.433
10.01M
0.544
10.07
0.547
10.0a0
0.481
10.011
0.606 10.018
LAPLACE PEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
I.L.D
1.000 20.0022
1,000
20003
1.000
2o.0ss
1.000
2oams
1.000 10.025
1.000 10.020
HALTON
0.721 #0.067
0.777
=0.081
0.779
=0.054
0.728
=0.015
0.721
=0.021
0.893 =0.025
ORTHOGONAL
0.418 20.00
0.546
20.026
0.614
20.00
0.527
20.033
0.489 20.016
0.360 10.010
PNC + ANTITHETIC
0.367
0.486
0.618
20.110
0.438
0.418
0.324
Table I: Performance of RFFs and RLFs on kernel estimation with UCI datasets with different coupling schemes, taking kd random frequencies for RFFs and 2kd for RLFs where & E N (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 equal to one. Lower is better. Error bars are standard errors on RMSES.
taking a multivariate Gaussian and pushing each of its coordinates forward first with the Gaussian CDF Fx. and then the Xa inverse CDF F=1 If the diagonal terms of the underlying Gaussian covariance matrix Σ E _mxx 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 Σ. 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 Σ, 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 differentiable. In doing so, optimise the RF coupling. Remarkably, this data-dependent optimisation does not to find couplings much better than PNC: see the training curves in Fig. 6. This suggests that scheme may already be close optimal for m * 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., 2019), this is (to our knowledge) their first application to learning a Monte Carlo coupling.
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 GP (RFFs) or selecting reasonable values for Performers (RLFs) (Choromanski et al., 2020). We then compute the kernel approximation RMSE (Frobenius norm error of K) test set. Full details are in App. B.I.
Results for variance reduction. Table shows the results. For RFFs, we take m = kd orthogonal frequencies, with & E N. For RLFs, we also include their antiparallel directions, giving frequencies. For each dataset, the RMSEs are normalised by the result with i.i.d. features, which means the results are the same for all k. As baseline, we include RFs constructed using Halton sequences (Dick et al., 2013; Yang et al., 2014a), a fixed, off thelhhh--hlf QMC scheme that can provide small gains but is clearly suboptimal. The third row shows orthogonal frequencies with independent
2kd
5

---

Published as a conference paper at ICLR 2025
norms (Yu et al., 2016). When also couple the frequencies' norms using 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 2 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., 2016; Rowland et al., 2018; Reid et al., 2023; Likhosherstov et al., 2022; Yang et al., 20: Le et al., 2013; Bojarski et al., 2017; Choromanski et al., 2017; Lyu, 2017: Shen et al., 2017: Dao al., 2017: Munkhoeva et al., 2018). 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. Fig. I demonstrates this on a train-fest split of the POWER dataset for different values of m (we use single split to highlight the small spread between different coupling methods on fixed data). We find that lower kernel estimator RMSE may not improve downstream performance. This is because, when optimising a coupling, we minimise the variance of pointwise kernel estimates {A= (x)) )) 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 concrete example, consider estimating attention, d.g = k(z,, //E/" T((,, 21) using random Laplace features (Choromanski et al., 2020). This normalises the kernel evaluation between the i and j tokens by the sum with all the other tokens. Taylor expanding. if the kernel ex
estimators have equal means Je, the average mean square error MSE(a.) == x دj=1 MSE(a,;) obeys
(SE( =
1
Var(k(a; x;)) N
Cov ajj), R(s,,, f1
+ O
(8)
N2
N
N8
By coupling the frequency norms, PNC reduces Var (ao æ;))) as intended. However, it also reduces the covariance (ovkk(z,,,,,
Table 2: Performer test accuracies on Im- ageNet with different coupling schemes. Counterintuitively, maximising the point- wise kernel estimator variance by posi- tively correlating feature norms boosts performance a different OT problem to the naive, obvious choice
((,,,,2))) so MSE();)
does not actually substantially improve overall. In stark contrast, if we instead take the positive mono- tone (PM) coupling where [wi] are all equal almost surely, then Var (kr,,, æx,) is maximised (see App. B.3). But these strong, positive correlations also increase Cov(k(a,, x,,,) ), k(w, x;,)) by an even greater amount. Hence, we find that MSE(a.) (a.) falls (see Fig. 7). This is surprising: maximising the pointwise kernel estimator vari- ance 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., 2009) by 40.8% whereas PNC makes no statistically significant difference. See Table 2. This demonstrates the limitations of simple variance reduction and invites more careful treatment, considering the downstream quantities of interest. App. B.3 gives further discussion and transformer training details.
LAPLACE FEATURES
TEST ACC.
ORTHOGONAL
0. 0.625 +0.000
ORTHOGONAL PNC
0.620 10.003
ORTHOGONAL. + PM
0.633 20.003
4 GRAPH RANDOM FEATURES
We now shift our attention from Ed to the discrete domain. Consider an undirected graph G(N, where N := {v1,... ..., } is the set of nodes and E is the set of edges, with (v;, 13y) E & if and only if there exists an edge between and Uj in G. Graph node kernels k. : x N + R are positive definite, symmetric functions defined on pairs of nodes of G. reflecting some notion of their 'closeness' via the graph edges and weights. h captures the structure of G, letting practitioners repurpose popular kernelised learning algorithms to the discrete domain (Smola and Kondor, 2003b;a). 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, 2002). We give short introduction in App. C.I.
Graph random features. As in the Euclidean setting, graph kernel methods scale poorly due to the O time complexity of inverting the Gram matrix K == [l<(1, -=) In fact, even computing K often incurs a cubic cost since it involves multiplying large adjacency matrices. Research has been dedicated to improving efficiency by approximating K, including graph random features (GRFs)
6

---

Published as a conference paper at ICLR 2025
matching between quantiles
graph random walks
Figure 2: Schematic overview of o-coupled GRFs with 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.
(Choromanski,, 2023; Reid et al., 2024b). These are sparse vectors (ØCRF((;) (c:) C ER that satisfy K, = E [(GRF(V;) (GGF(U;)) (v;)], so their dot product equal true graph kernel in expectation. GRFs are constructed in subquadratic time.
Coupled random walks. For GRFs, the 'frequencies' {w, are simple random walks: sequences of graph nodes (ot)f=1 with (u, 0(41) E. At every timestep, the walker chooses of its neighbours uniformly and at random. The length of the walk is also a random variable, drawn from geometric distribution ~ G(p), p (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 al., 2007; Zhou al., 2015). 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 must remain geometric. Reid et (2024c) proposed 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.
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. gives visual overview.
Constructing GRFs from random walks. To obtain (w;) E RN,, one samples m random walks (4) out of node U, E A and averages their 'projections', ( 1k=1
ØGRE
(w)
(9)
772
The projection function (-). Ω + R- maps from the set of graph random walks Ω (1:)//11 N,, 1241) EN to sparse N N-dimensional feature vector satisfying k(2;, U2) = Eni w wi) (+(w(() {(w(i))). v6() depends on the particular kernel being approximated. We direct the reader to the work of Reid et al. (2024b) for an introduction to GRFs, and also provide background in App. C. vb()) is complicated function; is difficult reason about analytically but straightforward to compute for particular walk. Moreover, its input walks are discrete random variables so (²() not differentiable with respect the lengths ((9) (where 1,.., m and vs N is the walker's ' start node). This precludes straightforward gradient- based optimisation (Sec. 3.2).
A pair of walkers. Initially consider just = 2 walkers. The kernel estimator Te(1r, v3)) is
7e(1, 0y) = (v. ( = (w)()( + (((w) (v(w)()) + v(wo)))) (10)
7

---

Published as a conference paper at ICLR 2025
(6.5) which is unbiased provided (i) the marginal distribution of each w1,2 is simple random walk with geometrically distributed length (hereafter denoted 1)a) and (ii) walks from node us are independent from walks from node By. The variance of the estimator depends on
- 16 1 ([(oteP) + m(w5") (((v(c)))+(()))))))
(W(W))) + v( ))
E (()))))))))
(II)
where the expectation is taken over both the directions and lengths of the random walks, Suppose the directions remain independent, but the lengths 1 and 1/" (and likewise (79) and (64) are to be coupled for variance reduction, analogously to the vector norms in Sec. Let Edirs denote the expectation over the walkers' directions. We want to minimise:
([(vtof)) + Mw(o))) +(
T (W(w))) + Ma.? (v((o)))+
I. uff 1mppo yop Baires
for E A2(ro)- (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 Monte Carlo approximation, we move Edirs inside the square. This is because, unlike the expression in Eq. 12, Ean. (1((w((2)) )) can be efficiently estimated by simulating random walks. Second, we must optimise amongst the class of couplings (2a)) 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 A2(r))). Taking inspiration from numerical OT, consider the family of measures (r() on [0, 1J2 described by the permutation densities Po (z. = nl[[[nx]) [ng]] with o a permutation of order n (that is, a bijection a : [m]] + [n]). The unit square is split into a n x n grid where each row and column has single 'tile' of probability density and is 0 otherwise (Fig. 2 left). Both marginal distributions of ((") are uniform on [0. 1] and are transformed to probability distribution by pushing forward with the inverse CDF, Fy () = in E l. (x) 2 J.. Transforming both coordinates in this way yields a joint measure ( E A2(1) which will give an unbiased estimator. n! such couplings exist for given permutation order ne; we aim to efficiently find the one with the lowest variance.
The permutations a E Sn 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 a, 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 a. the approximate OT problem can be written
o' = arg min
Σ
Σ
( (#(1))+ +(o()))))) T (oldl)
(13)
SES
11E 12
where ()) = E " .2)) (Bang (1)) U((a,) b)) is the uniform distribution on the interval (a, b]. Eq. 13 is a quadratic assignment problem (Finke et al., 1987; Burkard al., 1998). 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). Given a set of vectors ())g(()), D(gh))) CRN, Eq. 13 can be solved with high probability in time complexity independent of N. Moreover, in the special case where = j. can be solved in polynomial time in (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., 2010), which preserves pairwise dot products with high probability.
the special case = j. Fq. 13 can be rewritten as finding permutation a that minimises the L2 norm of some particular -imddimensoon vector. In App. C.3.1 we provide novel algorithm to achieve this efficiently by projecting onto a sequence of random Gaussian vectors, requiring only a mild geometrical condition called c-separation. See Lemma C.2 for full details. O
More pragmatically, one can set 91 = 92 to simplify to a related linear assignment problem, which can be solved efficiently in time Onz²) using e.g. the Hungarian algorithm (Kuhn, 1955). We empirically investigate how this final approximation modifies the objective in App. C.3. Taking the optimal permutation and corresponding coupling v1)), we define couupled GRFs as follows.
8

---

Published as a conference paper at ICLR 2025
Definition 4.2 (a-coupled GRFs). GRFs are a-coupled they are constructed using puirs of random walks with lengths drawn from the coupling م μ(0), with the optimal permutation a obtained by solving a matching problem between the quantiles of the distributions over walk lengths (specifically, Eq. 56 in App. C.3).
Besides being easy to optimise and sample from, there are also more rigorous OT motivations for the choice of o-couplings μ!o They relate to the asymptotic behaviour of μ/() as the permutation order + oo and the stability of OT plans (Villani, 2021). 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 combine m 1 permutations, finding a saaiuum-weight, m- mn-partite matching with all subgraphs constrained to be complete Kn- Another approach would be to approximately solve this multi-marginal OT problem using the Sinkhorn-Knopp algorithm (Sinkhorn and Knopp, 1967; Cuturi, 2013).
Broader applicability. As 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 graph proposed by Page al. (1998) rank websites in search engine results. r-couplings consistently match or beat the previous best algorithm for coupling walk lengths (Reid al., 2024c). See Fig. 9 in App. E for the full results.
4.2 EXPERIMENTS WITH COCUPLEED GRFS
We now empirically evaluate a-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 Phatt as a hyperparameter, setting the rate of decay of the geometric distribution over walk length. smaller value of Phah samples longer walks and gives more accurate kernel estimates, but takes longer to run. The optimal coupling depends on Phalt- We consider the values Phat E {0.1. 0.2, 0.3, 0.4, 0.5), finding the optimal permutation in each case. To find the optimal Ø,, we solve the matching problem (App. C.3) for a random Erdös-Rényi graph with N nodes, taking a permutation order n = 30 and choos- ing the 2-regularised Laplacian kernel as our target. We use the Hungarian algorithm, averaging the cost matrix over every possible node pair (vi, U5) E E
100
GRFs kernel estimator error
rolative Frob. norm error
Having computed couplings yu(o) for each Phalt, we then test the corresponding a-coupled GRFs variety of real-world graphs. Fig. 3 shows the re-
L.tdd antithetic
sults for cora (N 2708), with the rest left
0.50
o-coupled
App. F.I. We plot the relative Frobenius norm error of the Gram matrix approximation K K |e/| K|F with walkers that are i..dd, antithetic (Reid et al., 2024c) or a-coupled. For each Phalt, a c-culledd GRFs give equally good or smaller kernel estimator errors. Our OT approach outperforms antithetic termination: data-independent, hard-coded algorithm designed specifically to improve GRFs (Reid et al., 2024c). We include visualisations of the optimal permutations for different values of Phalr in the inset, verifying that the o-coupling adapts to different hyperparamaters.
0.25
0.1
0.2
0.3
0.4
0.5
termination probability
Figure Kernel estimator error vs termina- tion probability. Insets show permutations for Phalt E (0.1, 0.2, 0.3, 0.4,0.5), 0.5).
Novel application: r-coupled GRFs for scalable graph-based GPs. We now apply o-coupled GRFs to scalable graph-based Gaussian processes! I (GPs), where improved estimation of the covariance function permits better approximate inference. Scalable GPs are novel application of GRFs that may be of independent interest (Borovitskiy et al., 2021; Mostowsky et al., 2024).
Consider the task of probabilistic graph interpolation. This aims predict unknown graph function values, along with principled uncertainty estimates, from an observed set (Pfaff et al., 2020). Take
This may be better described by Gaussian distributions than Gaussian processes since for fixed G we have finite number of random variables, hut use the latter for consistency with recent literature.

---

Published as a conference paper at ICLR 2025
mesh graphs 9 where every node V, V has normal vector ne, E R (Dawson-Haggene) 2023).
Our task is to predict the z-components of a masked set. {(n:)= Mue with Nis = [0.05N]. To achieve this, we use graph-based GP with heat kernel covariance function. We compute sparse, unbiased approximation of this kernel using GRFs with {16,35 32, 64) walkers that are i.i.d., antithetic (Reid al., 2024c) a-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 >
teapot
fuze_ascii
sphere
idler riser
busted
torus
[| -ip
0.12
0.12
0.00
0.03
0.05
0,05
test RMSE
0.20
0.30g
0.16
0:11
0.05
0.20
0:12
KL. div.
20
DO. walkers
L.d.
antithetic
a-coupled
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. o-coupled GRFs give better predictions and uncertainty estimates, sometimes by a factor of > 2. Standard errors are shaded.
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. (2021) using data from Chen et al. (2001) and StenStreetMap. The graph has 1016 nodes, with speed only known at 325, We use 250 randomly-chosen nodes training data and the remainder as test data. With this sparse, noisy dataset, a-coupled GRFs again give substantially better predictions and uncertainty estimates. Fig. 11 in App. F33 shows the full results.
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 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 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., 2020), 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.
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 kerne] estimates determines downstream performance. Coupling to optimise e.g. the spectral properties of K (Choromanski et al., 2018; Avron et al., 201 7a) or the variance of row-normalised attention scores may prove better. These objectives rarely considered in the literature. Fortunately, OT provides 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.
10

---

Published as a conference paper at ICLR 2025
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 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.
IR acknowledges support from a Trinity College External Studentship and a Google PhD Fellowship. Part of the work was completed as student researcher at Google. SM acknowledges funding from the Vice Chancellor's and the George and Marie Vergottis scholarship the Cambridge Trust, and the Qualcomm Innovation Fellowship. RET is supported by the EPSRC Probabilistic AI Hub (EP/Y028783/1). AW acknowledges support from Turing Al fellowship under grant EP/V025279/1 and the Leverhulme Trust via CFI.
We thank Bruno Mlodozeniee 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.
ETHICS AND REPRODUCIBILITY
Our work is foundational with 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 as unintended consequences.
Ethics statement:
Every effort has been made to ensure the work's '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 the main text for clarity. All datasets are available online. We give links to suitable repositories in every instance. Where possible, results are reported with uncertainties to facilitate comparison. Code available at: https: // /github com/ camb rettrrddge widge-ngee mmaee learnable e-qmc.
Reproducibility statement:

---

Published as a conference paper at ICLR 2025
REFERENCES
Noga Alon, Itai Benjamini, Eyal Lubetzky, and Sasha Sodin. Non-backtracking random walks mix faster. Communications in Contemporary Mathematics, 9(04):585-60) 2007. URL https:
//doi.org//0/ 01t..0///////00099997777000
Haim Avron, Michael Kapralov, Cameron Musco, Christopher Musco, Ameya Velingker, and Amir Zandieh. Random fourier features for kernel ridge regression: Approximation bounds and statistical guarantees. International conference on machine learning, pages 253-262. PMLR, 2017a. URL https://dolt www.ooog./..oooo.oooo 48550/ arXiv. . 1804 09893.
Haim Avron, Michael Kapralov, Cameron Musco, Christopher Musco, Ameya Velingker, and Amir Zandieh. Random fourier features for kernel ridge regression: Approximation bounds and statistical guarantees. Proceedings of the 34th International Conference on Machine Leurning, ICML 2017, Sydney, NSW. Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 253-262. PMLR, 2017b. URL http: /tte:://rrooeee /ptp//eeiccccesseeeeeeeee/e/s.ee mll. pt....press//0/ avron17a. html.
Chandra R Bhat and Aupal Mondal. On the almost exact-equivalence of the radial and spherical unconstrained choleskybased parameterization methods for correlation matrices. Technical report, Department of Civil, Architectural and Envirommental Engineering, The University of Texas at Austin, 2021. https: /ttttts::///ppoeo http:////epppsttttrreeeeeeo 1ib. utwxas edu/ server/eveeerpee/ /core/
bitstreams 69c0499c 21t-aa--21552--222---2
www.aaaa45444055555 454445555500/555ttt5..5./ content.
Patrick Billingsley. Convergence of probability measures. John Wiley & Sons, 2013.
Mariusz Bojarski, Anna Choromanska, Krzysztof Choromanski. Francois Fagan, Cedric Anne Morvan, Nouri Sakr, Tamas Sarlos, and Jamal Atif. Structured adaptive and random spinners for fast machine learning computations. In Artificial intelligence and statistics, pages 1020-1029. PMLR, 2017. URL https: /ttp://oo...... org/ 488500 arxiv. 1610. 06209.
Viacheslav Borovitskiy, Iskander Azangulov, Alexander Terenin, Peter Mostowsky, Mare Deisenroth, and Nicolas Durrande. Matérn gaussian processes graphs. In International Conference Artificial Intelligence and Statistics, pages 2593-2600 PMLR, 2021. Uttpttps: //t:::///oo/// ht/ooi..sst org/ 10. 48550 arw.i...vv..... . 2010 110.111..3010.11 15538.
Rainer Burkard, Eranda Cela, Panos M Pardalos, and Leonidas Pitsoulis. The quadratic assignment problem. problem. Springer, 1998. URLhttps, /ww:///wwww///...///////////www opt mat. wuggaz.at//caee... papers/qap_bers cht pdf.
Colin Campbell. Kernel methods: survey of current techniques. Neurocompurinng. 48(1-4):63-84 2002. URL https: /oii org/ 1et.e/11111/955550 //////0925521111 (01) 00643-9.
Stéphane Canu and Alex Smola. Kernel methods and the exponential family. Neurocomputing, 69 (7-9):714-720, 2006.
2006. URL https: /ww. wttp:/o...rr.....0.............. 1016///...11/11 nwwco....mmmmmmmmm 2000..000..00 009.
Olivier Chapelle, Jason Weston, and Bernhard Schölkopf. Cluster kernels for semi-supervised learning. Advances in neural information processing systems, 15, 2002. URL https: //d1. acm. wwtp/oi/////////// /tt/:/11.55555555533 01155552266618.29685533 2968693.
Chao Chen, Karl Petty, Alexander Skabardonis, Pravin Varaiya, and Zhanfeng Jia. Freeway perfor- mance measurement system: mining loop detector data. Transportation research record, 1748(1): 96-102, 2001. URL https: : api.
Jinjin Chi, Jihong Ouyang, Ximing Li, Yang Wang, and Meng Wang. Approximate optimal transport for continuous densities with copulas. In IJCAI, pages 2165 2019. URL https: :/wbs:///www.. /www. ijcai. wttpa//reppp.oociigg....... 2019 9/ 0300. pdf.
Krzysztof Choromanski, Mark Rowland, Tamás Sarlós, Vikas Sindhwani. Richard Turner, and Adrian Weller. The geometry of random features. In International Conference Artificial Intelligence and Statistics, pages 1-9. PMLR, 2018. URL http: http:://pocceedings... mlr. press/////eee.. / choromanski Ba choromankkiiooiiccoocooo pdf.
12

---

Published as a conference paper at ICLR 2025
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXlv:2009.14794, 2020. URL https://doi.....
48550/arXiv. 2009. 14794.
Krzysztof M Choromanski. Mark Rowland, and Adrian Weller. The unreasonable effectiveness of structured random orthogonal embeddings. Advances in neural information processing systems, 30, 2017. URL https::tp......m.....m pentts://dee /doi. .rgg 48550/ axsivv 1703. 00864.
Krzysztof Marcin Choromanski. Taming graph kernels with random features. International Conference Machine Learning, pages 5964-599 5977. PMLR, 2023. URL httpps://dtt .rg/ 10. 40.54850/40 arxiv. w...i.2300.. 00156. 56.
Fan RK Chung. Spectral graph theory, volume 92. American Mathematical Soc., 1997.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in
neural information processing systems, 26, 2013. UtLhttps: / /wppers.ppsppppp//p/ wwwppccpppaaaaaaamomma ccppppppp.........ppppe files lale//
Tri Dao, Christopher M De Sa, and Christopher Ré. Gaussian quadrature for kernel f Advances in neural information processing systems, 30, 2017. URL arXiv. 1709. 605.
Anirban Dasgupta, Ravi Kumar, and Tamás Sarlós. A sparse johnson: Lindenstrauss transform. In Proceedings of the forty-second ACM symposium Theory of computing, pages 341-350, 2010. URL https: /oi...i..///// org. 101.:oo000011111111800666665
806737.
Sanjoy Dasgupta and Anupam Gupta. An elementary proof of theorem of johnson and lindenstrauss. Random Struct. Algorithms, 22(1):60-65 2003, doi: 10.1003//SA.10033 URL https: /doi.
org/10.1000 rsa. 1asaa.1000. 10073.
Michael Dawson- Haggerty Trimesh repository, 2023. URL htt.aatsss//gtthb.ccoo//at/ trimesh.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference computer vision and pattern recognition, pages 148-255. leee, 2009. URL https::tp......c..... doi org ot../..1///1///////......1..1111
www.:/1119........ CC/R///PRRR 2009. 520 6848.
Josef Dick, Frances Y Kuo, and lan H Sloan. High-dimensional integration: the quasi-monte carlo way. Acta Numerica, 22:133-288, 2013. URL https: http:///ooooorg//.../
$0962492913000044.
Alexey Dosovitskiy. Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprins
285.44455500//0 20100.2.01.111.. arXiv:2010.11929 2020. URL https: http:tt//oooooor//////
2010.
Gerd Finke, Rainer E Burkard, and Franz Rendl. Quadratic assignment problems. In North-Holland Mathematics Studies, volume 132, pages 61-82. Elsevier, 1987. URL battps://dol. //tp:::///oo...0...... org/10. 101 20116/00004-00 - 0208 (08) 73232-8.
Dániel Fogaras, Balázs Rácz, Károly Csalogány, and Tamás Sarlós. Towards scaling fully personalized pagerank: Algorithms, lower bounds, and experiments. Internet Mathematics, 2(3):333-335 2005.
URL www.ttttp/... /tw.///oiiiii////o/////oooo org/ / 10. 1007 0900077977-10 - --0---3--000000 00440302266-22
Paul Glasserman and David Yao. Some guidelines and guarantees for common random numbers.
Management Science. 38(6):884 908, 1992. URL httpp:://oi..rr/////// 1287 ///////sc... 38.
6. 884.
John H Halton. On the efficiency of certain quasi-random sequences of points in evaluating multi- dimensional integrals. Numerische Mathematik, 2:84-90, 1960. URL https: /ww. org/rogg.....0 /10. 1007 /FF 08F01386211
13

---

Published as a conference paper at ICLR 2025
JM Hammersley and KW Morton. A new monte carlo technique: antithetic variates. In Mathematical proceedings of the Cambridge philosophical sociery. volume 52, pages 449-475. Cambridge University Press, 1956. URL https: http://:///ii....1..////o
Martin Haugh. An introduction to copulas. quantitative risk management. Lecture Notes. New York: Columbia University, 2016. URL http:://twwtt http://wwwwwllllcoc/mm cow.mbia 8
Copulas pdf.
Vladimir Ivashkin. Community graphs repository, 2023. URL https::/gttthbbbe
v1 vlivashking deivashkin/commee hothi///noooorritrgaarph...
William B Johnson. Extensions of lipschitz mappings into hilbert space. Contemp. Math., 26:189- wdw/classccsllllccccaaaaaae re Johnsson.
206, 1984. URL hetp://stanfee pdf.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6900 6980, 2014. URL https: //w.iiioooooooo org.iorrg1000 485.0//0...00//.........
6980.
Risi Imre Kondor and John Lafferty. Diffusion kernels on graphs and other discrete struc- tures. In Proceedings of the 19th international conference on machine learning, volume 2002, pages 315-322, 2002. URL https: /wwwwww.... www..eeu/eeseeercc/aaaaaaaaappaa kondor dinioor-iffe diffusicrrr kernels. ..ff.
Leonid Kontorovich, Corinna Cortes, and Mehryar Mohri. Kernel methods for learning languages.
Theor. Comput. Sci., 405(3):223-233 236, 2008. doi: wtt.:/..cs.2008.66.. URL https:
10. 1016 tes. 2008 00....66.00 037.
Harold W Kuhn. The hungarian method for the assignment problem. Naval research logistics
quartenly, 2(1-2):83-97, 1955. URL https: /tt:::////oioi// http:///ooo...r/..//../ 10t///00..00/////../// 3800020003
Miguel Lázaro-Gredille Gredilla, Joaquin Camnonero-Candele Carl Edward Rasmussen, and Aníbal R Figueiras- Vidal. Sparse spectrum gaussian process regression. The Journal of Machine Learn- ing Research, 11:1865-1881, 2010. URL https://jml= csail ..m.t.e.uu /apers/ 0a wawlaaaao-ggreed wtt.gggeaaaaaaaaaaaaooorrreeerea llalla.ll.lll .
Quoc Le, Tamás Sarlós, Alex Smola, et al. Fastfood-approximattng kernel expansions in loglinear time. In Proceedings of the international conference on machine learning, volume 85, 2013, URL
https://doe http:://dooooo/g/1/// / 48550/ arxiv. 1408.34 00.00033060..
Valerii Likhosherstov, Krzysztof M Choromanski, Kumar Avinava Dubey, Frederick Liu, Tamas Sarlos, and Adrian Weller. Chefs' random tables: Non-trigomometric random features. Advances in Neural Information Processing Systems, 35:34559-34573, 2022. URL https: /wwii org/
10. 4054445500/00 arxivv 2205.15317
Fanghui Liu, Xiaolin Huang, Yudong Chen, and Johan A. K. Suykens. Random features for kernel approximation: A survey on algorithms, theory, and beyond. IEEE Trans. Pattern Anal. Mach. Intell., 44(10)::128-777 -1441 2022. doi: 10.1 10111099AM..20011000000 //doi. 110.11999 19/ TPAMI. 2021.3097011
URL https:
Siqiang Luo. Distributed pagerank computation: An improved theoretical study. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 4496-4503, 2019. URL https: //w.i//////////////////oo /ro//0.......... 1609 aaaii v33101. .301143044 4496.
Yueming Lyu. Spherical structured feature maps for kernel approximation. In International Con- ference on Machine Learning, pages 2256-2264. PMLR, 2017. URL http::/proceeeding
mlr. mlr.press/v70 1yu1 7a /yylllll1/ 7...f....p...f....
William J Morokoff and Russel E Caflisch. Quasi-monte carlo integration. Journal of computational physics, 122(2):218-230) 1995. URL https: //doi. /rg/10. 1ww.1000/jippi 1995 1209.
Peter Mostowsky, Vincent Dutordoir, Iskander Azangulov, Noémie Jaquier, Michael John Hutchinson, Aditya Ravuri, Leonel Rozo, Alexander Terenin, and Viacheslav Borovitskiy. The geometrick- emnels package: Heat and mat', emn kernels for geometric learning manifolds, meshes, and graphs. arXiv preprint arXiv:2407.08086, 2024. URL https: /oi.. worg/org...... org/10.485500 48550/ 2407.08086.
arXiv. 2407. 08086.
14

---

Published as a conference paper at ICLR 2025
Marina Munkhoeva, Yermek Kapushev, Evgeny Burnaev, and Ivan Oseledets. Quadrature- based features for kernel approximation. Advances in neural information processing systems, 31, 2018. URL httph::t/////.
A8w50/aaKKa... arXiv. 18018802222 03832.
Roger B Nelsen. An introduction to copulas. Springer, 2006.
Lawrence Page, Sergey Brin, Rajeev Motwani, and Terry Winograd. The pagerank citation ranking: Bring order to the web. Technical report, Technical report, stanford University, 1998. URL https: //www. cis upenn. /tttt//mmeaassseeeeeeeeeeee http:///aaa.eeeccaaaa/a///////aa wagerakk. pdf.
Victor Panaretos and Yoav Zemel. An invitation to statistics in Wasserstein spuce. Springer Nature, 2020.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, and Peter W Battaglia. Learning mesh- based simulation with graph networks. arXiv preprini arXiv:2010.03409, 2020. URL https: // //doi. , /tt....o//.......////////..o
2010. 03409.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007. URL https: httlhhttss//ptttt /ttp://ppppp....eeee eecs. berkeley.. r.r.rah.rr rew.rrec..pppp....... pdf.
edt/-brecht/paperz
Ali Rahimi and Benjamin Recht. Weighted sums of random kitchen sinks: Replacing mini- mization with randomization in learning. In Advances in Neural Information Processing Sys- tems 21, Proceedings of the Twenty-Second Annual Conference on Neural Information Pro- cessing Systems, Vancouver, British Columbia, Canada, December 8-11, 2008, pages 1313- 1320, 2008. URL https: https: //tp:::///ppp/aa//eeeee nips. wtwpppssccppppprrrpppppp papprr Defe32849d2222
pdf.
Isaac Reid, Krzysztof Marcin Choromanski, Valerii Likhosherstov, and Adrian Weller. Simplex random features. In International Conference Machine Learning, pages 2023. URL https: htt:///oi.iio/////// http::::///....o//..oooo//oo/ axii0/ 2301 13856.
PMLR,
Isaac Reid, Eli Berger, Krzysztof Choromanski, and Adrian Weller. Repelling random walks. In International Conference on Learning Representations, 2024a. URL http:://s////.ooo
48550/ arw.v..iiivvv. 2310. 04854.
Isaac Reid, Krzysztof Choromanski, Eli Berger, and Adrian Weller, General graph random features. In International Conference Learning Representations, 2024b. URL https: //c::///oiii org/ 10. 48550/ arXiv. .2r.23311 2310 04859.
Isaac Reid, Adrian Weller, and Krzysztof M Choromanski. Quasi-monte carlo graph random features. Advances in Neural Information Processing Systems, 36, 2024c. URL https: /ttp:///oo.../... oro/o//g/l0..o /10 48550/arXiv 220...iii.........mm 12470.
Mark Rowland, Krzysztof M Choromanski, François Chalus, Aldo Pacchiano, Tamas Sarlos, Richard E Turner, and Adrian Weller. Geometrically coupled monte carlo sampling. Advances in Neural Information Processing Systems, 31, 2018. URL https: //lg////mgg///// eng. cam. ac.ukk .uk / 8- gemc.pd..
Bernhard Scholkopf and Alexander J Smola. Learning with kernels: support vector machines, regularization, optimization, and beyond. MIT press, 2018.
Weiwei Shen. Zhihui Yang, and Jun Wang. Random features for shift-invarian kernels with moment matching. In Proceedings of the AAAI conference on artificial intelligence, volume 31, 2017. URL https: doi. .orgg//ggg... / 0. 609/ aaai v31 ....1..111........1..1111111.111 108.5.
Richard Sinkhom and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):343-348, 1967. URL https: //sppmppp./// .arsp.org/a http////gggp/////// 21 2112/pjmmv21--1
pd...
M Sklar. Fonctions de répartition dimensions et leurs marges. In Annales de I'ISUP, volume 8, pages 229-231, 1959. URL https:tttt.........t. /hal. science//eeeccccccc hha 040944433334 63/ccumen..
15

---

Published as a conference paper at ICLR 2025
Alexander J. Smola and Risi Kondor. Kernels and regularization graphs. In Computational Learning Theory and Kernel Machines, 16th Anmual Conference on Computational Learning Theory and 7th Kernel Workshop, COLT/Kemell 2003, Washington, DC, USA, August 24-27, 2003, Proceedings, volume 2777 of Lecture Notes in Computer Science, pages 144-158. Springer, 2003a. URL www.t:tpp/...o..... /oi .org/ 1ttprr//11100077--3-311111-
Alexander J Smola and Risi Kondor. Kemelss and regularization on graphs. In Learning Theory and Kernel Machines: 16th Annual Conference on Learning Theory and 7th Kernel Workshop, COLT/K- ernel 2003, Washington, DC, USA, August 24-27, 2003. Proceedings, pages 144-158. Springer, cs www.uccicaa. eeu/-riii -rist s / aKondor
2003b. URL htttss://ppptt pdf.
Matthew Thorpe.
Introduction to optimal transport.
Lecture Notes, 3, 2019. URL Opt. imal_
https: www.damtppppppoo cam ac. transport_Notest pdf.
Austin Tripp.
Sergio Bacallado, Sukriti Singh, and José Miguel Hemández-Lobato. Tanimoto random features for scalable molecular machine learning. Advances in Neural Information Processing Systems, 36, 2024. URL https::/p///o......mm /tts::///doo /doi..... org/10 /10. . 3550//ariiv 2306. 44899.
Cédric Villani. Topics in optimal transportation, volume 58. American Mathematical Soc., 2021.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2009.
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning, volume 2. MIT press Cambridge, MA, 2006.
Jiyan Yang. Vikas Sindhwani, Haim Avron, and Michael W. Mahoney. Quasi Quasi-monte carlo feature maps for shift-invarian kernels. In Proceedings of the 31th International Conference on Machine Learning, ICMI. 2014, Beijing, China, 21-26 June 2014, volume 32 of JMLR Workshop and Con- ference Proceedings, 2014a. URL http: /ttp:://zzee html.
h http:///pocceeeig....o mlr, v httpa////////aggll//// yangb14,
Jiyan Yang, Vikas Sindhwani, Quanfu Fan, Haim Avron, and Michael W Mahoney. Random laplace feature maps for semigroup kernels on histograms. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 971-978, 2014b. URL https://vikaa. sindhwani. org htt///annnmmaaaaaeeeee pdf.
Felix Xinnan X Yu, Ananda Theertha Suresh, Krzysztof M Choromanski, Daniel N Holtmann-Rice. and Sanjiv Kumar. Orthogonal random features. Advances neural information processing systems, 29, 2016. URL https: doi. org/10, 48550/ aarii..iii.v..... 1610 09072.
Yin-Cong Zhi, Yin Cheng Ng, and Xiaowen Dong. Gaussian processes graphs via spectral kemnll learning. IEEE Transactions on Signal and Information Processing over Networks, 2023. URL https: //ooi org/ 1 48550/ rxt 2006 07361.
Zhuojie Zhou, Nan Zhang, and Gautam Das. Leveraging history for faster sampling of online social networks, arXiv preprint arXi::7505.00079, 2015. URL https: ArXiv.150s 00079.
48550/
16

---

Published as a conference paper at ICLR 2025
A SOLVING THE OT PROBLEM FOR RFFS AND RLFS
In this appendix, we provide proofs for the thcoretical 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 I.I. Consider the following recently-derived result by Reid et al. (2023).
Lemma A.1 (Kernel estimator MSE for RLFs (Reid et al., 2023)), When estimating the Gaussian kernel k( y) == exp(-l< x 2l1/22 for datapoints {x, y) C R² using random features features (synonymously, positive random features), the mean square error of the kernel estimate k( y) is given by:
22
MSE (k:)
(m ()(p(2, y) -60))
(14)
m where m is the number of sampled random frequenics,,,,,, v == | + 2z ll2 is data- data-dependent scalar and plas; D.)) is the RF-conformity,
I(
24 22k (K + 4)
2 wi
p(x, y) = m(m
Σ
Euss
(15)
Liti
Here, Wij == l|w; + will2 la the norm of the resultant of a pair of distinct frequencies and the Gamma function.
Proof. Reid et al. (2023).
The simple derivation, reported in full by Reid et al. (2023), is based on rewriting the angular integral as a Hankel transform, yielding a Bessel function of the first kind with known Taylor expansion.
Supposing that w; 1 wy, we have that waj / wuf + wh where w, = ||, (wlll2 denotes the L2 Lormorm of the frequency w; Note that, for w; to be marginally Gaussian, we require that the marginal distribution over its norm is Xd. Chi distribution with d degrees of freedom. Substituting this into Eq. 14 and dropping terms unmodified by the coupling and irrelevant multiplicative factors, is straightforward arrive at the expression for the cost function CRLF Eq. 6.
Finding the cost function for RFFs is only slightly more difficult. We begin by citing another recent result by Reid al. (2023).
Lemma A.2 (Kernel estimator MSE for RFFs (Reid et al., 2023)). When estimating the Gaussian
kernel k(x y) exp(-Ilæ yll./2) 22) (2) for datapoints {æt y) C Rd using random Fourier features, the mean square error of the kernel estimate k(æ,, y) is given by:
1 (1- (1-e-xt)
MSE(K)
(m 1)(((z,)) - c-))
(16)
m
2
where is the number of samples,
y and Ç(æx, y) is defined by
1
((x,3) 3) == m(m 1) 1,77
Bun, [cos(w]] T z) cos(w)] x)]
(17)
Proof. Reid et al. (2023).
Note the close resemblance to Eq. 15. The only term that depends on couplings between the random
frequencies [w; is ((x, y), which we seek to suppress with carefully engineered correlations. cos (Iw; way 2). Defining the random variables 2/2 and integrating out the angular part,
From elementary trigonometry, cos w. cos w 2 = : (cos ((w; + +w; z)
and for any coupling scheme w; + w is isotropic.
w(+) = ||W + W/ and w|-) = ||w; - -will
((x,, 2zy) ((),,,))) = m(m - 1) 1.7
Σ
rdd/222 -B_ ()) 22 # Jg (( (
(w(+)))² -
4 -- (2))
(18)
17

---

Published as a conference paper at ICLR 2025
where (x) is Bessel function of the first kind, order If wi 1 Wj. w(+) w())) so their
distributions are identical. It follows that we can write
Ş( y) = (
1) 22 (w? + w9)
(19)
m mm f.
22 ET l + 2
k==
Again dropping multiplicative factors and terms unmodified by the coupling, we arrive at the RFF OT cost function CRFF specified in Eq. 6. This completes the derivation.
A.2 PROOF OF THM. 3.2
We now solve the OT problem formulated in Lemma 3.1 exactly in the special case that 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 Fxa()) the cumulative distribution function (CDF) of Xa= Consider m orthogonal frequencies with norms (21,W)). For RLFs, the OT problem in Eq. 5 is solved by the negative monotone coupling
Fxa(wi) + Fxa(w>)= 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, modifies arguments made for related problem by (among others) Thorpe (2019), to which 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 e R2 x Rd + R. we refer to a set I E R² as c-monotone if for all pairs (I1, 91), (222, y2) E T we have that
c(21,31 v1) + c(22, U2) S c(x1,82) + c(22,
(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 æ1 - 1/2 and - y1 instead of æ1 - Vi and 22 + y2. the plan would have done this instead. This formalised as follows.
Lemma A.5 (Support of optimal transport plan is c-monotone (Thorpe, 2019)). Consider n E P(R),
and assume that μ² E M(T) is the Kantorovich optimal transport for a continuous cost function elz, y). Then for all (I1, V1), (I2, y2) E supp(u") ) we have that
c(11,31) + c(22, v2) S o(x1, + c(22,
(22)
Proof. Thorpe (2019).
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
(-1)*=20 (wf + wB)
2 (wp w2)* 2II )
CRFF(WI, W2) =
CRLF(WY, W2)
(23)
22k kII $
0=
with = = ||0 2|12 and 0 = | + y 2|22 as usual.
First consider CRLF. We claim that for any I 22, V1. 32 2 0 satisfying Eq. 22 for this CRIF, 21 < 22 implies y/ 2 1/2 This seen by observing that
CRLF(Z1, v1) + CRLF(12, V2)
(RLF(F1, V2) (RLF(22,11)
2k
[<02
+ (w + ッのン
+ 2 -
(x + +
22 +
(24)
2k
Σ
[1 (
(1)))~1) ( - (x7)*)
22k
+
18

---

Published as a conference paper at ICLR 2025
Supposing that x1 < 22,, Eq. 22 is satisfied if and only if y1 2 y2, verifying the statement.
Denote by TRLE == supp(u") the support of the optimal transport plan for the cost function ORLF. and consider some point (Io. Mo) r. An immediate implication of the statement above is that
TRLF C {xx y) x I zo. y Z 3o U {( y) x 2 zo. y S yo
(25)
Let A = 10, ru] x [3o. 20), B = 2o) x 00, go). C = [To. 00) x [0, yol and D = (zo. so) x (Vo. oo]. Note that AUBUCUD = R+ R so u² (AUBLCUD) = since the measure is normalised. The subsets are also disjoint apart from AnC = (mo, Un). singleton of zero measure. Eq. 25 implies that μz (B) = 0 = 1 (D) whereupon 1= 1z (AUC) = μ² + +x((()) = 4 (AUB) + f BC. Now AUB
((0. zo x R (Hzo) x [0, yo)) and the set [xo] x (0, yo) is zero measure. Therefore, 1r" ""(AUB) = 1e" (1 rol x ( = Fxx (ro). Likewise, 1²(CU B) = (R+ x [0, yol) = Fxx (Vo).
Relabelling (Zo, vo) E T by (w1, w2), 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 μ² will instead be TRrp-monotone Unfortunately, Eq. 24 does not hold in general for CRFF, so we cannot immediately use the same arguments 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 μxm the negative monotone coupling. Consider coupling μ¹ E A2(')), (PMM), ie. any other feasible transport map which is not negative monotone. There exists some constant ,(1²) > O such that I(pNNs) < I(y() for all z < (where I(n) denotes the expectation of the cost function under p))
Proof. Recall that, for RFFs, the cost function is of the form
CREF(W1, w2) (-1)²²2 (w] w2)
(26)
kTT + ()
and that we would like to solve μ = arg we [E 2 2 4 4" 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 terms are trivial since
CFF (W1, W2]] In general. je
they do not contain 21,2 cross terms. We have
(-1)a22 2k
arg min (a))(a)
E ((w) + w2))) +
E (wwf + w2)
321 + ))
kll (k + )
(27)
We can solve the OT problem exactly for the first term, recovering negative monotone coupling. That is,
E (ww + wB}3)
arg min meA(()) 321 + 4))
= PNM
(28)
for any 2 > 0 and d finite. Note also that
mm (w) + w3) 2 (w)2)
(29)
because the coupling that maximises the expectation is positive monotone (this follows from our previous arguments almost automatically). This is evaluated as the kth moment of x² distribution. which is nothing other than 2 & + d) {)) This means we can upper bound the magnitude of the ktth term in the expansion by: 22k (BII(+)). So for any coupling (feasible transport plan) we can upper bound the magnitude of the sum on the right by: g(z) == Di _2k (I:T())). This goes to 0 as - 0 and increases monotonically.
Consider some coupling j' E Az)))), {/NM). Since the minimiser is unique (if jí' is not negative monotone, it is not a minimiser), there positive constant e such that
c = En (wf + wB)² Epnos (wf + w 2 > 0.
(30)
vo) Pe.g. since (zo, Uo) C {2o) IR+ and nt ((r0) x RR+) px(x mTo) = 0 since the marginal measure nonatomie.
19

---

Published as a conference paper at ICLR 2025
But we also have that
2- (w] + wg)
(-1) (wf + wB)
E
E
< <gg(2).
(31)
JNN
22 (A: + ())
22 + $
k=
Therefore, to guarantee that I(UNM) < I(u'), ), it is sufficient that
x
2g(z).
(32)
32 + 5)
Since g(=)/=² is also monotonically increasing and evaluates to 0 and (by inspecting its Taylor expansion), it is always possible to solve this to find 2 = 8(1²) such that, for 2 < d(")), I(UNM) < I(n') is guaranteed. This holds for any dimensionality d and any measure μ that is not negative monotone.
L
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.
only possible to prove optimality of (N when 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 at bigger = 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., 2023; Yu et al., 2016): a well-established and widely-use algorithm.
Comments on small 2 for RFFs We have seen that 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 frequencies, the sum Eijzs has d(d - 1) terms in total. Of these, 2 correspond are negative monotone norm couplings, and the remainder are independent. The independent terms are the same the pairwise norm-coupled and fully i.i.d configurations, 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 is small enough, so the result again follows.
A.4 PROOF OF THM. 3.5
We now drop the restriction that W1 1 We and consider the variance reduction problem for m frequencies whose respective direction is unconstrained. We will prove Thm. 3.5, which asserts that in this case the best possible coupling is antithetic, W1 -22
Proof of Thm. 3.5. Recalling the expression for RLF variance in Lemma A.1, the more general OT problem under consideration is
well 2T 4
2k W
1r arg min (EA
Ewi Ware
(33)
k=0
The term in square parentheses is expectation of an infinite sum, every term of which is greater than or equal to 0. The sum is manifestly minimised if ||w + wellz = O, which sets every term (apart from the first) to 0. This is achieved if and only if w1 wwy: valid coupling called 'antithetic sampling'. Any other joint distribution assigns nonzero probability to the event ||اا| w2ll2 > 0, so this optimal coupling is unique.
n
20

---

Published as a conference paper at ICLR 2025
Figure 5: Copula schematic for Random variables are drawn from a Gaussian distribution with correlation matrix Σ. They are pushed forward using Fxx then F-1 to obtain coupled variables with marginal Xa distributions. Σ is learned using gradient-basee optimisation, approximately solving the multi-marginal OT problem in Eq. 5.
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. is an open problem. In Sec. 3.2, briefly presented alternative numerical approach using copulas. Here, we discuss this in greater detail.
Copulas. A copula is multivariate cumulative distribution function (CDF) whose marginals are uniformly distributed on [0, 1]. By Sklar's theorem (Sklar, 1959), its joint distribution can be arbitrary. Given a copula, we can easily enforce the constraint that its marginals are Xa by pushing each coordinate forward with the inverse CDF F Fx.((), (.), whilst retaining necessary flexibility in the joint to reduce estimator variance. Fig. gives a schematic. Copulas can used to model dependencies between random variables and are popular tools in quantitative finance (Haugh, 2016).
Gaussian copulas. The general family copulas is still intractable to optimise and sample from so 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 {gi)m (O, Σ) where E E mmxx is correlation matrix, Le. a positive definite matrix with unit diagonals, and let Fx be the CDF of the standard univariate Gaussian. We say {m where u, = Fx((g)) is distributed according to Gaussian copula with covariance E. We use the notation ( ~ GC(2) to denote this.
Parameterising correlation matrices. Gaussian copulas are easy to sample from since they involve sampling multivariate Gaussian and applying the univariate Gaussian CDF We are therefore left with the task of finding appropriate correlation matrix E, for which turn numerical optimisation. The family m X m correlation matrices can be parameterised by a vector 0 E IR (m)(111//// In fact, there exist tractable bijections between unconstrained vectors of real numbers O mmmmm1)) and lower triangular Cholesky factors Lo such that Σ = Lg is a valid correlation matrix (Bhat and Mondal, 2021). In particular, suppose that for each 1, I...., N, and 1, ,,..., have Oi E R+, where 0g = 1. Then the parameterisation we use is
Os2 for i S. =
Lij
(34)
otherwise,
where s; - V sample from the associated Gaussian copula with (mm?) computational cost.
Ey- Note that, since we are directly parameterising the Cholesky factor, we can
Optimising correlation matrices. In order to pick an appropriate correlation matrix Σ, we optimise it directly to minimise the root mean squared error (RMSE) loss
(()) E [..)
Σ +1==1
(GRr(( (o:)))))(< k(x,, )2
(35)
A22
where {u:) ~ GC(L ). Note that ΦRF here depends on u, by pushing forward with F and using the result random frequency norms, though we have suppressed this dependence for notational
21

---

Published as a conference paper at ICLR 2025
RIF's
RLF's
train stops
oethogonal
(Jeunodd)
(fxod)
oethogonal
(Bood)
Figure 6: Training RMSE losses 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.
simplicity. Assuming that ORF differentiable with respect to wi; which the case in RFFs and RLFs,, optimise the copula parameters / by estimating the loss in Eq. 35, computing its gradients with respect to 0, and updating its values accordingly.
Training curves. Fig. 6 shows 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 an important avenue for future work.
B RFF AND RLF EXPERIMENTAL DETAILS
In this appendix, we supplement the discussion 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 downstream performance, We provide experimental details for the Performer (Choromanski et al., 2020) results in Table 2.
al., 2010), showing that in this case variance reduction does not help
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 e, an output scale variable o, and noise scale parameter on. takes the form
(-
1 (- x 2; 2;;13]
k(x;,,;) = of exp
Our baselines include standard methods for sampling random frequency vectors for within RFFs and RLFs: i.i.d. sampling and Halton sequences (Halton, 1960). 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 use ensembles of 2d vectors, including d orthogonal basis vectors and their d antiparallel vectors.
Selecting kernel hyperparameters. We want compare our coupling schemes using realistic kernel hyperparameter values, which we determine as follows. A realistic application setting for RFFs within GPs for probabilistic regression. Therefore, we first GP on a tractable subset of the data, specifically a maximum of 256 randomly chosen datapoints, to select appropriate parameters C, o and on. 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); e.g. Thm. 4 and App. F6 of Choromanski al. (2020) Thm. 4.3 of Likhosherstov et al. (2022), where the authors bound the Lg-norm queries and keys. This is anecdotally responsible for the deterioration in performance of Performers when networks become
22

---

Published as a conference paper at ICLR 2025
very deep. To reflect this fact and choose data regime where vanilla RLFs perform reasonably well (so we can assess any gains from our coupling), we set the lengthscale to two twice the average summed norm of the data, namely
Nipin
x; + 12,
(36)
train
KJ=1
over the training set. We train the rest of the kernel parameters (a, and (x) 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 training and 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 maximum of 256 points cach by sub-sampp 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, 2014), using learning rate of 10 The exact GP optimisation stage converges around 1000 steps, and 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 { (ÓRr (xx)) Group these in a large design matrix +
+
(2:)(
(37)
where (2;) of course depends on the ensemble frequencies {w.) (suppressed for notational compactness). For RFFs & E N x2m whereas for RLFs o IR Nxm, We estimate the Gram matrix by K TΦ
As noted by LázzaroGGreiilla et al. (2010), this RF kernel approximation is exactly equivalent to linear model, namely
y w + E,
(38)
where -N I) and (0, The prior covariance of this linear model is T + o which is, by construction, equal in expectation to the exact covariance produced by the kernel, namely K + of The predictive means of the approximate linear model and corresponding exact model are
Mapprox
D +
D,y.
(39)
02
Mexact K Kad(Kar (Kdd + -, -
(40)
whereas the predictive covariances are
Cappoox
$
p + I.
(41)
Cexect Kpp
Knd (Kdd + 0,,1)~1A
(42)
Here, 4, and t are the design matrices corresponding to the training inputs and prediction outputs respectively, Kad is the covariance matrix corresponding the training inputs, Kpp is the covariance matrix corresponding to the prediction inputs and Kpd and Kap are the cross-covariance matrices between the training and prediction datapoints. These models become exactly equivalent in the limit of infinite number of features since the kernel approximation becomes exact.
We have seen that pairwise norm coupling improves the approximation of the Gram matrices K. In particular, we 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
23

---

Published as a conference paper at ICLR 2025
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 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
Mapprox = (pa( + + Capprox = Bp - (Rad + o KKpp
(43)
where Kdd = D. Ld. Dd. K pd = - Rpd = $ Pd and Kyp = -b Dp. It is then straightfor- ward to compute the KL divergence between Gaussian distributions with means Mapproxs Mexact and covariances Cappoox. Cexact-
Surprisingly, 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 many trials this is not reflected in the data splits used for Table 3. This is also the case for predictive RMSEs (normalised and unnormalised by the i.i.d. result) and the predictive negative log likelihoods on held out test sets, also reported in Table 3.
routinely see variance reductions of 10-20%, but even with very
The reason for this experimental finding is that, as is clear in Eq. 43, the posterior is highly nonlinear in 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, Mappoox and C approx are in fact biased. We have achieved our objective of variance reduction of K with pairwise norm coupling, is both theoretically mandated and empirically observed. But clearly this does not necessitate variance (or bias) reduction for Mapрoox and Cappox- The relationship between the distribution of K and the distribution of the approximate posterior is more complex.
Variance reduction does not always help predictive performance. To sharpen this point, now take a single data split of the POWER dataset and plot the the kernel approximation RMSE (i.c. kernel estimator variance), well various quantities of predictive interest, against the number of features m. We exclude the Halton coupling (which consistently worse than 'orthogonal') for clarity of presentation. Fig. 1 (in the main text) 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 Table 3), (ii) the predictive RMSE on 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, 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 gives companion results to Table 3 but for = Sd (instead of d). In this regime, all the kernel estimators are closer to the groundtruth kernel value so the difference between their performance downstream tasks is smaller even between i.i.d. and orthogonal frequencies, which are separatdd when m = d in Table 3. Note that this is also behaviour is also clear from the upper range of number of features in Fig. where again the performances saturate.
24

---

Published as a conference paper at ICLR 2025
KL divergence (between approximate and exact predictive posteriors)
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOI
BOSTON
L.D.D
1.491 10.106
0.013
1a.00
1.570 10.n7
0.150
10.007
0.357 Loazs
2.128
HALTON
1,548 #0.104
0.014
+0.001
1.596
+0.419
0.138
+0.000
0.356
+0.024
2.263
=0.204
ORTHOGONAL
1,166 20.113
0.004
t0.00
1.635
#0.425
0.029
+0.002
0.235
+0.017
1.990
=0.191
PNC
1,168 #0.118
0.004
+0.000
1.589
#0.416
0.029
#0.002
0.235
#0.017
1.985
=0. 100
LAPLACE FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
LLD.
0.552
0.028
5.925
0.199
#0.011
0.230
0.759
HALTON
0.574 10.004
0.024
1a.0us
5.811
0.146
1o.cou
0.213
1o.u22
0.834
ORTHOGONAL.
0.486 20.009
0.014
10.002
5.494
11.77
0.059
1. 005
0.165 10.018
0.679 20
+ PNC + ANTITHETIC
0,482 $0.006
0.014
ta.02
5.468 11.780
0.050
+0.00m
0.165 +0.01
0.673
20.008
Predictive RMSE (normalised)
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
L.I.D
1.000 20.077
1.000
20.033
1.000 10.179
1.000
20.016
1.000
20.025
1.000 =0 007
HALTON
1.013 20.07
1.001 to.ue
1.005
10.77
0.991
a0.0u
0.999
10.025
1.021
20.008
ORTHOGONAL
0,917 #0.006
0.994
+0.044
1.019
#0.182
0.915
#0.019
0,930
#0.025
0.974
006
PNC
0.917 10.006
0.994 10.004
1.033
10.180
0.915
10.010
0.930
10.025
0.975
LAPLACE FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
L.L..
1.000 #0.028
1.000
+0.043
1.000 +0.192
1,000
+0.017
1.000
+0.029
1.000
=0.073
HALTON
1.007 20.008
0.997
20.044
1.003
10.100
0.964
20.017
0.989
20.029
1.019
=0 0m
ORTHOGONAL
0.979 20002
0.990
10.0e
1.003
10.10T
0.899
0.958
10.02s
0.979
20.07
PNC ANTITHETIC
0.981 10.027
0.991
10.004
1.006
10.187
0.909
10.018
0.958
10.02s
0.982
Predictive RMSE (umorrmalised)
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
11.333 10.304
2.587 10.112
64.259
111
4.840 10.078
3.468
10.0ss
4.976
HALTON
11.479
2.580
64.562
+11
4.799 +0.0m7
3.464
+0.088
5.081
ORTHOGONAL
10.398 40.292
2.572
+0.114
65.477
4.430
#0.008
3,225
+t.0ss
4.847
PNC
10.398 #0.261
2.572 #0.115
66.404
#11
4.429
#0.008
3,225
#0.055
4.852
LAPLACE FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
L.D.D
10.098 #0.284
2.587 10.112
85.772 #16 446
5.025 =:0 cea
3.243 =:0 084
4.224 310
HALTON
10.173
2.580
86,060
4.845 +0.088
3.209
4.303
ORTHOGONAL
9.882
2.562 10:115
86.042
4.518
3.108
4.136
PNC ANTHITHETIC
9.908 20.270
2.563
86,258
4.569 #0.com
3,108
+0.092
4.150 =0.
Predictive NLL
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
I.I.D.
4.609 20
2.388
20.002
7.077
10.036
3.037
20.030
2.764 10.033
4.808
407
HALTON
4.666 10.108
2.389
10.022
7.105
10a2s
3.025
10.020
2.762
20.45
4.9-2 20.440
ORTHOGONAL
4.301 #0.093
2.381
#0.042
7,381
+1.122
2.913
+0.022
2.640
#0.041
4.684
=0 432
PNC
4.304 #0.065
2.381
+0.042
7.543 +1.080
2.913
+0.022
2.639
+0.041
4.60 4.693 =0. 433
LAPLACE FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
I.L.D
3.930 #0.059
2.380
+0.042
12.981
#3.953
3.076
#0.023
2.644
+0.041
3,384
=0.296
HALTON
3.948 20.000
2.387
20.002
12.929
23 927
3.024
20.092
2.628
40.040
3.453
ass
ORTHOGONAL
3.880 20.056
2.379
12.722
13.
2.935
1u.022
2.581
1.0.us
3.312 20.20
PNC + ANTITHETIC
3.885 1oo05
2.380
10.092
12.729
133.94
2.948 10.022
2.582
1o.00s
3.319
L. 220
Table 3: Downstream performance on held out test sets when
d. Includes KL divergence between exact and approximate GP predictive posterions; predictive RMSEs (normalised to RMSE of I.I.D. and unnormalised); and predictive negative log likelihoods of the approximate GP predictive posteriors. Reported eиrоrs are equal to two standard errors, i.c. 98% confidence intervals, computed by averaging across splits. Note that differences berween data splits are responsible for large errors; within each split we take enough trials that the standard errors become small. Substantially lower kernel estimator variance does not general guarantee better predictive performance, though performance tends to be no worse. While = d features is not enough for good predictive performance absolute terms, exaggerates the difference between couplings and makes for easier comparison (see e.g. Fig. l). We average over 20 data splits.
25

---

Published as a conference paper at ICLR 2025
Kernel estimator RMSE (unnormalised)
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
L.D.D
0.527 #0.00
0.088
=0.012
0.264 20.035
0.137
#0.014
0.301 #0.018
0.107 =0.000
HALTON
0.550 10.001
0.089
10.015
0.262
20.038
0.123
10.014
0.281
10.019
0.126
10.011
ORTHOGONAL
0.340 10.014
0.061
20.007
0.160
20.033
0.094
10.010
0.181
10.012
0.074
10.005
+PNC
0.304 +0.013
0.055
+0.008
0.141
+0.030
0.079
+0.007
0.154
+0.011
0.071
+0.005
LAPLACE FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
LLD.
1.312 +0.118
0.382
#0.169
0.645
#0.212
0.376
#0.044
0.629
#0.062
0.392
=0.067
HALTON
1.336 20.380
0.297
20.047
0.453
20.076
0.249
20.030
0.539
2. 084
0.357 20.031
ORTHOGONAL
1.096 10.108
0.285
20.071
0.433
20.121
0.159
#0.015
0.466
20.002
0.433 10.166
+PNC ANTITHETIC
1.123 10.148
0.227 20.038
0.391 20.002
0.113 10.012
0.461
10.007
0.338 20.052
KL divergence (between upproximate and exact predictive posteriors)
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
LLD.
0.122 20.010
0.003
0.046
0.003
0.017
0.139 =0.030
HALTON
0.122 10.010
0.503
20.95
0.045
10.12
0.003
10.000
0.017 20.004
0.150 20.001
ORTHOGONAL
0.117 #0.000
0.002
=0.001
0.042
#0.011
0.003 #0.000
0.016
#0.004
0.133 =0.029
+PNC
0,115 +0.008
0.002 +0.001
0.041
#0.011
0.003 +0.000
0.015
+0.004
0.135 =0.029
LAPLACE FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
LLDD.
0,234 #0.058
0.067
=0.050
1.827
#1.302
0.003
#0.001
0.051
#0.023
0.742
=0.319
HALTON
0.231 10.048
0.066
20.050
1.992
12.002
0.003
10.000
0.047
10.003
0.754 10.321
ORTHOGONAL
0.227 20.049
0.066
20.050
1.889
21.839
0.003
20.001
0.048
20.023
0.713
20.310
+PNC ANTITHETIC
0.232 10.003
0.067
20.052
1.846 117770
0.003
10.001
0.051
10.024
0.713
20.315
Predictive RMSE (normalised)
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
L.D.D
1.000 #0.057
1.000
=0.073
1.000
#0.438
1,000
#0.051
1.000
#0.054
1.000 =0.174
HALTON
0.997 20.004
1.001
20.075
0.992
29.429
1.001
$0.051
0.999
20.056
0.996
20.173
ORTHOGONAL
0.996 #0.056
1.001
=0.078
1.004
#0.443
1.000
#0.001
0.990
#0.056
0.999
=0.175
+PNC
1.001 +0.056
1.001
=0.078
1.008
#0.445
1.000
#0.001
0.997
#0.056
0.994 =0.171
LAPLACE FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
LLD.,
1.000 +0.00 40.000
1.000
=0.050
1.000
#0.487
1.000
#0.061
1.000
#0.061
1.000 =0.100
HALTON
0.992 10.000
1.001
200081
1.020
20.084
1.000
10.051
0.996
10.000
1.004 10.201
ORTHOGONAL
0.993 +0.050
1.000
1.031
+0.507
1.000
+0.051
0.995
+0.060
0.995
+0.108
+PNC ANTITHETIC
0.996 +0.082
1.002
2:
1.022 20.002
1.000
#0.051
0.999
#0.061
0.992 20.108
Predictive RMSE (unnormalised)
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
L.D.
8.016 +0.458
2,595
#0.189
56.066
524.574
4.321
+0.221
2.689 #0.146
3.954
=0.089
HALTON
7,996 42.
2.597
#0.190
55.601
#24.094
4.326
+0.21
2.687
+0.150
3,938
=0.08s
ORTHOGONAL
7.987 20, 440
2.597
10.100
56.312
#24 s40
4.321
#0.223
2.686
10.350
3.948
20.000
+PNC
8.026 10.446
2.597 18.189
56.506
131
4.323
19.22
2.682
10.150
3.928
10.877
LAPLACE FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
LLD.
8,244
2.710 20.218
81.045
4.319
1:.222
2.773
17.170
4.339
HALTON
8.175
2.713 10.219
82.627
4.320
10.222
2.763
10.168
4.356
20.8711
ORTHOGONAL
8,183 +0.411
2.710
#0.224
83,571
+41.127
4.320
#0.220
2.761
#0.166
4.319 =0.800
+PNC ANTITHETIC
8,212 +0.428
2.715
#0.225
82.838 +39.877
4.320
+0.219
2.770
#0.160
4.305 =0.857
Predictive NLL
FOURIER FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
LLD.
3.467 +0.085
2.375
=0.078
5,021
2,888 +0.057
2.407 #0.048
2.793 =0.258
HALTON
3 10.084
2.376
20.073
5.021
2.888 10.057
2.408
10.018
2.797
ORTHOGONAL
3.459
2.375
5.023
2.888 20.057
2.406
20.018
2.787
+PNC
3.462 10.00s
2.375 20.075
5.018
20.202
2,888
10.057
2.407 20.048
2.784 20.202
LAPLACE FEATURES
CONCRETE
ABALONE
CPU
POWER
AIRFOIL
BOSTON
L.I.D.
3.541 #0.113
2.438
#0.104
9.694
#7.105
2,888
#00.57
2.442
#0.050
3.693
=0.ars
HALTON
3.542 10.111
2.437
20.104
9.722
27.220
2,888
100.57
2.438
10.0ss
3.708
20.009
ORTHOGONAL
3,544 #0.112
2.436
=0.102
9,836
#7.410
2,887
#0.007
2,440
#0.050
3.637 =0.047
+PNC ANTITHETIC
3.543 20:114
2.435
10.102
9.686 17.370
2.888
10.057
2.437
20.058
3.648 20.907
Table 4: Equivalent results for m 8d. Companion results to Table 3 taking m Sd features, so
the quality of kernel approximation and performance on downstream tasks are both much better. In all cases, the approximate kernel converges the true kernel so becomes more difficult distinguish the different algorithms. We still observe that our methods provide substantial variance reduction (see top row).
26

---

Published as a conference paper at ICLR 2025
B.3 NORM-COUPLED RLFs IN PERFORMERS
In Sec. 3.3, we used random Laplace features for attention approximation in Performers (Choromanski al., 2020), type of efficient transformer that estimates softmax using low rank decomposition. Here, we discuss the results in greater detail.
Derivation of Eq. 8. Let Xıy denote the random variable k( xy), which is an unbiased estimate of exp(x; 24)) constructed using features {exp(] 22,11/22)@LF((2 H Unbiasedness follows trivially from the fact that RL.Fs give an unbiased estimate the Gaussian kernel. Normalising by the sum of attention scores, let us define قij == Xij/ Σ, Xij. Let Xij = Pij + ōig with pij == RB(x,)) and
B(6,y) = 0. Then
a (mij + 244jdbi +,,5) Σ Mik +
Σ dak)=
k
(44)
pīj 24ssdis + +62
(1-2
Ex I-
dsk Eu dakda
1 Na )²
+
O
22
Njs
N
where we defined p4 are targeting Naa:"
2/2 E, Mig, the average groundtruth attention across tokens. Since
Psi
MSE( NZ
83
Apiydes
Σ
8ek +
3p N2
Σ Eu
+ O
1 N73
(45)
Nps
As in the main text, denote ()) = Σ, MSE(agy). the average mean squared error over the
tokens to which i attends. To better see the intuitive behaviour, suppose also that (4j = jās V j. Then
MSB(a.)=
Σ
dejosk
+0
N2 2
N
N2
N3
jk
(46)
N33 (N²))
Var (<,;; x3))
Σ J.k
Cov(K((z,, 2.)) ((,,, #x)) +
= N2p,
N
N²
as reported in Eq. 8.
Positive monotone coupling. In contrast to Eq. 7, we say that random variables (WI, w2) are postive monotone-coupled if w1 = wg almost surely. This is valid transport plan, the identity. Clearly, all m frequencies in an ensemble can be simultaneously positive monotone-coupled by making them equal. 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 CRLF - CRLF 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 atfention approximation error. Modifying the norm coupling between frequencies changes both Var(k(#;, æ.,) and Cov(k(æ;,,0) æ.)).. MSH(a.) change unpredictably. Fig. 7 shows the results for randomly N = 16 synthetic, normally distributed 1 6-ddimensional keys, (0, la). Strikingly, maximising the kernel estimator variance with positive monotone (PM) coupling ends 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 paradigm.
Performer experimental details. To obtain the results in Table 2, we train a Performer ViT (Dosovitskiy et al., 2020; Choromanski et al., 2020) on the ImageNet (IM) dataset (Deng et al., 2009). We estimate attention using m 128 random features that are orthogonal with norms that are: (i) i.i.d. (ii) PNC (Def. 3.3), (iii) positive monotone-coupe 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 x 16 patches, and train with the Adam optimiser for 90 epochs with compound learning rate (10* steps linear warmup, constant, then cosine decay, with base LR 3 x 10-3 and final LR x 10 5). The batch size is 4096. To get the standard errors, 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.
27

---

Published as a conference paper at ICLR 2025
10
kernel covariance
16"1
10
kornel variianee
atlention MSE
10
10
to
num features (i)
i..dd
orthogonal
PNC
PM
Figure 7: Kernel variance Var(k( æ2)), kernel covariance Cov(k( æz)) and attention mean squared error MSE(a,) plotted against number of frequencies for different coupling schemes, As intended, PNC decreases the kernel estimator variance, but it also decreases the kernel estimator co- variance so the attention MSE remains essentially unchanged. On the other hand, positive monotone coupling (PM) drastically increases both kernel variance (which maximises) and kernel covariance. The latter offsets the former (Eq. 46), that attention MSE actually drops: remarkable, counterin- tuitive finding.
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 o-coupling.
C.1 CONSTRUCTING GRAPH RANDOM FEATURES
For the reader's convenience, we begin by providing brief introduction to GRFs.
Graph node kernels. Recall that graph node kernels k N x + IR are positive definite, symmetric functions defined on pairs of nodes 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 weighted adjacency matrices (Smola and Kondor, 2003b; Chapelle et al., 2002). In particular, they are often functions of the graph Laplacian matrix,
L = D W.
(47)
where W is a (weighted) adjacency matrix and D is a diagonal degree matrix (satisfying Da Σ, W W.)). It also common consider its normalised variant,
L = D LD
(48)
whose spectrum is contained [0,2] (Chung, 1997). We provide examples in Table 5.
Name
Form
d-regularised Laplacian p-step random walk
(IN o²L)-d
(aIN LL", 2
Diffusion
exp(-o
Inverse Cosine
cos (LT/
Table 5: Examples of graph node kernels. The exp and cos mappings are defined via Taylor series expansions rather than element-wing e.g. exp(M) = lim (IN + M/n)" and cos(M) = Re(exp(iM)). and are regularisers.
Formula for graph random features. Computing graph node kernels exactly is computation- ally expensive because of the ((N3) cost of e.g. matrix inversion exponentiation. Inspired by
28

---

Published as a conference paper at ICLR 2025
the success of random Fourier features in the Euclidean domain (Rahimi and Recht, 2007), the class of graph random features features (GRFs) permits unbiased approximation of &s in subquadratic time (Choromanski, 2023; Reid et al., 2024b). Intuitively, GRFs interpret the powers of weighted adjacency matrices in the Taylor expansions of the expressions in Table as weighted sums over walks 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, GREs compute set of estimators (dGRF(U:)) such that k(vi, Uj) B((cRr(V:) (CRF(1))] by taking:?
ΦGRF(U:) (2)) = m v( (2))
(49)
where (0 the kth walk (out of a total of m) simulated out of the starting node uy We remind the reader that by 'walk' we mean a sequence of neighbouring graph nodes. The projection function (b). Q + RN maps from the set of graph random walks ?= (ma)/= | E N,, (0t, 2441) EELE to sparse N N-dimensional feature vector. It is com- puted as follows:
(wwrg)) ) ((ee((wr))) I(win w!!), q= I......,N
(w(1))
(50)
p(wigg
Wiy ERty
In this expression,
(a)(1), is the qth component of the feature projection of the random walk (0) (itself sequence of neighbouring graph nodes beginning with oz).
Sig is the set of all graph walks between nodes v. and lip. of which each Wig is member.
I(-) is the indicator function that evaluates to if its argument is true and 0 otherwise.
Win E w+( is the condition that Wig. particular walk between nodes vs and Up. is prefix subwalk of (1), the random walk we actually sample.* 4
f: N R. is the modulation function, which controls the particular graph node kernel we choose
to approximate.
len(wig) is the length of graph walk wig.
p(wig) is the marginal probability of sampling random walk Wig. which is known from the sampling strategy.
a(wyg) is function that returns the products of weights of the edges traversed by Wig-
Intuitively, to construct (+)( using random walk w()), we look the node that it visits each of the (en ()) + 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) 'modulation function' ʃ applied the number of steps taken so far. We refer the reader the original works of Choromanski (2023) and Reid et al. (2024b) 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., 2024c): a previously-proposed variance reduction algorithm that couples the lengths of pairs of random walks.
Consider a random walker that terminates with probability at every timestep. In practice, this can be achieved by drawing a 'termination random variable' t from a uniform distribution on [0, 1], t ~ U(10, 1)). and ending the walk if < p. Now consider two such walkers with corresponding termination random variables t1, 12. If they independent, and tz drawn independently. Antithetic termination instead proposes to induce a nontrivial joint distribution by first drawing
'Strictly, for unbiased estimation of diagonal kernel entries {k(v,, e,,), we should construct two independent sets of features, but this is found to make little practical difference (Reid et al., 2024h) so omit further discussion in this manuscript.
'Meaning that the walk (6) from node initially follows wig. then optionally continues to visit further nodes. Note the subtle difference in usage of the symbol E compared to in the expression Wig pig. where it mcans 'a member of the set of walks' rather than 'a prefix subwalk'.
29

---

Published as a conference paper at ICLR 2025
~0 1J) as usual, and then setting 22 mod 1(1 2). t2 still follows a uniform marginal distribution so we preserve unbiasedness, but this coupling forces walkers to terminate at different timesteps (if Phal < 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.
is also possible to improve the accuracy of estimators using graph random walks by coupling walker directions (Reid et al., 2024a; Luo, 2019). 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 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:
minimise E (1 )-)))) 1 yaflym
Edirs
T (v(W()) ((wS)))
v(w!
(51)
for μa E (2
where (2())) the set of joint distributions on N22 with geometrically distributed marginals, ( is the length of walk (3) out of node e4, Edis averages the directions of walkers of particular length, and (((w))) 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, will be convenient to introduce the direction-averaged featur feature vector v0.) N RN satisfying
(((()) E (v (O)).
(52)
This computes the average feature projection over all walks of a given length (00) out of node U/- 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
(509)+ (Mal))+5(()))
minimisc (150)) (11 of yuu
for p E A2(1e). (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 o-couplings. Unfortunately, Eq. 53 remains intractable. To make progress, must restrict search-space of joint distributions from (T1G) the set of all joint distributions with geometrically- distributed marginals - to tractable subclass. Taking inspiration from approaches in numerical OT, we have considered the class of permutation densities defined by P= y) == (ll(([==])== [ny]+ with a a permutation of order n (that is, a bijection o [m] + [m]), 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 matching problem which, under further simplifying assumptions,
frames the minimisation 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., 2024c).
We now replace (2(1)) in Eq. 53 by the set of measures (ulo))oss. obtained by pushing permutation densities through the required inverse CDF. Sn is the set of permutations [n] [m]. Denote by 3(-): [m] - N the function
() = E (
(5 (u)
(54)
30

---

Published as a conference paper at ICLR 2025
This takes the direction-average feature vector for node (see Eq. 52) and evaluates further average over walk lengths corresponding to the qth quantile of the geometric distribution v/.- Again, this is straightforward to estimate by simulating walks. With further approximation of the objective, we can consider
IILiI PESA
Σ 11 [2
Σ
(
(Dlvl") (000)))
(55)
91
which is 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., 1987; Burkard al., 1998). 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, discuss possible approaches directly solve the matching problem in Eq. 55. In particular, we prove that can be solved with high probability at time complexity independent of the number of nodes N, and give a polynomial time algorithm for the special case of diagonal kernel entries (e u;). Notwithstanding this progress, as a simpler first recourse we can just make a further approximation by restricting the sum to qi q2 diagonal terms. This is the approximation used to get the results in the main text. Doing so, arrive
arg min
E oE[n]
(o(0 (())))
(56)
This is now a minimum- bipartite matching problem with linear weights. Even though the search space of possible permutations of order n is 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, 1955)) in time O(n"). As stated in Sec. 4.1 of the main text, we can use this optimal permutation to construct o-coupled GRFs.
This approximation might seem unreasonable since it involves discarding all the off-diagonal qi * ga terms from the objective, but for modest permutation order 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 es- timates of the 2-regularised Lapla- cian kernel on the eurosis graph, comparing the approximate and ex- act objectives. The former can be computed efficiently but the latter is evaluated at cost exponential in n, so we limit to n = 10 and smaller. We plot the kernel approximation er- ror (normalised by the i.id. result) for different permutation orders. We also plot the Cayley distance between the permutations, which measures the difference between them. As grows, the permutations become more different so we are clearly not finding the same coupling. Nonethe- less, we do not see a statistically sig- nificant difference variance reduction. This suggests that our linear proxy objective is rea- sonable. We defer more detailed analytic study of this curious phenomenon future work.
Approximating the quadratic objective
dietance box and approx Copt
relative Flob. norm crror
0.0666
i..d
r-coupled (approx)
r-coupled (exact)
0.0660
Cayley distance
0.0654
0.0645
10
permutation order, 7
Figure 8: Relative Frobenius norm error (normalised by i.i.d. result) with permutations of order n learned using ex- act and approximate objectives. 'Cayley distance' shows the difference between the learned permutations. Although drop- ping off-diagonal terms from the objective in Eq. 55 changes the a-coupling, it does not significantly detriment variance reduction.
the amount of
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 computationally efficient
31

---

Published as a conference paper at ICLR 2025
coupling that offers much greater variance reduction compared prrristialll-mooivatte techniques. Analytic intractability is common in OT (Sec. being an obvious exception), but fortunately this framework also equips with 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 qi = 22 approximation in Eq. 56, in this section we discuss progress towards solving the optimisation problem in Eq. 55 exactly. In partieular, prove Thm. 4.1: that can be solved with high probability and
1. at time complexity independent from the dimensionality of vectors ("(g())))) (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 (1)) )==1)
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 (1 (3 q- 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 given c > 0 and vector
Rd, define the set N.(v) {w E IR |y 7 I
<), where av,w w. is an angle between v and w. Then define the property of c-separation as follows.
Definition C.1. For a given e > O, a set of vectors V and a vector v E V, we say that is e-separated in V if
Nv n U Nx(x)= =0.
(57)
xEV((v)
Simplifying notation, we can rewrite Eq. 55 as the following optimisation problem: (1)
o- argmin
Σ
[(v!9 +v volg.)) [v
(1) T + +oo(g=)]
(58)
(
where a is the permutation of the sequence (1, (I. n) and V. (b,f) = (t:)) With every permuta-
tion o of the sequence (L, (L., n),, we will associate the perfect matching (o) in its corresponding bipartite graph with monochromatic classes of size each. Denote
Σ u
(59)
M(o)
cEM(O)
where u. is given by
(8) (k))
vec [(v]]
(1) o(k) @
(v(0)
(0)
(60)
Here, represents the outer product and vec is the 'vectorising" operation that flattens its input to vector. Note that that (k))) E R N'"
It is straightforward to see that, in this notation, the optimisation problem becomes
(() ((()
argmin M(a) M(a)
(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 u via the Lindesss-s Transform (JLT) (Dasgupta and Gupta, 2003). Specifically, compute
(i.)
i (i.j)
Gu (i-)
(kko(k)) =
(k,o()))
(62)
32

---

Published as a conference paper at ICLR 2025
where N denotes the number of random projections and G E R is a Gaussian matrix with entries taken independently at random from from.N(O, (O,, 1). From the well-known properties of JLT (Dasgupta al., 2010), we have that
( (u (kk))()))) (k) T
E (1))) ) (ko)()))
(63)
since it unbiased. We also have that
(k)) + ())
(ko())) 11+c,
--c c
(64)
(
(k,(RR))
+ (R,0(19)22
with r = O log(m)) random projections, with probability p = - neg(n). Here, neg denotes the negligible function of
Eq. 64 describes the ability of the JLT to (approximately) preserve Lz-norms. The analogous property for preservation of dot- not-products also holds, which of interest to given Eq. 61. It follows because for any x,y E RN
y
yl1?
x y
(65)
4
The properties above lead us directly to an algorithm to decouple the time complexity of the optimi- sation from N. One simply replaces every u (k)) with its corresponding dimensionnllit-redueee (((j))
(k)) for r O log(n) and >> 0 small enough. Then, by the concentration result listed above and the union bound over all pairs (7, j)) we conclude that the optimal a (potentially not unique) will remain optimal after applying the JLT (since all dot- not-products are approximately preserved), This completes the analysis. We stress that does not depend directly on the dimensionality u 0.)) (k,o(k)))
and therefore on N.
Polynomial in n algorithm for i i j. We can make further progress in the special case that = J. Le. if are minimising the variance of diagonal kernel entries k(u; u,). Denote by T(n) time complexity of an efficient polynomial-time algorithm for solving the minimum weights matching problem bipartite graphs with monochromatic parts of size n each. The following is true.
Lemma C.2. If Eq. 58 has unique solution a² and Surthermore (o is e-separated in F (fM(() o E S,) for some 0 < <<< (where S, denotes the set of all permutations of the sequence (I,-..g n)), then, for any k > 0, a² can be found in time k T(n) with probability 1 (1 -p.) Here, Pe is given by the following formula:
π
P. P P
arccos
S
(66)
ot + + 22
2
with g E R a rundom isotropic vector.
Proof of Lemma C.2. if = j. the expression in Eq. 61 is simply ||f Thus, the goal is find permutation that minimises the L2-norm of the vector fM(() We propose the following algorithm. At every iteration, replace each vector u, with the projection u, u g for g E N I=).. Note that, from e-separability, we have that if g E Nv)),, then the permutation corresponding to the minimum weight matching in the bipartite graph with weights given by ue is also equal to o². 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 g is taken from the isotropic distribution, we know that the probability that g E (v(v)) is exactly Pe, where p. is 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 Vectors 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 o* with probability 1 - (1 - p²) which completes the proof.
Having considered both claims, the proof of Thm. 4.1 is complete.
33

---

Published as a conference paper at ICLR 2025
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 Po po(2;, y) nl [lo[[nz] ([[[z]))((x) 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 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 A = {ula) specified by permutations of order n, given by permutation densities Pa pushed forward coordinate-wise using the (leftccontnuous) inverse CDF F. Suppose also that the marginal measure 7 is absolutely continuous with respect to Lebesgue measure. Consider also continuous function f + R whose expectation is to be estimated using m = 2 coupled samples. In the limit + oo, the class A contains measures that converge weakly the optimal transport plan that is, the sample coupling which provides the smallest possible estimator variance.
Proof. Consider the set {u, == (i 1) for some N. Consider also a random variable X with measure y P(R) and corresponding distribution function F. Define its left-continuous inverse F-1 (y) = inf{x E R: F(x) 2 v). Given some permutation order n E N, consider the set
z; F = (-)
(67)
which we will refer to as the quantile midpoints. Then consider the discrete measure
1 71
n da
Th
(68)
where 8z, (A) = = x4 E A and O otherwise (with A C IR). Denote by F, the distribution associated with 7/2- Clearly, 7/ -> weakly as n + oo since Fx(x)) + F(x) for any continuity point of F
(Billingsley, 2013). In this sense, the measure Th a discrete approximation to 77-
Now consider the variance reduction OT problem for estimating EE [[w]] with m = 2 samples,
= arg min MEAz(n)
I ( [w2)]
(69)
Solving this analytically for arbitrary f is not in general tractable. On the other hand, if we instead take the discretised marginals Th., we must solve
M arg min (2) (w])
(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 Villani (2021), which the interested reader is directed for full details.
OT with equal-mass discrete marginal measures. Any measure A2(1). can be represented as a bistochastic x matrix B [Bis]]]= E Bn, meaning that all Bıj are nonnegative and satisfy
Σ BB = i: Σ 1
(71)
Therefore, the Kantorovich problem reduces
inf BEB,
Σ Ci 35
(72)
where we have encoded the transport costs in the cost matrix C = (=x By is
interpreted as the mass of the singleton (I,, 2,). This is a linear minimisation problem on a convex bounded set. It is well known that 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 the extremal points of B,, (points that cannot be written nontrivial linear combination
34

---

Published as a conference paper at ICLR 2025
of other points in B,). By Birkhoff's theorem these are exactly the permutation matrices B [6 with o E Sn. 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 (2021) for further details. Choosing just one of these optimal permutations a (since any convex combination of the corresponding measures will give the same amount of variance reduction), we have that
1 = 10
6 (c))
(73)
Stability OT plans. Another important result in optimal transport is the stability of transference plans: namely, that if 7/ - yf weakly then the optimal coupling an E A2(mh) also converges to the optimal x E (2))) weakly provided the cost function e(X,, X₃) (in our case fXX)) is continuous and
/ c(a1, 22) (2r, v) < co.
lim sup 2=+00
limsup/
(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 (2020) for proof sketch and Thm. 5.20 by Villani et al. (2009) for more detailed exposition. follows that Pn converges weakly to the true optimal coupling μ²"
The OT plan is (asymptotically) in our search class. Lastly, we have that our search class of couplings A(a) (corresponding to permutation densities pushed forward by F) contains measures when n + oo. In this mathematical sense, the asymptotic limit the class of couplings amongst which 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.
that converge in distribution to (6) 1
To see this, consider the measure on [0. 1]2 deseribed by the permutation density Po y) [la[[[=])=<<m9] Consider also the discrete measure on [0. 1]2 given by Ei= da, with the set {u))= the quantile midpoints defined previously. These measures converge in distribution (their corresponding joint CDFs can differ by at most at any point, which goes to as n - oo). They will also converge in distribution when pushed forward coordinate-wise by F,- by continuity of v,. It follows that asymptotic p.( and plo) converge in distribution, which completes the proof. O
We remark that not all the assumptions in Thm. D.1 hold in the GRF setting. In particular, neither the marginal measure 7) nor the function is continuous. The intention of including Thm. D. is to build intuition for the reader and provide some motivation for our choice of permutation densities. We defer rigorous investigation into relaxing these restrictive assumptions a tough measure-theoretic problem to future work.
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., 1998). This popular measure of the relative importance the nodes A of a graph is expensive to compute exactly so is often estimated by sampling random walks (Fogaras et al., 2005). 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. SETUP: ESTIMATING PAGERANK WITH RANDOM WALKS
The PageRank vector the stationary distribution of Markov chain whose state space is the set of all graph nodes N, with a transition matrix
P = (1 Phatl +
Phalt
E.
(75)
N
35

---

Published as a conference paper at ICLR 2025
Here, Phalt (0. 1) is scalar, N is the number of nodes and E [1]a is matrix whose entries are all ones. P is the transiton matrix simple random walk,
Ps if (i,)) otherwise
(76)
with de the degree of the ith node. Since P is stochastic, aperiodic and irreducible, we define the unique PageRank vector p IR::
P pT, p11=1,
(77)
where normalised the sum of vector entries to 1. Computing P analytically is expensive. Rear- ranging and Taylor expanding (1 - Phalt) P) p1) it is straightforward to see that
8
P:
Phalt
Σ
(1
Phalt) *P
(78)
N
jEN
This is simply sum over all walks from each of the graph nodes v to node vt, weighted by their respective probabilities that is, the expected number of random walkers ending at node e/ if they terminate with probability Phalt at every timestep which invites the estimator proposed by Fogaras et al. (2005),
EE 1 Iath walk from node vy terminates at node uz].
1
P Nm m=
(79)
An interesting and practically important question is whether the variance of the estimator p can be improved by coupling random walks. As with GRFs, this can be achieved using antithetic termination (Reid et al., 2024c) (see Sec. C.2). However, we will see that our OT length coupling approach does equally well or better.
F.2 OT FORMULATION OF PAGERANK VARIANCE REDUCTION
Given m = 2 walks, the variance of the PageRank estimator Var(p:) depends on the quantity
B
Σ
Σ
L|r,th walk from node vy terminates at node ex]
2
UJ. m1,92=l
(80)
1n2th walk from node ex terminates at node wx]
After few lines of algebra, the OT problem minimise this quantity is
minimise E d2)
Σ PRW of length 4 from node terminates at node uy]
By EN
(81)
PRR of length 12 from node by terminates at node v.) for μ E A2(11)
without making any approximations. P(-) denotes the probability of the event brackets. As with GRFs, for tractability we restrict our search space of joint distributions to permutation densities pushed forward coordinate-wiss with the geometric distribution inverse CDF, and this becomes
o² = arg min
ΣΣ
P[RW with length in gth quadrant from node vy terminates at node ei]
By EV VE[M
(82)
PRR with length in ar (g)th quadrant from node 2, terminates at node v.).
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 usual be solved efficiently with the Hungarian algorithm (Kuhn, 1955).
Note that we made fewer simplifications than for GRFs (Sec. C.I). They only requirements for tractability are restricting the class of considered joints o-couplings and computing MC estimates of the terms in the OT cost matrix. We do not need to e.g. move Egire inside a square or approximate quadratic matching problem by liights counterpart.
36

---

Published as a conference paper at ICLR 2025
karate
dolphins
polbooks
football
ncwsgroup
0.4
0.3
0.3
0.4
0.3
0.2
0
0.1
0.25
0-775
0. 25
50 0.75
0.25 0.50 0:75
0.25 0.50 0.75
0.25 0
50
0.70
eu-core
arenas-email
eurosis
citeseer
cora
RMSE
0.3
0.2
0.2
0.3
0.2
0.26
0.50
0.25
termination probability
Lid.
antithetic
a-coupled
Figure 9: PageRank estimator errors p مl1 for range of real-world graphs when random walkers are taken to be i) i.i.d., ii) antithetic or iii) o-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.
E.3 EMPIRICAL RESULTS FOR PAGERANK
We now provide experiments using our OT coupling improve the convergence of PageRank estimates. Fig. 9 shows the results. For termination probabilities Phalt E (0.1, 0.3, 0.4, 0.5, 0.6. 0.7, 0.8, 0.9), we compute the optimal permutation by solving the matching problem in Eq. 82 on an Erdós-Rènyi graph with N = 100 nodes, taking permutation order 10 as a hyperparameter. We then these couplings on range of real-world graphs (Ivashkin, 2023) of different sizes, plotting the estimator error || ø (where is the true PageRank vector and Φ its MC estimate) against the termination probability Phalt We include walkers that are i.i.d., antithetic (Reid et al., 2024c) and o-coupled. At every value of Phalts method performs at least as well as the baselines, and in some cases (c.g. on the very large cora a graph) it gives much greater variance reduction.
Note that we only consider connected subgraphs and take 1000 repeats for standard errors. Read- ing 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 to obtain a big reduction in PageRank estimator variance by coupling walkers (e.g. footbal and newsgroup), so neither antithetic termination nor our o-coupling can provide a big gain. detailed investigation of how these observations relate to the mathematical graph structure is deferred to future work.
the partieular graph being considered. For some graphs it is not possible
GRFS FOR GPs: ADDITIONAL DETAILS AND EXPERIMENTAL RESULTS
In this appendix, supplement Sec. 4.2 by providing further experimental results for GRFs, including Gram matrix estimation wider range of graphs and a scalable graph-based based GP on 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 the first part of Sec. 4.2 for more real-world graphs. We approximate the 2-regularised Laplacian kernel (I L with regularisation parameter and m walkers. Fig. 10 shows the results, with the kernel approximation emor normalised by the i.i.. result (unlike in Fig. 3) for clarity of comparison. The a-coupling consistently reduces the estimator
37

---

Published as a conference paper at ICLR 2025
karate
dolphins
polbooks
football
newsgroup
norraalseed relative Frob. normn error
1.00
0.9
0.9
0.95
0.95
0.85
0,,00
0.90
0.8
0.2
0.4
0.2
0.4
0.2
0.4
0.2
0.4
0.2
0.4
eu-core
arenas-email
eurosis
citesser
cora
1.00
1.00
1.0
0.9
0.95
0.95
0.8
termination probability
Lid.
antithetic
o-coupled
Figure 10: Relative Frobenius norm error K - ( | between true and approximated Gram matrices for different termination probabilities when walkers are i) i...d,, ii) antithetic (Reid et al., 2024c) iii) o-coupled using OT approach. All results are normalised by the i.idd. variance for easy comparison. Lower is better. The o-coupling is consistently as good or better across values of Phaht and graph topologies and sizes. One standard error is shaded.
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.
F.2 SCALABLE GPS ON GRAPISS
very short introduction to graph- GPs. 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 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 al., 2021; Zhi et al., 2023).
Like their Euclidean counterparts, exact GPs on graphs suffer O(N3) time complexity, making them impractical for very large structures. Techniques 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 1 v-parameter. at the expense of lost flexibility (Borovitskiy et al., 2021). 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, (CRF(U;) constructed by simulating m random walks out of iú and adding weighted contributions only at the coordinates corresponding to visited nodes. For walks with seometrically-distriie
lengths this typically small subset of N. Therefore, the matrix K (Ø(c;) o(v;) ) is sparse, unbiased estimate to any kerne,, not just specific families.5 5 We can use established numerical routines for sparse linear algebra to speed up training, providing computational savings on big graphs. This is analogous the sparse spectrum GPs discussed in App. B (Láázaro-Gredilla et al., 2010); 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 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 future self-contained paper.
"Strictyy: any graph node kernel that expressible function of weighted adjacency matrix.
38

---

Published as a conference paper at ICLR 2025
technical details. Here, provide details to supplement Sec. 4.2. We use mesh graphs made available by Hawson-Haggerty (2023), 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 = x exp(-x> L). Here, L is the normalised graph Laplacian (Eq. 48) and K; y learnable parameters corresponding the signal variance and kernel lengthscale respectively. The observation noise o is also learnable. For every graph, we construct kernel estimate by sampling m = i..dd 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, 2014) for 1000 epochs (after which both the objective value and the hyperparameters are empirically found to have converged). We freeze K; 7, 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 a-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 Phali = 0.4. As reported in the main body, coupling the lengths of graph random walkers permits better approximate Bayesian inference with GPs on graphs.
Training GRF-GPs:
Intuition for the effectiveness of coupling lengths for GRFs. When we couple the lengths of random walks, we 'diversify' the ensemble, sampling walks with different number of hops. This is a very direct way to ensure we explore the graph more effectively, avoiding (by chance) simulating pairs of very similar walks. is intuitive that improving 'coverage' of the graph means we estimate the kernel more effectively, letting us model more interactions between different graph nodes. This may explain why our method is effective in this setting. However, we stress that coupling schemes for estimators defined on discrete spaces like graphs were only introduced very recently (Reid et al., 2024c), so further work is needed for a rigorous understanding of this phenomenon.
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., 2021) 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 small subset of the nodes (325) so good uncertainty quantification essential.
Results for GRFs. We randomly divide the nodes into training and test datasets of sizes Nusin 250 respectively. As described App. F.2, we train scalable GRF-GP with m Li.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 {4, 8, 16} i..dd, antithetic (Reid et al., 2024c) and o-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 posterior). Fig. 11 shows the results. a GRFs consistently do best on all metrics, even with a severely restricted sampling budget.
and Nes
RMSE (mph)
0.50
KL div.
IK-KIp
10.3
n0. walkers
i.i.d.
antithetic
a-coupled
Figure 11: Results for probabilistic traffic interpolation experiment. Plots show the kernel approx- imation accuracy, test RMSE and KL divergence to the true posterior for graph kernels estimated using i..dd, antithetic and cocupledd GRFs. Lower better. o-coupling gives the best results for approximate inference. One standard emror over random draws of walkers is shaded.
39

---

Published as a conference paper at ICLR 2025
G ADDENDUM: FURTHER EVIDENCE OF GENERALISABILITY
Here, we give further evidence of the generalisability of our scheme. We show that our PNC algorithm (Def. 3.3) provides strong empirical improvements for the convergence of estimates of other radial basis function (RBF) kernels, beyond the Gaussian example considered in Sec. 3. In particular, we consider approximating the popular Matérn class of covariance functions (Williams and Rasmussen, 2006), defined by
Matém (I,, D.y) = 21-v
22 I; x x3l12
v 2v x
x xU112
K2
(83)
l(v)
1
with positive parameters v and I, where K2 is a modified Bessel function. is well known that this covariance function has spectral density
S(w) 2d 4/ TTV u + (20)"
2v 12
-(++d/2)
4m w
(84)
[(u)
in d dimensions. We can sample frequency vectors {wi from this density and use them to construct RFFs (Eq. 3), thereby unbiasedly approximating the Matérn kernel.
Clearly, S(w) is isotropic, it is possible condition that frequency vectors mutually orthogonal. Moreover, we are free to couple their norms using our PNC scheme, though now each obeys a different marginal distribution compared to the Gaussian case.
Experimental results. We construct RFFs to approximate the Matérn kernel with smoothness parameters E { 3 2), taking random frequencies marginally distributed as in Eq. 84. As before, 2
we consider sampling frequencies that are (i) i..d., (ii) orthogonal with independent norms, and (iii) orthogonal with PNC norms. For the dataset, we randomly generate N 64 vectors at; ~ (O, La/d) of dimensionality
8. For simplicity, we set the kernel lengthscale = 1. Fig. 12 shows the error on the Gram matrix approximation as the number features grows. Even for this different kernel and across different choices for the smoothness parameter PNC continues to substantially suppress kernel estimator variance, This provides strong evidence that our OT-driven norm coupling scheme generalisable across RBF kernels, and is not specific to the Gaussian example.
Matérn kernel spproximation
10
10
Kiz
7
1.d.
orthogonal
PNC
Figure 12: Kernel estimation RMSE of the Matérn kernel with different smoothness parameters v.
taking different coupling schemes for the frequency vectors. Once again, we see that PNC provides substantial variance reduction, outperforming all the baselines.
40


Reviews:
[{"review_id": "MIeysgB9bG", "reviewer": "Reviewer_kXYZ", "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.\n\nNumerical experiments are presented for each case that show the effectiveness of these algorithms.", "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?\n\nThe 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.\n\nMore specific questions regarding the presentation:\n\n1- I am not familiar to the notion of negative monotone coupling and for me (7) does not make sense. Can you provide more explanation?\n2- 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.\n3-In line 290, what is the definition of \\hat{a}_i?\n4-In line 387, does \"walker's direction\" mean the random walk conditioned on its length?\n5-In Theorem 4.1, I do not understand the definition of q^{i,j}. Also q^{i} is not exactly defined.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.", "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. \n\nI also have a more specific questions/concerns about the justification of this framework which is explained in the next part.\n\nI am willing to raise my score if my concerns are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hz2E0yDOCm", "reviewer": "Reviewer_GwbX", "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.", "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.\n\nAre 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.\n\nOne 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.\"", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The manuscript is well presented, structured, and clearly written.\n\nThere is a high degree of novelty in the application of optimal transport in determining the feature couplings. \n\nThe proposed method is tested on a good selection of real world datasets.", "weaknesses": "I have some concerns regarding some of the experimental results presented in the manuscript:\n\nRegarding 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. \n\nWhile 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).\n\nRegarding 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "uK8mN5bNIJ", "reviewer": "Reviewer_93BC", "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.", "questions": "1. What are the computational implications of implementing these optimal transport techniques at scale?\n2. Can the theoretical findings on variance reduction be generalized to other kernel methods or models?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 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.\n2. **Comprehensive Coverage**: It addresses both Euclidean and discrete input spaces, providing a unified strategy applicable across different domains.\n3. **Theoretical Guarantees and Empirical Validation**: This paper offers theoretical guarantees and empirical validations, enhancing the understanding of variance reduction in RF performance.", "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.\n2. 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.\n3. Despite variance reduction, the downstream impact on model performance can be inconsistent, especially for transformers.\n4. 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.\n\n\n[1]. *General Graph Random Features*. Nips 2023.\n[2]. *Simplex Random Features*. ICML 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "52y6nhhxpu", "reviewer": "Reviewer_huMM", "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.", "questions": "Here below some comments and observations:\n\n- 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. \n- 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.\n- 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "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.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "TejCjP2fkP", "reviewer": "Reviewer_L1dC", "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.", "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? \n\n2. 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? \n\n3. 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. \n\n4. 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. \n\n5. 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? \n\n6. 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The connection between variance reduction of RFs and OT is clearly presented.\n2. RFs for different settings including RFFs, RLFs and GRFs are deeply studied. \n3. A through simulation study has been conducted and some interesting phenomenons show up in the numerical experiments.", "weaknesses": "Please see following questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)