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

Summary

This paper addresses the problem of data selection for finetuning large pre-trained models. The key contributions are: 1. A theoretical analysis of data selection for finetuning that reveals a variance-bias tradeoff in high dimensions. 2. A provable result showing that gradient sketching can efficiently find a low-dimensional subspace that preserves fast-rate generalization. 3. A practical two-stage algorithm called Sketchy Moment Matching (SkMM) that uses gradient sketching to explore the parameter space and moment matching to exploit the low-dimensional structure. 4. Empirical validation on synthetic and real datasets demonstrating the effectiveness of the approach.

Strengths

1. The paper provides a rigorous generalization analysis for data selection in both low and high-dimensional settings. The proofs are detailed and appear sound. 2. The proposed SkMM method is simple to implement and scalable to large models/datasets. Experiments on both synthetic and real data demonstrate the effectiveness of the approach.

Weaknesses

1. Some of the theoretical results rely on assumptions (e.g., low intrinsic dimensionality) that may not always hold in practice. More discussion of the implications when these assumptions are violated would be valuable. 2. The method introduces new hyperparameters (e.g., sketching dimension, moment matching strength) without much guidance on how to set them optimally.

Questions

1. Does the approach extend naturally to other finetuning scenarios beyond linear probing (e.g., adapters, full finetuning)? 2. How does the computational cost of SkMM compare to other data selection methods as the dataset/model size increases?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

yes

Summary

The authors study the task of data selection. They extend the classical variance reduction to the high dimensional case and provide a variance-bias tradeoff analysis. Based on the theoretical results, they propose sketchy moment matching, which first utilizes gradient sketchy to form a low-dimensional space and then uses moment matching to reduce the variance.

Strengths

The proposal is a reasonable improvement over the baselines which often only consider bias or variance reduction. The theoretical analysis is also a decent contribution of the paper.

Weaknesses

The experiment focuses on linear probing, which already limits the scope of the evaluation. Furthermore, even under this limited scope, the setting does not seem to be challenging. For the synthetic setup, the sample count is 2000 while the rank is 2500, so it seems not to be a very high-dimension setup (the rank is not so much larger than the sample count). Also, the cluster count seems to be low for both tasks, 8 for synthetic, while the number of class is 10 for Cifar-10.

Questions

For the convenience of the readers, could you list the number of parameter fine-tuned for the Cifar-10 linear-probing task and the number of samples of Cifar-10? (I know this is something that can be looked up online, but these seems to be important numbers to show that the experiment setting is high dimensional) Also, could you scale up the synthetic experiment, like increasing sample count/rank/number of clusters? Testing on Cifar-100 is also another way to evaluate the performance on a larger number of clusters.

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

Yes, the authors discuss the limitations.

Summary

This paper concerns the data selection problem: given a collection of $N$ embeddings of dimension $r$ for $r\gg N$, the goal is to pick a subset $S$ of points of size $n$ so that one could run any downstream algorithm on $S$ with a regularization term, so that the empirical risk is small even on the entire finetuning set. Assuming the model is $y=\phi(X) \theta_*+z$ where $\phi: \mathbb{R}^d\rightarrow \mathbb{R}^r$ and $z$ is an i.i.d. noise vector with zero mean and bounded variance, then there exists a subspace that one could project onto and decompose the empirical risk as a bias and a variance term. Further, under the assumption that the second moment matrix has low intrinsic dimension, then one could find a good subspace via gradient sketching: draw a JL matrix $\Gamma\in \mathbb{R}^{r\times m}$ for $m\ll r$, then as long as one has $\Gamma^\top \Sigma^{\phi} \Gamma \preceq c_S \cdot \Gamma^\top \Sigma^{\phi}_S \Gamma$, then the error could be decomposed into a bias, variance and a sketching error term. A sketching gradient, moment-matching algorithm is proposed, involves applying sketching to the gradient, form the Jacobian and solve a quadratic relaxation. Experiments are performed on both synthetic datasets and CIFAR10.

Strengths

The main theoretical contribution is that for over-parametrized setting where $r\gg n$, one could provably show the existence of a subspace that one could project onto and perform data selection on that subspace. Moreover, if the second moment in addition has low intrinsic dimension, then one could use standard dimensionality reduction techniques (in $\ell_2$ norm) to sketch the high-dimensional gradient. In the sketchy moment-matching algorithm proposed in the paper, the authors first sketch the gradient then use uniform sampling to construct $S$.

Weaknesses

The core results of this paper are not technically very novel and surprising, the algorithm could be interpreted as a generalization of the leverage score sampling via JL trick due to Spielman and Srivastava, STOC'08. The analysis largely draws inspirations from the over-parametrization literature, which makes sense as finetuning is essentially training in an over-parametrized setting. Another point that is a bit unsatisfactory is the sketchy moment-matching algorithm utilizes quadratic relaxation to solve the program efficiently with projected gradient descent, but all analysis is based upon *not solving the quadratic programs*. The authors should try to provide some theoretical justifications of sketchy moment-matching, as that's one of the key contributions of this paper.

Questions

What is the runtime efficiency of your proposed method? It seems the performance is slightly better than ridge leverage score sampling, but ridge leverage score sampling could be implemented in input sparsity time, see the algorithm due to Cohen, Musco and Musco, SODA'17. Their algorithm is based on recursive uniform sampling, so could be implemented efficiently in practice.

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

Yes.

Summary

This paper studies the problem of data selection in the over-parametrized fine-tuning regime, i.e. when the number of fine-tuning parameters $r$ is larger than the amount $N$ of available examples. We want to subsample $n\ll N$ examples that form a representative set to train on, and hopefully achieve quality as close as possible to fine-tuning on the whole set. The idea is to compute the gradients $G\in \mathbb{R}^{N\times r}$ of all examples wrt the fine-tuning params and then select a subsample $S\subseteq [N]$ such that the Gram matrix of the gradients is approximated: $c\cdot \Sigma_S := c \cdot G^\top I_S G \approx G^\top G := \Sigma$. However, this is not possible to achieve since the model is over-parameterized. Fortunately, if the spectral approximation holds on a low-dimensional subspace of the parameter space, this is good enough, so the authors project the gradients on a random low-dimensional space. The proof goes through under the assumption that the singular values of the gradient matrix are well-concentrated on a small enough (<10%) support. The experimental results include fine-tuning on a synthetic linear task, as well as fine-tuning a vision transformer on CIFAR-10 image classification.

Strengths

- The authors study the data selection for fine-tuning problem from first principles - The writing is overall good and math looks sound, even though I didn't check details. - The experimental results look promising since SkMM beats a variety of algorithms including leverage scores. - The idea of spectral approximation on a subspace of the parameter space is interesting.

Weaknesses

- Important details on the experimental setup are missing or unclear. Specifically, what is the optimization process after the data is subsampled? For the image classification experiments, what is being fine-tuned, is it all the ViT parameters? For how many epochs? - The algorithm requires computing the gradients of all samples, which can be computationally expensive. Besides, if we are computing all gradients, why can't we just train one epoch on all datapoints? Why is data selection useful in this case? - The literature review could be expanded, including relevant papers such as BADGE [1], Coreset-based sensitivity sampling [2]. - In the experimental results, the authors should also compare with margin sampling (in addition to entropy sampling), as well as uniform sampling for the image classification task. - Computing the moment-matching subset in Algorithm 3.1 seems overly complicated, see questions [1]: Deep Batch Active Learning by Diverse, Uncertain Gradient Lower Bounds [2]: Data-Efficient Learning via Clustering-Based Sensitivity Sampling: Foundation Models and Beyond

Questions

In Remark 3.2 the authors write that their goal is to achieve the constraint $\tilde{G}^\top \tilde{G} \preceq c_S \cdot \tilde{G}^\top I_S \tilde{G}$ (1). They subsequently relax this problem and solve the resulting constrained convex optimization problem using projected gradient descent. However, it seems to me that (1) might be equivalent to $U^\top I_S U \succeq c_S^{-1} \cdot I$, where $U\Lambda^{1/2} V^\top$ is the SVD of $\tilde{G}$. Here $U\in {N\times \bar{r}}$ is a tall and thin matrix. This is a spectral sparsification task which could be solved using leverage score sampling on the rows of $U$. Furthermore, this is the same as sampling examples proportional to the squared $\ell_2$ norms of the rows of $U$. Maybe, I'm missing something, so please correct me if I'm wrong.

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

N/A