Supervised Kernel Thinning

Thank you for the time you’ve spent reviewing our work and for your thoughtful feedback. We now address each of your questions and concerns in turn.

$\blacktriangleright$ $**Comparisons to existing methods for efficient non-parametric regression**$

• "The application are only to toy examples or standard small datasets..."
• "In the experiments, the baseline is chosen to be NW / KRR with naive thinning..."

Thank you for this suggestion.

• We have added an experiment on the SUSY dataset ($d=18, N=5\times 10^6$) from Diaz et al. (2023). We sample $4\times 10^6$ points to use for training. (Full KRR is now infeasible as we can only store 3k-4k columns at a time, e.g., on a machine with 100 GB RAM).
• We include the RPCholesky preconditioning method from Diaz et al. (2023) and the FALKON method from Rudi et al. (2017) as SOTA baselines. As a sanity check, we also include Conjugate Gradient method (without preconditioning).
• In Fig. 2 of the attached PDF, we have added RPCholesky-NW (where $S_{NW}$ are the pivot points from RPCholesky) and Restricted KRR with RPCholesky for landmark selection (see box plots in red). While it outperforms the KT method in the KRR setting, RPCholesky is less suitable for speeding up the NW estimator as shown in Fig. 2a of the attached PDF.

$\newline$

$\blacktriangleright$ $**Kernel assumptions**$

• "In the beginning of Sec 3..."
• "The theory makes the sub-Gaussian tail decay assumption on the kernel..."

Thank you for raising this question.

• Our theory in fact permits a large class of base kernels. This strength is borrowed partially from the Kernel Thinning results of Dwivedi & Mackey (2021).
• In our revision, we will use $\mathbf{k}$ exclusively to denote the base kernel and use $\mathbf{k} _{ALG}$ to denote the kernel for thinning. Given a valid base kernel $\mathbf{k}$, we can construct $\mathbf{k} _{NW}$ and $\mathbf{k} _{RR}$ following (4) and (6), respectively.
• For NW, we only showed the results when the kernel is compactly supported or admits sub-Gaussian tails (for ease of presentation); remarkably, the same results hold when the kernel has sub-exponential tails or poly-tails as well. In fact, $\mathbf{k}$ can be any radial kernel $\mathbf{k}(x_1,x_2) := \kappa(\lVert x_1-x_2\rVert_2/h)$ such that $\kappa(u)$ is a bounded, L-Lipschitz function and decays (at any rate). This includes kernels with sub-Weibull tail decay. The kernel density estimation guarantees from Dwivedi & Mackey (2021), which we use in our analysis, indeed use an Orlicz bound.
• For KRR, $\mathbf{k}$ can be any kernel such that the log-covering number satisfies the PolyGrowth (e.g., finitely-times differentiable kernels) or LogGrowth (e.g., analytic kernels, finite-rank kernels) property (see Assump. 4). For KRR, our results do degrade with increasingly less smooth kernels.

$\newline$

$\blacktriangleright$ $**Multi-dimensional response data**$

• "The method (or at least the theory) is developed for one-dimensional response data..."

The reviewer’s comment on multi-dimensional response data has direct and interesting applications for speeding up multivariate response regression problems.

• Roughly speaking, KT with input kernel $\mathbf{k}_{ALG}$ provides a coreset that well approximates the averages of functions in its RKHS $\mathcal H$.
• Dwivedi & Mackey (2022), Thm. 4 showed that thinning with $\mathbf{k} _{ALG} := \mathbf{k} _1 + \mathbf{k} _2$ produces a coreset that simultaneously well approximates averages for any $f\in \mathcal{H}(\mathbf{k} _1)$ or $f\in \mathcal{H}(\mathbf{k} _2)$ (aka the "aggregate kernel trick").
  • We use the aggregate kernel trick to tackle multi-variate regression via one-dimensional kernel thinning on an aggregated kernel.
• To be precise, let $f = (f _1, …, f _m)$ be multi-dimensional response function, and for all $j = 1, …, m$, $f _j : R^d \to R$ are coordinate functions in some RKHS $\mathcal{H}( \mathbf{k} _j)$ generated by some base kernel $\mathbf{k}(j)$.
• Now define the aggregated kernel $\mathbf{k} _{ALG}((x_1,y_1),(x_2,y_2)) := \sum _{i=1}^{m} w _j \mathbf{k}(j) ((x _1,y _{1,j}),(x _2,y _{2,j}))$ (for NW and use $\mathbf{k} _{RR}(j)$ instead of $\mathbf{k} _{NW}(j)$ for KRR), where weights $w _j$ reflect prior beliefs.
• Then $\mathbf{k} _{ALG}$-kernel thinned core-sets are shared across all coordinate functions, that also balances information across all coordinates.
• Notably, aggregate kernel tricks have been used in Generalized Kernel Thinning by Dwivedi & Mackey (2022) for improving integration error guarantees and Compress-then-Test by Domingo-Enrich et al. (2023) for increasing power of kernel-based hypothesis tests. Here we use it for speeding multi-variate non-parameteric regression.

$\newline$

$\blacktriangleright$ $**Minimax optimal rates**$

• "In the discussion section, it is mentioned that the rates on the excess risk could be sub-optimal..."

We will include the expanded discussion in the revision:

• The following is classically known (see Tsybakov (2009)):
  • When $f^\star$ belongs to a Hölder class with smoothness $\beta > 0$, the minimax optimal rate is $\mathcal{O}(n^{-\frac{2\beta }{2\beta + d}} )$.
  • When $f^\star$ belongs to a Sobolev space with smoothness $\nu=1/\alpha$ (for $\alpha\in (0,2)$), the minimax optimal rate is $\mathcal{O}(n^{-\frac{2}{2+\alpha}})$.
• As shown by Dwivedi & Mackey (2021), KT produces a coreset that achieve the minimax rates (as if all $n$ points were used) for a large class of kernel density estimation tasks—e.g., employing any radial kernel $\mathbf{k}(x_1,x_2) = \kappa( \lVert x_1-x_2\rVert _2 / h)$ with bounded, L-Lipschitz $\kappa: \mathbb{R} \to \mathbb{R}$. However, our current proof techniques do not yield a minimax rate for supervised learning task. We believe a modified proof technique should yield the same rate as obtained by $n$ points for some class of kernels.

Thank you for the time you’ve spent reviewing our work and for your thoughtful feedback. We now address each of your questions and concerns in turn.

$\blacktriangleright$ $**Comparison with large-scale kernel methods**$

• "The comparison with other large-scale kernel methods is limited to RPCholesky"

Thank you for this suggestion. We have added new results in the global response pdf.

• It includes an experiment on the SUSY dataset ($d=18, N=5\times 10^6$) from Diaz et al. (2023). We sample $4\times 10^6$ points to use for training. (Full KRR is now infeasible as we can only store 3k-4k columns at a time, e.g., on a machine with 100 GB RAM).
• We include the RPCholesky preconditioning method from Diaz et al. (2023) and the FALKON method from Rudi et al. (2017) as SOTA baselines. As a sanity check, we also include the Conjugate Gradient method (without preconditioning).
• In Fig. 2 of the attached PDF, we have added RPCholesky-NW (where $S_{NW}$ are the pivot points from RPCholesky) and Restricted KRR with RPCholesky for landmark selection (see box plots in red). While it outperforms the KT method in the KRR setting, RPCholesky is less suitable for speeding up the NW estimator as shown in Fig. 2a of the attached PDF and is outperformed by KT-NW.

$\newline$

$\blacktriangleright$ $**Kernel mis-specificiation**$

• "the authors acknowledge the problem of kernel mis-specification but do not elaborate or provide any resolution"

We will elaborate our discussion in the revision.

• The supervised KT guarantees apply when $f^\star$ lies in the RKHS $\mathcal{H}(\mathbf{k})$, where $\mathbf{k}$ is the base kernel.
• In practice, choosing a good kernel $\mathbf{k}$ is indeed a challenge common to all prior work.
• Our framework is friendly to recent developments in kernel selection to handle this problem:
  • Dwivedi & Mackey (2022), Cor. 1 provide integration-error guarantees for KT when $f^\star \notin \mathcal{H}(\mathbf{k}).
  • Moreover, there are recent results on finding the best kernel (in testing; Dominigo-Enrich et al. (2023), Sec. 4.2).
  • Radhakrishnan et al. (2024) introduce Recursive Feature Machines, which use a parameterized kernel $\mathbf{k} _{M}(x_1,x_2) := \exp(-\frac{(x_1-x_2)^\top M (x_1-x_2)}{2\sigma^2})$, and propose an efficient method to learn the matrix parameter $M$ via the average gradient outer product estimator. An exciting future direction would be to combine these parameterized (or "learned") kernels with our proposed KT methods for non-parametric regression.
• Formally combining these results, potentially with the recursive feature machines of Radhakrishnan et al. (2024) to build a practical strategy with theoretical guarantees is an exciting future direction.

Thank you for the time you’ve spent reviewing our work and for your thoughtful feedback. We address each of your concerns below.

$\blacktriangleright$ $**Comparison with RPCholesky**$

• "The empirical advantage..."
• "KT-KRR does not demonstrate..."

Thank you for raising these points. The main focus of this work was to build a stepping stone towards bridging two literatures: distribution compression and non-parametric regression, by designing algorithms and developing theory. Tuning to optimize empirical performance, e.g., by combining the pre-conditioning step of RPCholesky with KT, is a great future direction.

As another sanity check, we have added an experiment on the SUSY dataset ($d=18, N=5\times 10^6$) from Diaz et al. (2023). KT-KRR's (22%) performance lies between RPCholesky (20%) and ST-KRR (22.7%). Thinning 4 million points with our method took only 1.7 seconds on a single CPU core, with further speed-ups to be gained from parallelizing on a GPU.

• "Fig 2: Right panel: why not include RPCholesky?"

In Fig. 2b of the attached PDF, we have added Restricted KRR with RPCholesky for landmark selection. In Fig. 2a of the attached PDF, we have added RPCholesky-NW, where $S_{NW}$ are the pivot points from RPCholesky.

$\newline$

$\blacktriangleright$ $**Choice of kernel**$

• "Eq 4 and 6: why do you choose the kernel in this way?..."
• "Sec 4.1: ... the kernels listed in 1.2.3 used... are the kernel for thinning?"

This is a great question central to the novelty of our work while bridging the two literatures. Directly applying generic kernels would not work in theory and our experiments confirm that. Overall the choice we make are informed by the following points:

• Roughly speaking, KT with input kernel $\mathbf{k}_{ALG}$ provides a coreset that well approximates the averages of functions in its RKHS $\mathcal H$.

• Dwivedi & Mackey (2022), Thm. 4 showed that thinning with $\mathbf{k} _{ALG} := \mathbf{k} _1 + \mathbf{k} _2$ produces a coreset that simultaneously well approximates averages for any $f\in \mathcal{H}(\mathbf{k} _1)$ or $f\in \mathcal{H}(\mathbf{k} _2)$ (aka the "aggregate kernel trick").

• The NW estimator is a ratio of averages: an average of $f(x)=k(x, x_0)$ in the denominator and an average of $f _{num}(x,y)=\langle y,1 \rangle k(x,x_0)$ for the numerator.

  • $f _{num}$ lies in the RKHS corresponding to the kernel $\mathbf{k} _{num}((x_1,y_1),(x_2,y_2)) := \mathbf{k}(x_1,x_2) y_1 y_2$, so thinning with this kernel provides good approximation of the numerator.
  • Applying aggregate kernel trick to $\mathbf{k} _1 = \mathbf{k}$ and $\mathbf{k} _2= \mathbf{k} _{num}$ yields the $\mathbf{k} _{NW}$ from (4).

• For KRR, we approximate the training loss, which is an average of $\ell_f(x,y) = (f(x)-y)^2 = f^2(x) -2 f(x)y + y^2$.

  • Assuming $f\in \mathcal{H}(\mathbf{k})$, $\ell_f$ lies in the RKHS of $\mathbf{k}^2(x_1,x_2) + y_1 y_2 \mathbf{k}(x_1,x_2) + y_1^2 y_2^2$, where the last term can be dropped (see Sec. 3 of Compressed Empirical Measures by Grünewälder (2022)), yielding (11) (thanks for pointing out the typo).

• Remarkably, our ablation studies in Fig. 2 of the original manuscript and Fig. 2 of the attached PDF confirm this theory.

  • Thinning with $\mathbf{k} _{ALG}((x_1,y_1),(x_2,y_2)) = \mathbf{k}(x_1,x_2)$ (orange; Ablation #2 in Sec. 4.1) performs poorly because it only exploits the $x$ information.
  • Thinning with $\mathbf{k} _{ALG}((x_1,y_1),(x_2,y_2)) = \mathbf{k}(x_1\oplus y_1,x_2\oplus y_2)$ (i.e., ) (brown; Ablation #1 in Sec. 4.1)—despite incorporating information from both $x$ and $y$ by concatenating the two—also yields poor empirical performance.
  • Thinning with $\mathbf{k} _{ALG} = \mathbf{k} _{NW}$ and $\mathbf{k} _{ALG} = \mathbf{k} _{RR}$ perform the best for the NW and KRR settings, respectively.

$\newline$

$\blacktriangleright$ $**Improvements to notation and writing**$

• "Notation ... instead of for KT-NW, could use...KT-KRR part."

We will use $S_{NW}$ to denote the coreset for the thinned NW estimator and $S_{RR}$ to denote the coreset for the thinned KRR estimator. We will use $\mathbf{k}$ exclusively to denote the base kernel and use $\mathbf{k} _{ALG}$ to denote the kernel for thinning.

• "There is not an algorithmic description for Kernel Thinning algorithm."

We apologize for this oversight. We will add a paragraph in the revised paper and an algorithm with details in the appendix.

• "Fig 3 and Fig 4: why not use the function in (11)?"
• "Fig 3 and Fig 4: what is the unit of y axis?..."

In the attached PDF, we have reproduced Fig. 3 and 4 with (11) and improved the formatting. The y-axis is in seconds and we use log-scale to make the difference between ST, KT, and RPCholesky clearer.

• "Sec 4.2: what is the conclusion from Fig 5?"

We will revise the discussion for Fig. 5 as follows: On the California Housing dataset, KT-KRR lies between ST-KRR and RPCholesky in terms of test MSE. In terms of train time, KT-KRR is quite a bit faster than RPCholesky (0.0153 s vs 0.3237 s).

$\newline$

$\blacktriangleright$ $**Discussion on potential failure cases**$

• "Missing a discussion on potential failure cases of the method..."

We will expand our discussion:

• The supervised KT guarantees apply when t $f^\star$ lies in the RKHS $\mathcal{H}(\mathbf{k})$, where $\mathbf{k}$ is the base kernel. In practice, choosing $\mathbf{k}$ is indeed a challenge common to all prior work.
• Dwivedi & Mackey (2022), Cor. 1 provide integration-error guarantees for KT when $f^\star \notin \mathcal{H}(\mathbf{k}). Moreover, there are recent results on finding the best kernel (in testing; Dominigo-Enrich et al. (2023), Sec. 4.2).
• Combining these results, potentially with the recursive feature machines of Radhakrishnan et al. (2024) (also see our response to Rev PSUT) to build a practical strategy is an exciting future direction.

1. Thank you for the catch. RPCholesky-KRR achieves an estimated test loss of 0.9980 when $n=2^{14}$ (see Table B below)—close to the Bayes optimal loss of 1. We have estimated the population test loss by sampling 10,000 points from our data generating distribution (as a proxy for integrating), so the estimated test loss can actually be less than 1 with some probability. When we use log-scaling on the y-axis, negative values of excess risk become undefined, so they don't appear on the plot. Note that Full-NW and Full-KRR for $n=2^{14}$ also face this issue. We will make sure to add a note about this in the revision.
2. In Fig. 2, Full-NW and Full-KRR are the black box plots (always the first in each group when reading from left to right). Full is deterministic, so the box plots appear as lines. Note that Full-NW and Full-KRR indeed perform the best across all $n$. For additional clarity, we have attached tables of the test loss (mean ± standard deviation across 100 trials). (Subtract each value by 1 to get the excess risk values used in Fig. 2.)

Table A: Test loss for NW estimators in Figure 2a of the attached PDF

$\newline$

Table B: Test loss for KRR estimators in Figure 2b of the attached PDF

Thank you for your follow-up questions.

1. The values in Table A and B are test loss while the values in Figure 2 are excess risk. Hence a test loss of 0.9980 becomes an excess risk = test loss - 1 = -0.0020, which is not a valid value on log scale; hence the missing marker for RPCholesky in Figure 2.

2. Since the Tables are test loss, a lower value is better. So the opposite trend to what you wrote is true. In particular, in Table A, RPCholesky performs worse than our proposed method. On the other hand, RPCholesky is tends to do better than our proposed method In Table B. Overall the main takeaways from Table A/B (and Figure 2) is that our Kernel Thinning based estimator

   • improves upon standard thinning for both NW and KRR estimators across all coreset sizes
   • improves upon RPCholesky for NW estimator across all coreset sizes
   • improves upon RPCholesy for small coreset sizes for KRR, and remains competitive (but a bit worse) in large coreset sizes for KRR.

RPCholesky’s superior performance over KT for KRR is not too surprising since RPCholesky a spectral method tuned to perform better in such settings, while NW estimator does not benefit immediately from an improved spectral approximation.

We remark that these results are consistent with our original motivation of designing a unified framework for speeding up general non-parametric regression estimators using a single recipe.