TSDS: Data Selection for Task-Specific Model Finetuning
This addresses the data selection bottleneck for practitioners finetuning models on specific tasks, though it is incremental as it builds on existing finetuning paradigms.
The paper tackles the problem of selecting appropriate training data for task-specific finetuning of foundation models, and shows that using their TSDS method with a 1% selection ratio often outperforms using the full dataset and beats baseline methods by 1.5 points in F1 score on average.
Finetuning foundation models for specific tasks is an emerging paradigm in modern machine learning. The efficacy of task-specific finetuning largely depends on the selection of appropriate training data. We present TSDS (Task-Specific Data Selection), a framework to select data for task-specific model finetuning, guided by a small but representative set of examples from the target task. To do so, we formulate data selection for task-specific finetuning as an optimization problem with a distribution alignment loss based on optimal transport to capture the discrepancy between the selected data and the target distribution. In addition, we add a regularizer to encourage the diversity of the selected data and incorporate kernel density estimation into the regularizer to reduce the negative effects of near-duplicates among the candidate data. We connect our optimization problem to nearest neighbor search and design efficient algorithms to compute the optimal solution based on approximate nearest neighbor search techniques. We evaluate our method on data selection for both continued pretraining and instruction tuning of language models. We show that instruction tuning using data selected by our method with a 1% selection ratio often outperforms using the full dataset and beats the baseline selection methods by 1.5 points in F1 score on average.