NeurIPS2024
Sketchy Moment Matching: Toward Fast and Provable Data Selection for Finetuning
Yijun Dong, Viet Hoang Phan, Xiang Pan, Qi Lei
摘要
We revisit data selection in a modern context of finetuning from a fundamental perspective. Extending the classical wisdom of variance minimization in low dimensions to high-dimensional finetuning, our generalization analysis unveils the importance of additionally reducing bias induced by low-rank approximation. Inspired by the variance-bias tradeoff in high dimensions from the theory, we introduce Sketchy Moment Matching (SkMM), a scalable data selection scheme with two stages. (i) First, the bias is controlled using gradient sketching that explores the finetuning parameter space for an informative low-dimensional subspace S; (ii) then the variance is reduced over S via moment matching between the original and selected datasets. Theoretically, we show that gradient sketching is fast and provably accurate: selecting n samples by reducing variance over S preserves the fast-rate generalization O(dim(S)/n), independent of the parameter dimension. Empirically, we concretize the variance-bias balance via synthetic experiments and demonstrate the effectiveness of SkMM for finetuning in real vision tasks. * Equal contribution. 2 Throughout this work, we refer to "low-dimension" as the setting where the number of finetuning parameters r is smaller than the selected downstream sample size n, while "high-dimension" refers to the opposite, r > n.