Sketchy Moment Matching: Toward Fast and Provable Data Selection for Finetuning

We would like to thank the reviewer for their insightful questions and suggestions. We are glad that they found our theory solid and our method effective. On the questions raised in the review:

1. Low intrinsic dimensions of fine-tuning: Recalling the references [2,72] from the introduction, we highlight that the low intrinsic dimension of fine-tuning is an extensively observed phenomenon in practice with theoretical rationale. In particular, [2] demonstrates the surprisingly low intrinsic dimensions of fine-tuning language models in practice; while [72] provides a theoretical justification for the ubiquity of low-rank structures in natural data. One of the main goals of this work is to leverage such low intrinsic dimensions via data selection and enable learning with a sample complexity independent of the high parameter dimension $r$.

   As explained in footnote 7 (page 5), such a low intrinsic dimension is necessary for sample-efficient learning. Intuitively, if all the $r$ directions in the high-dimensional parameter space are equally important, when the coreset size $n < r$, the learned parameter $\theta_S$ must fail to capture the orthogonal complement of the space spanned by the coreset and lead to $\mathbb{E}[ER(\theta_S)] \gtrsim r-n$. We will further elaborate on these assumptions in the revision.

2. Choices of hyperparameters: The two hyperparameters in SkMM (Algorithm 3.1) are the sketching dimension $m$ and the constant $c_S$ that controls the moment matching strength.

   • Theorem 3.1 provides theoretical guidance for the choice of sketching dimension: $m = O(\overline{r})$ where $\overline{r}$ is the low intrinsic dimension. In practice, choosing $m$ to be a small constant multiple of $\overline{r}$ (eg, $m = 2\overline{r}$) is generally sufficient. (Notice that such pessimistic constants in theory compared to practice are ubiquitous in sketching, cf. [30,49,81] in submission.)
   • On the choice of $c_S$, we first recall from Remark 3.3 that the lower $c_S$ corresponds to the stronger moment matching, resulting in harder optimization of (5), but up on convergence, leading to a better generalization guarantee. In practice, we start by choosing $c_S$ close to one to ensure the solvability of (5) and then gradually decrease $c_S$ until the optimization of (5) fails or takes too long to converge. Specifically, we fix $c_S = 0.999$ in the synthetic experiments as reported in Sec 4.1 (line 264); while for experiments on real data, we explore $c_S = 0.99, 0.9, 0.8, 0.7, 0.6$ in the respective order for the minimum feasible $c_S$.

   We will include the above discussion on hyperparameter tuning (especially for $c_S$) in the revision.

3. Extension beyond linear probing: Thanks for the constructive question. In the revision, we will extend the empirical evaluation beyond linear probing. In particular, we demonstrate the effectiveness of SkMM in data selection for finetuning the last two layers of an ImageNet-pre-trained ResNet18 on two datasets: CIFAR-10 with 10 well-balanced classes and StanfordCars with 196 imbalanced classes. Please refer to General Response 3 for the results and details of the additional experiments.

4. Computational efficiency of SkMM: SkMM is efficient in both memory and computation while scaling well with the model and data sizes. Please refer to General Response 1 for a detailed discussion on the computational efficiency of SkMM.

We are happy to answer any further questions you may have. Thanks again for the helpful feedback.

We would like to thank the reviewer for their helpful questions and suggestions, and we are glad that they found this work well-presented and theoretically sound. Nevertheless, we believe there have been misconceptions regarding some key notions and the focus of this work. Before diving into the specific questions, we want to emphasize the following two main theoretical/algorithmic contributions of this work:

• a generalization bound of data selection for high-dimensional fine-tuning with a low intrinsic dimension that unveils the possibility of learning a fine-tuning model with a sample complexity proportional to the low intrinsic dimension (instead of the high parameter dimension), and
• SkMM, a practical data selection framework inspired by the analysis that finds such coresets efficiently via gradient sketching + moment matching.

Meanwhile, we recall the key notion of "high-dimensional finetuning" as explained in the introduction (footnote 1): the over-parametrized setting where the number of finetuning parameters $r$ is larger than the selected downstream sample size $n$. That is, “high-dimensional” refers to the relative magnitude $r>n$, instead of the absolute data/model size.

On the questions raised in the review:

1. Scaling up linear probing on synthetic data does not change the takeaway information: While we agree with the reviewer that more experiments on larger scales would intuitively be preferred, scaling up the synthetic data experiments is not a reasonable choice because (i) our current synthetic data experiments readily lie in the “high-dimensional”/over-parametrized regime with $n < r$ even though the absolute sample/model sizes are not huge; and (ii) scaling up linear probing experiments on synthetic data generally does not change the takeaway information, as shown below.

   Consider a set of $N=2000$ samples in a higher dimension $r=4000$ generated from a GMM with $\overline{r}=12$ clusters where ridge regression over the full dataset achieves $L(\theta_{[N]})=1.18e-2$. We show the results in the table below, which are consistent with the synthetic data experiments in the submission. The main takeaways remain unchanged: SkMM achieves among the best empirical risks across different coreset sizes $n$, especially for small $n$; while methods that facilitate variance-bias balance (SkMM and T/R-leverage) tend to outperform the rest baselines. (Notice that for $n \ge 120$, data selection can bring even lower empirical risk than learning over the full dataset, which coincides with the recent theory and observation in [41] in submission.)

   n	48	64	80	120	400	800	1600
   Herding	4.04e+03	4.04e+03	4.04e+03	3.87e+03	3.74e+03	3.67e+03	1.76e+03
   Uniform	(1.24 $\pm$ 1.19)e2	(0.81 $\pm$ 1.18)e2	(0.76 $\pm$ 1.11)e2	(0.74 $\pm$ 1.11)e2	(9.92 $\pm$ 0.18)e-3	(9.56 $\pm$ 0.03)e-3	(1.64 $\pm$ 0.39)e-2
   K-center	(6.64 $\pm$ 0.22)e-1	(4.72 $\pm$ 2.79)e-1	(3.98 $\pm$ 0.81)e-2	(1.71 $\pm$ 2.62)e-1	(3.12 $\pm$ 0.67)e-2	(1.29 $\pm$ 2.29)e-1	(2.04 $\pm$ 0.62)e-1
   Adaptive	(2.19 $\pm$ 1.27)e-2	(2.78 $\pm$ 3.34)e-2	(2.51 $\pm$ 2.02)e-2	(1.88 $\pm$ 1.63)e-2	(2.39 $\pm$ 4.25)e-2	(9.52 $\pm$ 0.05)e-3	(1.71 $\pm$ 2.72)e-1
   T-leverage	(1.06 $\pm$ 8.15)e2	(7.21 $\pm$ 8.25)e1	(5.44 $\pm$ 7.70)e2	(5.44 $\pm$ 7.70)e2	(8.27 $\pm$ 1.62)e-3	(8.66 $\pm$ 1.79)e-3	(1.09 $\pm$ 0.37)e-2
   R-leverage	(3.52 $\pm$ 7.63)e1	(1.92 $\pm$ 1.77)e-2	(1.19 $\pm$ 2.18)e-2	(1.03 $\pm$ 0.05)e-2	(9.48 $\pm$ 0.50)e-3	(9.54 $\pm$ 0.02)e-5	(1.17 $\pm$ 0.47)e-2
   SkMM	(2.13 $\pm$ 1.28)e-2	(1.65 $\pm$ 1.30)e-2	(1.11 $\pm$ 0.06)e-2	(1.03 $\pm$ 0.03)e-2	(9.34 $\pm$ 0.13)e-3	(9.41 $\pm$ 0.11)e-3	(1.08 $\pm$ 0.42)e-2

   To further investigate the effect of problem size on the data selection performance, instead of scaling the synthetic data experiments, we will provide more extensive empirical evaluations of our method on real data, as elaborated below and in General Response 3.

2. More comprehensive experiments on real vision tasks: Thanks to the constructive suggestion, in the revision, we will extend the empirical evaluation to

   • real vision tasks with a larger number of classes (ie, StanfordCars with 16,185 images in 196 imbalanced classes of cars, in addition to CIFAR-10 with 60,000 images in 10 balanced classes) and
   • fine-tuning settings beyond linear probing with parameters of much higher dimension (ie, finetuning the last two layers of an ImageNet-pre-trained ResNet18 with $r=2,364,426$ parameters, versus linear probing on CLIP with $r=5,130$ parameters, vide Table 5 in the attached PDF for detailed configurations). We will make sure to include these experimental details in the revision.

   Please refer to General Response 3 for the results and more details of the additional experiments.

We are happy to answer any further questions you may have. If our responses above help address your concerns, we would truly appreciate a re-evaluation accordingly.

We appreciate the insightful questions and suggestions from the reviewer. However, we believe there have been misunderstandings regarding the focus and contribution of this work. We hope that the following responses will help clarify these confusions.

1. Our theoretical contributions are explanatory instead of instrumental: As discussed in the related works, prior theoretical studies on data selection are conducted either in the low-dimensional linear regression setting or in the asymptotic regime, neither of which aligns with the fine-tuning in practice that manages to learn high-dimensional parameters with much fewer samples. (Although we analyze fine-tuning in the kernel regime following [83], the over-parametrization literature is generally out of the data selection setting.) To fill the gap, we provided

   • a generalization bound of data selection for high-dimensional fine-tuning with low intrinsic dimension in the non-asymptotic regime that unveils the possibility of learning a fine-tuning model with a sample complexity proportional to the low intrinsic dimension (instead of the high parameter dimension), and
   • a practical data selection framework inspired by the analysis, SkMM, that finds such coresets efficiently via gradient sketching + moment matching.

   Instead of proposing an instrumental theory, we use the existing theoretical tools from sketching and statistical learning to provide an explanatory theory for the sample efficiency of data selection for fine-tuning that generally admits low intrinsic dimensions in practice and an efficient method for selecting such data.

   While we agree that gradient sketching shares a similar high-level idea as the JL trick in Spielman and Srivastava, STOC'08 (which is ubiquitous in the sketching literature, cf. [30,81] in submission), leveraging such a fundamental theoretical framework should not compromise the novelty of an explanatory theory.

2. The quadratic relaxation is a practical realization of SkMM instead of the only solution: As elaborated in General Response 2, after gradient sketching, moment matching in the resulting low-dimensional subspace can be realized via various methods, including leverage score sampling and existing polynomial-time heuristics of discrete optimization for experimental design:

   • Our main motivation for exploring beyond leverage score sampling is the empirical observation that (truncated/sketched) leverage score sampling can perform poorly, especially in the low data regime (cf. Table 1). This can be explained by the dependence of the moment matching guarantee of leverage score sampling on the matrix coherence (vide General Response 2).
   • Discrete optimization for the V-optimality is an alternative to leverage score sampling. However, such optimization is known to be hard: the best available polynomial-time heuristic with a theoretical guarantee ([4] in submission) involves matrix inversion in each optimization step, bringing instability issues and compromising the practical application.

   Our proposal of SkMM can be interpreted as a quadratic relaxation of moment matching $\widetilde{\Sigma} \preccurlyeq c_S \widetilde{\Sigma}_S$ by assuming that $\widetilde{\Sigma}, \widetilde{\Sigma}_S$ commute (i.e., are simultaneously diagonalizable). This relaxation improves computational efficiency and avoids numerical instability by eliminating the need for matrix inversion. Although the assumption that the two matrices commute does not hold in general, it is a valuable heuristic for designing efficient algorithms, as demonstrated empirically across various domains. For example, in optimization, this approach is used to simplify certain SDP relaxations and approximate solutions effectively (Lu, Monteiro, 2005). In handling distributional shifts, it is used to design optimization formulations to obtain minimax risks efficiently (Blaker, 2000; Lei et al., 2021). We will further clarify this in the revision.

3. Efficiency of SkMM: As ridge leverage score approximation via recursive sampling, gradient sketching in SkMM can be computed efficiently in input sparsity time, and the following moment matching stage takes place in the low dimension $m$ (vide General Response 1). As discussed in General Response 2, while (ridge) leverage score sampling also facilitates moment matching as SkMM and enjoys a slightly better runtime in the moment matching stage (in a lower-order term), with experiments (cf. Table 1), we show that SkMM tends to provide better performance as it is tailored for optimizing moment matching.

We are happy to answer any further questions you may have. If our responses above help address your concerns, we would truly appreciate a re-evaluation accordingly.

References:

• Lu, Z., Monteiro, R. D. C. (2005). "Primal-dual interior-point methods for semidefinite programming using inexact step directions." Mathematical Programming, 103(2), 453-485.
• Blaker, Minimax estimation in linear regression under restrictions, annals of statistics 2000
• Lei et al, Near-optimal linear regression under distribution shift. ICML 2021

We would like to thank the reviewer for their constructive questions and suggestions, and we are glad that they found this work interesting and well-presented. On the questions raised in the review:

1. Cost of gradient computation and SkMM: First, we kindly emphasize the ubiquitous role of gradients in most data selection methods. For linear probing, the gradients (of the finetuning model instead of the loss) are essentially the last-layer features from the pre-trained model, which can be computed without backward propagation. Such information is necessary for all unsupervised data selection methods discussed in "1.1 Related Works" including leverage score sampling and influence function. Beyond linear probing, gradient computation can bring asymptotically similar costs as one epoch of training on the full data set. Notice that this is more efficient compared to methods with several epochs of warmup training (eg, [1,15,38,52,56,69,83] in submission, including all methods we compared against in the real-data experiments except uniform sampling and herding). Such cost of gradient computation or warmup training is acceptable because the selected data have their own importance (eg, finetuning similar models with limited memory and computation). Moreover, we refer to General Response 1 that SkMM based on gradient sketching is highly efficient in both memory and computation. As mentioned in Algorithm 3.1, the gradients can be computed and compressed (via sketching) in parallel and on the fly without storage, while all the subsequent steps for moment matching are conducted in the lower dimension $m \ll r$. We will make sure to clarify these points further in the revision.

   Second, we highlight that training for one epoch on the full dataset is generally insufficient for finetuning. For example, when finetuning CLIP via linear probing on CIFAR-10 (Table 2), training for one full epoch leads to test accuracy of only $91.94 \pm 0.17$ (cf. $92.96 \pm 0.07$ learned with as few as $1000$ selected data via SkMM).

2. Sketchy moment matching v.s. sketchy leverage score sampling: Up to proper scaling $U^\top diag(\frac{N}{n} s) U \succcurlyeq c_S^{-1}$ (where $s \in \{0,1\}^N$, $\|s\|_0 = n$), we agree with the reviewer that leverage score sampling based on the sketched gradients $\widetilde{G}$ can also provide good control over $c_S$ with high probability. However, such control comes with nuance dependence on matrix coherence and can render the upper bound vacuous in the worst case (vide General Response 2).

   This is exactly the major motivation for the comparison with the truncated and ridge leverage score sampling (T/R-leverage) in the synthetic data experiments. In particular, leverage score sampling based on the sketched gradients is effectively equivalent to T-leverage (ie, when the original gradient dimension $r$ is too high for the exact computation of SVD for $G \in \mathbb{R}^{N \times r}$, randomized SVD based on sketching is generally used as an approximation, in which case leverage scores of $\widetilde{G}$ is exactly T-leverage). In experiments, we also confirmed the indistinguishable behaviors between the two, and therefore only reported one of them. As discussed in Sec 4.1, while SkMM and T/R-leverage both encourage variance-bias balance (by controlling $c_S$) and therefore provide among-the-best risks in Table 1, SkMM shows better performance on average as it is tailored for optimizing $c_S$. Nevertheless, we agree that with very limited computation (ie, even optimizing (5) in low dimension $m$ is infeasible), sketchy (ridge) leverage score sampling would be a more affordable alternative to SkMM.

   Motivated by this insightful question, we will extend Sec 3.2 based on the above discussions on the connection between SkMM and leverage score sampling in the revision.

3. Improvements on related works, experiment setup, and more comprehensive experiments: Many thanks for the constructive suggestions on related works and experiments. We will include discussions on the suggested literature in the next version. In addition, we will provide more experiments on real data with more comprehensive baselines (including margin sampling) in the revision. (Please refer to General Response 3 for the results and details of the additional experiments.) We will also provide more detailed experimental setups in the revision as follows.

   • After data selection, we fine-tune the model in two different settings: (a) linear probing over CLIP-pre-trained ViT (vide Tables 1 and 3 in the attached PDF) and (b) non-linear finetuning of the last two layers in an ImageNet-pre-trained ResNet18 (vide Tables 2 and 4). Please refer to Table 5 in the attached PDF for detailed configurations like parameter dimensions.
   • The fine-tuning parameters are optimized via Adam with a learning rate $10^{-2}$, for 200 epochs over CIFAR10 (following [29] in submission, Tables 1 and 2) and 50 epochs over StanfordCars (Tables 3 and 4).

We are happy to answer any further questions you may have. If our responses above help address your concerns, we would truly appreciate a re-evaluation accordingly.

As the deadline for the discussion phase approaches, we would greatly appreciate it if the reviewer could provide some feedback on our responses, which we believe have addressed all the questions in the review. As a summary of important updates we made thanks to the valuable suggestions from the reviewer:

• In Response 1 and General Response 1, we provided further clarification on the computational efficiency of SkMM and comparison with other data selection methods.
• In Response 2 and General Response 2, we discussed the connection and difference between SkMM and leverage score sampling in terms of moment matching, where such difference comes from, and why we may want an alternative to leverage score sampling.
• In Response 3, General Response 3, and the attached PDF, we provided stronger empirical evidence on (i) more real vision tasks and (ii) fine-tuning settings (beyond linear probing) with much higher dimensions, with (iii) more comprehensive baselines (including margin sampling) and (iv) more detailed experimental setup.

We value the opportunity to improve our work based on these constructive suggestions and are always happy to provide further clarifications if needed. Thanks again for your time.