Variance-Reducing Couplings for Random Features

We thank the reviewer for their comments on the manuscript. We are pleased that they find the paper clear and thorough, and that they find the numerical experiments interesting. We address all their questions in detail below.

1. Data dependence of the optimal coupling. Thanks for the great question. In general, the cost function will depend on $(x,y)$ so the optimal coupling may indeed depend on the data. Remarkably, as the reviewer correctly identifies, for RFFs and RLFs with $m=2$ the optimal coupling turns out to be independent of $(x,y)$, and therefore independent of the data distribution. For general $m$, the exact solution to the (generically intractable) multi-marginal OT problem may indeed depend on the data. However, our proposed algorithm (Def 3.3) is data-independent, guaranteed to be better than i.i.d., and empirically still close to optimal (see Sec 3.2 for numerical evidence).
2. Meaning of ‘provided $z$ is sufficiently small. Thanks for the question. We precisely characterise what we mean by ‘provided $z$ is sufficiently small’ in Lemma A.6 (line 990). The discussion is mathematically involved, so we omit it from the main text to preserve its flow and readability. Eq 32 and the following discussion explains how to choose $z$ to guarantee that it is small enough. We invite the reviewer to read this section for full details. Note that in our experiments we find that this condition tends not to matter in practice: for a variety of real-world datasets, our algorithm provides substantial variance reduction for RFFs (see Table 1). We do not expect practitioners to actually need to check Eq. 32 when using the algorithm; it is a technical detail.
3. RMSE and $d$. By RMSE, we mean the Frobenius norm error between the true and approximate Gram matrices. We will make this clearer – thanks. $d$ is not chosen by the practitioner; it is the dimensionality of the data. The standard UCI datasets we use have a range of dimensionalities. We do not see a consistent trend for how this modifies the effectiveness of our coupling, which depends more subtly on the data distribution.
4. Bound for the difference in variance. Thanks for the great suggestion. Using Lemma A.6, it is actually very straightforward to extend Eq. 32 to show that the reduction in variance is lower bounded by $z^4 c / 32 \Gamma(2+d/2) - 2g(z)$, where $c$ is defined in Eq. 30 and $g(z)$ is defined on line 1019. We have added this comment to the paper. Thanks again for prompting it.
5. Downstream tasks for Table 1. We respectfully invite the reviewer to inspect App. B.2, which gives a detailed account of the experiment they suggest here. They are right to note that the fact that variance reduction may not always improve downstream performance, and that OT couplings should be optimised for a different objective, is an important conclusion of our work. Indeed, we already show how a different OT coupling can improve the accuracy of Performer ViTs by +0.8% (see Table 2). We hope our paper will inspire future research in this crucial new direction.
6. Moving $\mathbb{E}_\textrm{dirs}$ inside the square. Thanks for the great question. We agree that moving the expectation inside the square modifies the objective. This is one of several simplifications needed to make the problem tractable, all discussed in detail in App. C.3. This particular step is essential for scalability – without it, we cannot write the objective in a way that can be efficiently estimated using random walks, so the problem becomes intractable. However, we emphasise that, even with these approximations, our method still achieves SOTA variance reduction for length-coupled GRFs compared to previously-proposed alternatives [3].
7. Intuition for GRFs. Thanks for the great question. When we couple the lengths of random walks, we ‘diversify’ the ensemble, sampling walks with a 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. It is intuitive that improving ‘coverage’ of the graph means we estimate the kernel more effectively, letting us model more interactions between different graph nodes. The benefits of coupling the norms or orthogonal frequencies in RFFs and RLFs – effectively modelling the lengthscales of perpendicular components in the data – are less clear. We hope this helps build intuition for the difference between these cases, and have added this discussion to the paper.

We again thank the reviewer. Having answered all their questions (including adding a bound for the amount of variance reduction), we very much hope they will consider raising their score. We warmly invite them to respond if anything remains unclear.

CONTINUED BELOW.

We thank the reviewer for their comments on the manuscript. We are pleased that they find our approach innovative and practical, and that they appreciate the benefits of a unifying perspective applicable to both Euclidean and discrete spaces. We address all their questions and concerns below.

1. Relationship between OT and variance reduction. The Kantorovich formulation of OT seeks to find a joint measure that minimises the expectation of some cost function, conditioned on fixed marginal measures. Meanwhile, a MC estimator will be unbiased if each random variable’s marginal measure is equal to the target measure under which we would like to estimate the expectation. Therefore, simply letting the OT cost function be the variance of the MC estimator, it follows that variance reduction can be considered an OT problem. We will be sure to make this clear in the text.
2. Theorem 3.2. Thm 3.2 exactly solves the OT problem for RFFs and RLFs with $m=2$. Corollary 3.4 extends this to propose a coupling scheme that is strictly better than i.i.d. for any $m$, not just $m=2$. Sec 3.2 then runs numerical experiments to show that this coupling is actually close to optimal – i.e. the (generically intractable) solution to the multi-marginal optimal transport problem – for any $m$. Respectfully, we do not believe any of the above to be a weakness of the paper. We will be happy to respond if the reviewer can further clarify.
3. Variance reduction does not always improve downstream performance. The reviewer is correct to identify that a surprising finding of the paper is that optimising a coupling for variance reduction may not always improve downstream performance. In fact, in Table 2 we give evidence that optimising a different property of the coupling – i.e. solving a different OT problem – can actually improve downstream performance by more. This leads to a +0.8% accuracy gain with a Performer ViT on ImageNet [1]. We believe these observations to be underappreciated in the literature, which almost exclusively focuses on variance reduction [2,3,4,5,6,7,8,9,10,11]. Highlighting these problems and presenting effective alternatives is a major contribution of this work. Respectfully, we do not consider this important part of our work to be a weakness.
4. Relationship to other papers. Thanks for the question. ‘General Graph Random Features’ [12] first proposed the GRFs algorithm for estimating graph node kernels. However, the walkers were taken to be i.i.d. – here, we improve on it by coupling walkers’ lengths. Observe that we comprehensively beat this baseline in Figs 3 and 4. Meanwhile, ‘Simplex Random Features’ [4] improved RLFs by coupling the directions of the frequency vectors. Importantly, it did not couple their norms. The reviewer is correct that we build on the theoretical analysis in this paper, but substantially extend it by introducing OT and for the first time coupling vector norms.
5. Computational implications of implementing these techniques at scale. PNC (Def 3.3) is trivial to implement at scale, and is no more expensive than i.i.d. sampling of norms even though it guarantees lower kernel estimator variance. Meanwhile, the GRF coupling is learned very efficiently using the Hungarian algorithm (line 430), and in our experiments we scale it to massive graphs with as many as $8700$ nodes (see Fig. 4). Indeed, our algorithms are designed to scale well to massive datasets – this is one of their key strengths. We will be sure to make this clear.
6. Can the results be generalised to other kernel methods and models? We already consider kernels defined on both Euclidean and discrete spaces, and GRFs can approximate any graph node kernel that is written as a function of a weighted adjacency matrix. This is already a very broad class. Whether the negative monotone coupling (Eq 7) is optimal for any other Euclidean kernels beyond Gaussian and softmax is an excellent open question, which we defer to future work. Thanks for raising it.

Once again, we thank the reviewer. Having provided detailed clarifications to all their questions, we hope they will consider raising their score. We warmly invite them to respond.

CONTINUED BELOW.

We thank the reviewer for their comments on the paper. We are pleased that they find the connection between Monte Carlo and optimal transport interesting, and that they appreciate our detailed numerical results. We address all their concerns and questions in detail below.

'OT does not obviously lead to a simple rule'. We agree that solving optimal transport problems can in general be very difficult, and that the multi-marginal (MMOT) case is even harder. However, we respectfully bring the reviewer’s attention to the following observations.

1. For RFFs and RLFs, there *is* a simple rule. For $m=2$, we are able to solve the OT problem exactly for RFFs and RLFs (asymptotically for the former). Even better, the optimal joint measure is trivial to sample from, and is independent of the data. This is a great example of a lightweight, simple coupling scheme derived using OT, which guarantees substantial (in fact, optimal) variance reduction. Although the MMOT problem is harder, Sec 3.2 provides strong numerical evidence that, in spite of its simplicity, the $m=2$ scheme is still almost optimal in this more general setting. This shows that OT sometimes gives couplings that are simple and work well.
2. Coupled GRFs are cheap and give SOTA variance reduction. Whilst we agree that the series of approximations needed in Sec. 4.1 may not be obvious at a first glance, we emphasise that the final coupling algorithm is very efficient – one simply simulates random walks on a test graph, uses them to set up a linear weights bipartite matching, and efficiently solves it with the Hungarian algorithm. Once again, the (approximately) optimal measure is very cheap to sample from. Importantly, it also achieves SOTA variance reduction for GRFs. Whilst the reviewer says the approximations might take a bit of `ingenuity’, the resulting GRFs algorithm is both practical and effective, across graph topologies and tasks.
3. The counterintuitive results are important. We agree that the fact that variance reduction may not improve downstream importance is counterintuitive. We believe it to be underappreciated in the literature: an enormous number of papers have been published at top ML conferences seeking to reduce the variance of kernel estimates [1,2,3,4,5,6,7,8,9,10], but to our knowledge none have yet sought to optimise a different property of the coupling like the MSE of the row-normalised attention scores. We agree that more work is needed to fully understand how practitioners should best achieve this and we agree that certain OT objectives will turn out to be intractable, but we have already seen in Table 2 that other simple couplings exist that perform better than previous approaches. In particular, we have shown that the positive monotone coupling is very effective for attention approximation in Transformers, giving accuracy gains of +0.8%. More work in this direction is certainly needed; we hope our broader perspective beyond variance reduction will help guide it. As such, we believe that the counterintuitive nature of the findings is not a weakness of our paper, but rather an important strength.

For the reasons above, whilst we agree that some of our mathematical arguments are technically involved, we respectfully suggest that our algorithms (Def 3.3 an Def 4.2) are actually both simple to implement and empirically very effective.

As a minor point, we are unsure what the reviewer means by ‘guarantees can be proven only for close enough arguments of the kernel function to estimate’. For RFFs, our theoretical guarantees for improving estimates of $k(\boldsymbol{x},\boldsymbol{y})$ hold provided $\|\boldsymbol{x} - \boldsymbol{y}\|_2$ is sufficiently small. This is a very standard asymptotic condition for proofs with RFFs in the literature [1,3], and in our experiments we do not find it to be restrictive. If we misunderstand this comment, please may they clarify?

CONTINUED BELOW.

Other questions and minor points

1. Notation and 'Monte Carlo sampling’. In Eq. 1, each feature indeed depends on the $m$ random frequencies, so we agree that $\phi(\boldsymbol{x}; \\{ \boldsymbol{\omega}_i \\} )$ might be clearer. Eq. 1 is not yet a Monte Carlo estimator; it becomes one when we replace the expectation by the average of empirical samples drawn from $p(\\{\boldsymbol{ \omega }\\})$. We have now phrased this more carefully. Thanks for the suggestions.
2. No second term in $\mathcal{I}(\mu)$. We can omit the $\mathbb{E}(\phi^\top \phi)^2$ term from our objective because, for any coupling with the correct fixed marginals, the estimator is unbiased. Hence, $\mathbb{E}(\phi^\top \phi)=k$, the true kernel evaluation. Since this term in the objective will be the same for any valid coupling, we omit it. This is not specific to the RFF and RLF case. It will be true for variance reduction with any MC random feature.
3. Intuition for GRFs. Thanks for the interesting question. When we couple the lengths of random walks, we 'diversify’ the ensemble, sampling walks with a 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. It 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 so effective in this setting. However, we stress that coupling schemes for estimators defined on discrete spaces like graphs were only introduced very recently [11], so further work is needed for a rigorous understanding of this phenomenon.

We again thank the reviewer for reading the manuscript. We hope that, having clarified the practicality of our method and addressed minor concerns, they will consider raising their score. We warmly invite them to respond with any further questions.

[1] Orthogonal Random Features, Yu et al., NeurIPS 2016, https://doi.org/10.48550/arXiv.1610.09072.
[2] Geometrically Coupled Monte Carlo Sampling, Rowland et al., NeurIPS 2018, https://mlg.eng.cam.ac.uk/adrian/NeurIPS18-gcmc.pdf.
[3] Simplex Random Features, Reid et al., ICML 2023, https://doi.org/10.48550/arXiv.2301.13856
[4] Chefs’ Random Tables, Likhosherstov 2023, NeurIPS 2022, https://doi.org/10.48550/arXiv.2205.15317
[5] Quasi Monte Carlo Maps for Shift Invariant Kernels, Yang et al., ICML 2014, http://proceedings.mlr.press/v32/yangb14.html [6] FastFood, Le et al., ICML 2013, https://doi.org/10.48550/arXiv.1408.3060
[7] Structured adaptive and random spinners for fast machine learning computations, Bojarski et al, AISTATS 2017, https://doi.org/10.48550/arXiv.1610.06209
[8] The unreasonable effectiveness of structured random orthogonal embeddings, Choromaksi et al., NeurIPS 2017, https://doi.org/10.48550/arXiv.1703.00864
[9] Spherical structured feature maps for kernel approximation, Lyu et al, ICML 2017, http://proceedings.mlr.press/v70/lyu17a/lyu17a.pdf
[10] Random features for shift-invariant kernels with moment matching, Shen et al, AAAI 2017, https://doi.org/10.1609/aaai.v31i1.10825
[11] Quasi Monte Carlo Graph Random Features, Reid et al., NeurIPS 2023, https://doi.org/10.48550/arXiv.2305.12470

We sincerely thank the reviewer for their swift and detailed reply. We are grateful for their interesting perspectives. We are happy to share our thoughts below.

Generality of the framework and lack of simplicity. The reviewer is correct to note that, since we take the RF estimator variance as an OT cost function, the OT map will in general differ depending on the specific RF construction. Therefore, the ‘difficulty’ of the problem is upper bounded by the difficulty of OT, and by constructing an ‘adversarial’ RF one could in principle make the variance reduction problem very difficult or intractable. However, we respectfully raise the following points.

1. We already consider two Euclidean RF constructions (RFFs and RLFs) and two different kernels (Gaussian and softmax). Up to an asymptotic assumption for RFFs, the OT formulation yields the same, simple coupling. It is not a priori obvious that this should be the case, since the two RFs considered use totally different basis functions (exponentials vs sine and cosine). In this sense, the PNC coupling is already actually quite general.
2. Moreover, in the literature, it is often found that coupling strategies that are effective for the Gaussian kernel – like conditioning orthogonality [1] – are also effective for a range of other radial basis function kernels such as the Matérn kernel [2]. Whilst we defer rigorous study to future work, we are optimistic that this norm coupling strategy may also reduce the variance of other kernels. It seems intuitively reasonable that trying to sample the norms of frequency vectors more diversely might improve estimator convergence for a wider variety of kernels.
3. Lastly, since the Gaussian and softmax kernels are so ubiquitous in machine learning (both for attention in Transformers and in Bayesian ML), we suggest that variance-reducing couplings specifically for these cases may still be of general interest. All the papers cited in our previous rebuttal focus on these examples; it is common for researchers to tackle this special case. For these reasons, we suggest that our method will be of general interest.

Counterintuitive results for variance reduction. Thanks for the balanced perspective. We agree that there are different possible interpretations of our findings, although (predictably!) we tend to prefer the latter suggested by the reviewer. This is for the following reasons.

1. Evidence that variance reduction is not the correct strategy is very important for the field. As cited above, very many papers and hours of research have been dedicated to variance reduction. Our smoking gun evidence that variance is not the last word for applications (as the reviewer says, we actually ’solve’ the variance reduction problem for RFFs/RLFs in particular special cases) gives compelling motivation to redirect and broaden these efforts, thinking about ‘optimal couplings’ in a more general sense than just optimising the estimator variance. We believe this to be practically important for the field’s future direction, even if in some sense it could be considered a null result.
2. OT is not specific to variance reduction. Moreover, whilst we mostly focus on coupling for variance reduction (apart from the brief section about Transformers where we maximise variance), we stress that the OT framework is not specific to this case. The reviewer also pointed out this important observation earlier. In principle, the OT framework can be used to optimise any property of the estimator. Whilst we agree that the OT problem may in general be difficult depending on the choice objective, we believe that our perspective and research toolbox – both theoretical (Thm 3.2, Corollary 3.4, Thm 3.5, Thm 4.1) and numerical (Sec. 3.2, Sec. 4.1) – will be helpful.

As such, we respectfully suggest that the counterintuitive findings are not a problem with the OT reformulation per se (read: the optimisation method), but rather with variance reduction as a goal (read: the objective function). Our results will be important in the quest for better couplings.

Once again, we sincerely thank the reviewer for the response. Their thought-provoking questions have improved the manuscript. If satisfied with the above, we again politely ask that they recommend acceptance.

[1] The Geometry of Random Features, Choromanski et al., AISTATS 2018, https://proceedings.mlr.press/v84/choromanski18a.html
[2] Orthogonal Random Features, Yu et al., NeurIPS 2016, https://doi.org/10.48550/arXiv.1610.09072.

We thank the reviewer for continuing to engage so wholeheartedly in the reviewing process.

Generalisability. Following the reviewer’s comments, to further demonstrate the generalisability of our approach, we have now added an extra experiment using OT-coupled RFs to approximate the Matérn kernel. Across a range of smoothness parameters $\nu$ describing very different function behaviours, we again find that our algorithm provides excellent variance reduction, outperforming all baselines. This demonstrates that our approach can be used to improve estimates of a totally different, very popular class of covariance functions, beyond the Gaussian example. Please take a look at Sec G.1 (page 39), especially Fig. 12.

To summarise, this means that our framework has now been successfully applied:

1. To RFFs and RLFs;
2. To estimate the Gaussian kernel and the Matérn kernel (the latter for smoothness parameters $\nu \in \\{1/2,3/2,5/2 \\})$;
3. For GRFs to estimate kernels on graph nodes;
4. For both real and synthetic graphs (where variance reduction gives big downstream gains);
5. And for estimating the PageRank vector (a totally different application, which isn’t even about RFs! See App. E for full details).

We firmly believe that this shows that our approach generalises well, across RF basis functions, target kernels, and even input domains. It is very typical for papers to focus on just one of these instances listed above, e.g. Yu et al. [1] just consider approximating the Gaussian kernel with RFFs, and Reid et al. [2] just consider improving approximating of the regularised Laplacian graph kernel with GRFs. On the other hand, we provide substantial gains for all of them.

Variance reduction as a goal. We remind the reviewer that, for GRFs, variance reduction gives big, SOTA gains in downstream inference tasks (see e.g. Fig. 4 and Fig. 11). Clearly, variance reduction can bring benefits and should still be considered an important goal for Monte Carlo research -- hence its continued popularity. However, for RFFs and RLFs, we do indeed find that the benefits of variance reduction saturate. Alternative properties for the coupling are found to be important e.g. in efficient Transformers, for which we suggest an interpretation in terms of OT. To be clear, we still believe variance reduction to be an important research enterprise. However, pushing it to its limits using OT has revealed that we may also need to consider optimising other properties of the coupling. We will make this nuanced point clearer in the text. Thanks for the suggestion.

Summary of novel contributions. To avoid getting bogged down in the specific part of the paper about the limitations of variance reduction for RFFs/RLFs, we thought it might be helpful to recapitulate the paper’s main novel contributions.

1. First application of OT to improving the convergence of RFs;
2. First analytic solution to OT problem for orthogonal RFFs/RLFs (may be of independent interest to OT researchers, Thm 3.2);
3. Novel numerical methods for when the OT problem is analytically intractable:

• Numerical copula approach for optimising couplings in Euclidean space (Sec. 3.2);
• Permutation densities and Hungarian algorithm for optimising couplings in discrete space (Sec. 4.1).

4. First data-based optimisation of a coupling between random walks, and first application of GRFs to Gaussian processes on graphs (where our variance reduction techniques unlock better approximate inference and beat all baselines -- see Sec. 4.2).
5. (Already discussed in detail -- demonstration that benefits from variance reduction can saturate with RFFs/RLFs, and proposal of other properties of the coupling to optimise (Sec. 3.3)).

To run with the hammer and nail analogy suggested by the reviewer: we have substantially upgraded the hammer (points 1-3), introduced some new nails (point 4), but found that at some point existing nails may stop benefiting from a better hammer (point 5). We believe all these contributions to be integral to our work.

Once again, we sincerely thank the reviewer for their engagement and for encouraging us to spell out these points more clearly. With the additional Matérn kernel experiments and discussion in mind, we politely ask whether they might consider a further score increase.

[1] Orthogonal Random Features, Yu et al., NeurIPS 2016, https://doi.org/10.48550/arXiv.1610.09072.
[2] Quasi Monte Carlo Graph Random Features, Reid et al., NeurIPS 2023, https://doi.org/10.48550/arXiv.2305.12470

We sincerely thank the reviewer for their thoughtful comments on the manuscript. We are pleased that they believe it to be well written and highly novel, and that they agree we test on a good selection of real world datasets. We address all their concerns and questions below.

Number of features plot. Thanks for the suggestion. We actually already include this in Fig 6: a plot showing the kernel estimator RMSE, prior and posterior KL divergence, and predictive RMSE against numbers of features $m/d \in \\{1,2,4\\}$. This is exactly analogous to Fig. 1 in 'Orthogonal Random Features’ by Yu et al. [1]. Note also that, for all methods considered, the kernel estimator RMSE goes down as $\sim (m/d)^{-1/2}$ as we introduce extra `blocks’ of size $d$. This is because the couplings are within each block rather than between different blocks. Therefore, the ratio of respective variances between the different methods does not change as $m/d$ grows, since they all just scale down together. In this sense, just considering $m=d$ to show variance reduction is sufficient (though of course other downstream quantities will saturate as the estimators converge to the true kernel).

Normalisation of RMSEs. Thanks for the suggestion. As the reviewer says, we normalised the predictive RMSEs for simplicity of comparison. However, we agree that including the overall scales might be useful for reproducibility, so have added these to the manuscript. (With that said, we do emphasise that the purpose of this experiment is to demonstrate greater variance reduction compared to previous variance reduction approaches, rather than necessarily achieving SOTA performance on the UCI regression tasks).

KL and NLL. Thanks also for this suggestion. Our intention with the experiments in Sec. B.2 is to compare the quality of approximate inference when the kernel is estimated with different methods. To achieve this, we report the KL divergence between the true and approximate GP predictive posteriors (Table 3), which captures how close the approximate posterior is to the result with the full rank kernel. We also report the predictive RMSEs on held out test sets (Table 4), which straightforwardly measures the quality of the predictions. Note that we are more interested in quantifying the accuracy of our approximate predictor compared to the groundtruth with an exact kernel, rather than making a statement about the predictive performance of the exact kernel in absolute terms. That said, other choices like the predictive NLL are also possible. When we compute the NLL In experiments, we find that it shows similar behaviour as our existing metrics: variance reduction by coupling does not always improve approximate inference. We initially omitted these NLL results for brevity – the paper is already 38 pages! However, we agree that they might be a nice addition, so we have now added them to the paper. Thanks again for the comment.

Are there applications where kernel reconstruction is the primary focus? We certainly find that the importance of kernel approximation quality depends on the particular setting and downstream task – e.g. for RFFs and RLFs variance reduction by coupling norms appears less effective, whereas for GRFs we see very big gains. Understanding this complicated behaviour is challenging. We hope our results will prompt future work in this direction.

On being 'surprised’. Thanks for the extra pointer to Liu et al. [2004], which we have incorporated into the manuscript. However, we still respectfully suggest that the complicated relationship between variance reduction and downstream performance is not always fully appreciated in the literature. An extraordinary number of papers have been published on variance reduction, including in top ML conferences [1,2,3,4,5,6,7,8,9,10]. But here we show that a modification as simple as normalising attention scores by their row-wise sum can mean that the maximum variance coupling actually works better (from line 287). More detailed, thorough work is needed to optimise the `right’ properties of the MC coupling. We hope that our work will prompt further investigation, and that it will help provide mathematical and algorithmic tools by making the connection to optimal transport.

Once again, we thank the reviewer for their helpful comments and suggestions. We believe that we have addressed all their concerns, and respectfully ask that they consider raising their score. We warmly invite them to respond with any further questions.

CONTINUED BELOW.

We appreciate the reviewer’s response and their thoughts about the best setting to test the downstream effects of variance reduction.

$m=d$ choice. We initially chose $m=d$ features because at this point the absolute difference in variance between the algorithms is at its greatest. (As mentioned in our previous response, the relative difference does not change as $m/d$ grows in integer multiples because the correlations are within each ‘block’ of size $d$). Hence, we expected that any difference in downstream performance to be most obvious in this regime. This is indeed reflected in Fig. 6, where the improvements in predictive RMSE and KL divergence are found to be greatest to the left of the plots when $m=d$. This behaviour is also well-established in the literature. To give a few examples, see:

1. Table 2 by Yu et al. [1], where orthogonal random features only consistently appear to outperform i.i.d. features for SVM classification accuracy when $m \leq 4d$.
2. Fig 8 by Reid et al. [3], when the gap between i.i.d., orthogonal and simplex random features rapidly drops with $m$ in the KRR experiment when $1 \leq m/d \leq 10$.
3. Table 1 by Yang et al. [5], where the gap between the different methods on the regression task is found to drop as the number of features grows.

The reason is that, with more features, all the estimates approach the true kernel evaluation. Therefore, the performance of each method saturates at the performance of the exact kernel. They soon become indistinguishable. It is not the case that one should only expect gains from variance reduction at large $m$.

That said, we do agree that it may be helpful to also compare the algorithms in the setting with more features, where all the kernel approximations are better so the overall performance on downstream tasks is more competitive/realistic. Therefore, we have repeated the experiments with $m=8d$ features and added the results to App. B.2. Please see Table 4, page 26. The trend is the same: despite providing variance reduction compared to the baselines, the downstream benefits from PNC (and in this case actually variance reduction more generally) are less clear. Note that the results are consistent with Fig. 6 for a split of the power dataset, where the difference between all methods becomes small as they approach the groundtruth kernel. The reviewer will also note that with more features we outperform linear regression, which gives predictive RMSE $10.43 \pm 0.56$ on 10-fold cross validation on the concrete dataset (script to get this result temporarily uploaded to supplementary material). Thanks for prompting this addition to the paper.

(Note that these preliminary results are averaged over fewer random seeds than Table 3 because of time constraints; in the final draft the standard errors will be smaller and minor discrepancies will be ironed out).

Fig 6. Thanks for the comments about Fig. 6. We agree that it’s a nice contribution! We found it difficult to squeeze it into the main body with space restrictions: we would like to leave enough space for our important GRF results, where we do see very substantial benefits from variance reduction, and our extensive theoretical contributions. However, we have updated the text to signpost interested readers to this figure more explicitly.

Thanks again for the feedback, and for confirming that you will be happy to reconsider your score at the end of the discussion period. We look forward to hearing your response, and warmly welcome any further questions.

[1] Orthogonal Random Features, Yu et al., NeurIPS 2016, https://doi.org/10.48550/arXiv.1610.09072
[2] Geometrically Coupled Monte Carlo Sampling, Rowland et al., NeurIPS 2018, https://mlg.eng.cam.ac.uk/adrian/NeurIPS18-gcmc.pdf
[3] Simplex Random Features, Reid et al., ICML 2023, https://doi.org/10.48550/arXiv.2301.13856
[4] Chefs’ Random Tables, Likhosherstov 2023, NeurIPS 2022, https://doi.org/10.48550/arXiv.2205.15317
[5] Quasi Monte Carlo Maps for Shift Invariant Kernels, Yang et al., ICML 2014, http://proceedings.mlr.press/v32/yangb14.html.
[6] FastFood, Le et al., ICML 2013, https://doi.org/10.48550/arXiv.1408.3060
[7] Structured adaptive and random spinners for fast machine learning computations, Bojarski et al, AISTATS 2017, https://doi.org/10.48550/arXiv.1610.06209
[8] The unreasonable effectiveness of structured random orthogonal embeddings, Choromaksi et al., NeurIPS 2017, https://doi.org/10.48550/arXiv.1703.00864
[9] Spherical structured feature maps for kernel approximation, Lyu et al, ICML 2017, http://proceedings.mlr.press/v70/lyu17a/lyu17a.pdf
[10] Random features for shift-invariant kernels with moment matching, Shen et al, AAAI 2017, https://doi.org/10.1609/aaai.v31i1.10825

Thanks for the response. Let us be clearer on these points.

Previously, the only experimental results in the main text for coupled RFFs and RLFs were kernel estimator RMSE (Table 1). This shows the precise amount of variance reduction – after all, our chief goal – so it is probably the most important result to include. Since we normalise each kernel estimator RMSE by the value for i.i.d. features, the results are the same for any $m=kd$ (or $m=2kd$ for RLFs), where $k \in \mathbb{N}$ is a positive integer. Table 1 is not specific to $m=d$. For example, on the concrete dataset, using orthogonal + PNC RFFs instead of i.i.d. RFFs will always reduce the variance by the same ratio, whether you take $m=d$ or $m=1000d$. We have made this clearer in the text.

We agree that considering downstream performance (predictive RMSE, KL divergence between true and approximate posterior, NLLs, etc.) on a range of values of $m/d$ may be interesting, including when $m$ is big so performance is better/more realistic. This is why, following the reviewer's comments, we now include full results not only for 1) $m=d$ (where absolute differences in kernel estimator variance, and hence performance, are biggest), but also 2) $m=8d$ (where differences between methods become very small because all approximations are close to the exact kernel). The latter is of course substantially better than linear regression for all datasets. It would also be very simple to include more values of $m$ between these regimes (or yet bigger $m$) if the reviewer thinks this would be a nice addition – this is just a matter of space.

Therefore, our paper already includes all the results the reviewer is asking for. It is simply a case of where they are located in the text – the main body or the appendix. We agree that this is an important decision. On reflection, we agree that the reader might benefit from seeing downstream results in the main text, so we have moved the ablation over $m$ (now Fig. 1) into the main text. To make space, we were forced to move the copula schematic into the Appendix – a tricky choice. Thanks for this suggestion, which has probably made our message clearer overall.

Could the reviewer please confirm that this latest change – including downstream performance for multiple $m/d$ in the main text, rather than in the appendix – has resolved their remaining concerns?

We thank the reviewer for their comments on the manuscript. We are pleased that they enjoyed reading it, and that they found it well written and easy to understand. We clarify all points of confusion and address their questions below.

Computational aspects and pairs of points. Thanks for the question. As written, Eq. 2 indeed considers variance reduction for a single pair of points, and the optimal transport map might in general depend on this specific pair. However, in Sec. 3.1 we find that, after working through the maths, the optimal RLF and RFF coupling is actually independent of the data (Eq. 7). Hence, the negative monotone coupling in Eq 7 is optimal for any data distribution. Meanwhile, for GRFs, we average the cost function over different pairs of nodes of the graph to learn a coupling that works well for the whole kernel matrix (line 463). This is the same as the reviewer’s suggestion to ‘minimize the variation over the process’. We also find that the optimised coupling generalises very well to new graph topologies (line 466). To be clear, one does not need optimise the coupling separately for every pair of datapoints – the coupling we use is either data-independent or the cost function is averaged over a training distribution. Therefore, our methods continue to scale well to very large datasets, including massive graphs with 8700 nodes (see e.g. the experiments in Fig. 4).

Questions regarding the presentation. Thanks for the questions. We are happy to clarify them all below, and will make sure the manuscript is as clear as possible on these points.

1. Negative monotone coupling. Draw a random variable $u_1$ from the uniform distribution on $[0,1]$. Push it forward with the inverse CDF of the Chi distribution, defining a new variable $\omega_1 = {F^{-1}}(u_1)$. The variable $\omega_1$ will follow the Chi distribution. Define $u_2 = 1-u_1$, and let $\omega_2 = {F^{-1}}(u_2)$. The marginal distribution of $u_2$ is still uniform on $[0,1]$ so the marginal distribution of $\omega_2$ is still the Chi distribution. However, $u_1$ and $u_2$ are now negatively correlated, so $\omega_1$ and $\omega_2$ are also negatively correlated. This is the negative monotone coupling, shown in Eq. 7.

2. Theorem 3.2 and asymptotic $z$. For random Laplace features, the negative monotone coupling is always optimal. For random Fourier features, it is optimal provided $z$ is sufficiently small. Formulating this asymptotic statement precisely is technically involved so we put the details and proof in the Appendix (Lemma A.6). We also emphasise that, in experiments on real data, the negative monotone coupling performs extremely well even when $z$ is not small (see Table 1).

3. Definition of $\textrm{MSE}(\widehat{a}_i)$. We defined attention $a_{ij} = k(x_i,x_j)/ \sum_l k(x_i, x_l)$, the row-normalised kernel evaluations, on line 287. The estimator $\hat{a}_{ij}$ is obtained by replacing the exact kernel evaluations $k(x_i, x_j)$ with their stochastic Monte Carlo estimates $\hat{k}(x_i, x_j)$. Then $\textrm{MSE}(\hat{a}_i)$ averages the MSE of this estimator over the keys indexed by $j$, as shown in line 290. We have made this clearer in the text – thanks.

4. ’Expectation over walkers’ directions’. Yes, $\mathbb{E}_\textrm{dirs}$ averages the walkers’ directions – i.e. which of their neighbouring nodes they choose at every timestep. It does not consider their lengths, which are instead shown in the first expectation of Eq. 12.

5. Notation in Thm. 4.1. The vector $\widehat{\psi}(q^{(i)})$ is defined on line 414. In Thm. 4.1, it would be clearer if we had written $\\{\widehat{\psi}(q^{(i)}), \widehat{\psi}(q^{(j)})\\}^n_{q=1}$ so we have updated the notation. Thanks for the suggestion.

We again thank the reviewer for their time. We have clarified the efficiency of our scheme (especially, that we do not need to solve the OT problem separately for every possible pair of datapoints). We have also improved the presentation of the text where there were points of minor misunderstanding. We hope that the reviewer will consider raising their score, and invite them to respond with any further questions.

This is a brief note to let the reviewer know that we have uploaded a new version of the manuscript. Changes are shown in red. They may be especially interested to see the updated notation on lines 288 (shows $\widehat{a}_{ij}$) and 419 (shows $\widehat{\psi}(q^{(i)})$) following their helpful comments. We look forward to hearing their thoughts, and hope that they will consider raising their score.