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

TL;DR

This paper introduces Sketchy Moment Matching (SkMM), a scalable data selection method for finetuning that balances bias and variance reduction through gradient sketching and moment matching, with theoretical guarantees and empirical validation.

Contribution

The paper proposes SkMM, a novel two-stage data selection approach that efficiently balances bias and variance in high-dimensional finetuning with provable guarantees.

Findings

SkMM preserves fast-rate generalization independent of parameter dimension.

Gradient sketching is fast and provably accurate for bias control.

Empirical results show SkMM improves finetuning performance on vision tasks.

Abstract

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 $\mathcal{S}$; (ii) then the variance is reduced over $\mathcal{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 $\mathcal{S}$ preserves the fast-rate generalization $O(\dim(\mathcal{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.