TSDS: Data Selection for Task-Specific Model Finetuning (2410.11303v2)

Abstract: 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 LLMs. 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.

Data Selection for Task-Specific Model Finetuning: A Technical Analysis

In contemporary machine learning, the finetuning of foundational models tailored to specific tasks has become a prevalent strategy. The efficacy of such task-specific finetuning hinges crucially on the appropriate selection of training data. The paper by Liu et al. addresses this challenge by proposing a novel framework for data selection that maximizes the efficiency and effectiveness of the finetuning process.

Core Contributions

The authors articulate a data selection paradigm for task-specific finetuning, framing it as an optimization problem. This is characterized by two primary objectives: distribution alignment and data diversity.

1. Distribution Alignment: The paper employs optimal transport as a measure for the alignment of the selected data's distribution with that of the task-specific representative examples. Optimal transport provides a robust metric for quantifying distributional discrepancies and ensures that the finetuned model can efficiently learn the target distribution.
2. Diversity: To avoid overfitting, the framework incentivizes diversity in the selected dataset. Kernel density estimation is introduced into the regularization, effectively mitigating the risks posed by near-duplicate data points inherent in web-crawled datasets.
3. Efficient Algorithmic Realization: The authors delineate an efficient algorithm that adapts nearest-neighbor search techniques for scalable data selection, overcoming the computational challenges posed by massive data repositories.

Experimental Evaluation

The framework is empirically validated across several natural language processing tasks, including both instruction tuning for LLMs and domain-specific continued pretraining. Remarkably, the proposed method, even with a mere 1% selection ratio, surpassed baseline full dataset training and established benchmarks by an average of 1.5 F1 score points. It also demonstrated robustness against data duplicates, maintaining stable performance even when a significant amount of near-duplicate data was present.

Theoretical Implications

This paper advances the theoretical understanding of data selection as it pertains to foundational model finetuning. The integration of optimal transport in the optimization framework represents a robust methodological innovation that bridges the gap between data characteristics and model performance. The authors’ decision to leverage both model-agnostic and model-specific metrics reflects a nuanced approach to finetuning that acknowledges the complexity of model behavior across diverse training regimes.

Practical Implications and Future Work

Practically, this work provides a scalable solution for real-world datasets, which are expansive and often plagued with redundancy. The computational efficiency of the proposed solution, taking merely 28 hours to preprocess and an hour to execute task-specific selections on a 150M-example corpus, renders it viable for industrial applications.

Looking forward, there's promising scope for further optimization in computational efficiency through variants such as Sinkhorn distances. Moreover, the framework's reliance on representative examples invites exploration into more autonomous methods for example generation or augmentation, potentially alleviating biases introduced during manual example selection.

In conclusion, Liu et al. deliver a sophisticated, yet practical approach to data selection for task-specific finetuning, offering both theoretical enhancements and actionable insights for the future development of AI systems. The acknowledgment of limitations and potential biases underscores the responsible and mindful approach taken by the authors toward impactful AI research.