TSDS: Data Selection for Task-Specific Model Finetuning

Thank you for your feedback. We address the concerns and answer the questions below:

W1 - connection to optimal transport

As mentioned in lines 38-44, we want the selected data to match the distribution of the representative data from the target distribution, and optimal transport is a powerful tool to measure distribution alignment. We will highlight this connection in Section 3.1 when we introduce the optimization problem.

W2 - discussion of LESS

We will add more discussion on LESS to related works. Specifically, LESS computes the gradient similarity between each candidate and all the query examples, and the maximum similarity is the score for ranking. Then the top-ranked candidates are selected. A major difference between our method and LESS is that our method matches the distributions, while LESS takes the top ones based on aggregated statistics which can make it focus on a particular set of query examples.

W3 - need M query examples

In practice, M can be a budget limit on how much human effort a user is willing to put into the selection process. We will expand our discussion in the conclusion section to include this point.

W4 - how to set parameters

Note that we only have two effective parameters ($C$ is a constant to make sure that the transport cost and $G(\gamma)$ are on the same scale). The way we set the hyperparameters is as follows:

• We set $C$ to 5 when the embeddings are normalized.
• We set $h$ to the maximum distance between 10 hand-crafted near-duplicates. The intuition is that the points within the distance of h will be considered as near-duplicates and the probability assigned to them will be reduced.
• $\alpha$ can be any value between 0.05 - 0.95, and the performance is not sensitive to it as long as it is not too small or too large (see Appendix E.3.1). In practice, we can also use a validation set and a small surrogate model to guide the parameter selection.

We will add the details to the appendix.

W5 - evaluation using different sample sizes for continued pretraining

We did experiments that selected 100K and 300K examples for domain-specific continued pretraining (see Table 2 in the PDF attached to the global response). The observation is similar to the 1M sample size: our method is either better than or comparable to the baseline. We will add this result to the appendix.

W6 - baseline systems

We compare with the current state-of-the-art for each setting. The discrepancy between our reported numbers and the numbers in the LESS paper is due to the updates to open-instruct. We followed LESS to use open-instruct for the evaluation, and we used the updated version. In fact, we discussed this finding with the authors of LESS and confirmed the discrepancy reasons using their old open-instruct scripts. We will make the evaluation script public so that people can reproduce our numbers.

We believe that you have missed key points of our work and we would like to correct several factually incorrect points in the provided review.

W1 - novelty

To the best of our knowledge, our work is the first one that presents a unified framework for task-specific data selection which considers distribution matching, diversification, and deduplication. In addition, our formulation is general, supporting any metric spaces and distance functions where efficient nearest-neighbor search algorithms exist.

W2 - extension to VLMs

We disagree that the focus on language models poses a weakness of our work. For task-specific data selection, it is common practice to use large language models for the evaluation [1][2]. In addition, we already show the effectiveness of our framework in diverse experimental settings. Evaluation using vision-language models is an interesting extension that can be studied in future works.

W3 - performance of random selection

This statement is factually incorrect. Our method is significantly better than random selection with large gaps. In Table 2, we can see that our method consistently outperforms random selection, with a gap of 3.2 F1 points on average for task-specific finetuning. In Table 4, we show that our method consistently outperforms random selection, with a gap of 2.9 F1 points on average for domain-specific continued pretraining.

W4 - finetuning cost

Finetuning on the full dataset takes 80 hours on an A100 GPU with 40G memory. The cost of finetuning on the selected set is based on the selection ratio. For example, if the selection ratio is 1%, the finetuning time is 80 * 1% = 0.8 hours. We will add the details to the Appendix.

[1] Xia, Mengzhou, et al. "Less: Selecting influential data for targeted instruction tuning." International Conference on Machine Learning. PMLR, 2024.

[2] Xie, Sang Michael, et al. "Data selection for language models via importance resampling." Advances in Neural Information Processing Systems 36 (2023): 34201-34227.

Thank you for the feedback. We address the concerns and clarify the questions as follows:

W1 - ablation study using embeddings

We proposed a framework where the data can be embedded in any metric space with any distance function that supports efficient nearest neighbor search. Finding the best choice of embedding and distance is orthogonal to our study. Our choice of using gradients for task-specific finetuning is motivated by LESS [1], a state-of-the-art method for task-specific finetuning at the time of submission, which shows that similarities of gradients are better approximations of influence on task-specific performance than similarities of embeddings. Based on their observations, we further show that applying our framework using exactly the same encoding and distance function yields even better results due to our distribution matching and diversification. For the experiments on domain-specific continued pretraining, we use sentence embeddings and also show strong results.

W2 - additional methods

Thanks for the pointers. These methods target different settings in the broader data selection space and we find them not directly comparable to the setting we consider, i.e., task-specific data selection where not all candidates are relevant. We will add a discussion on these techniques in related works, as their diversification methods can be of interest to the reader.

W3 - hyperparameter setting

We follow the convention of the data selection literature [1, 2, 3] to use the same hyperparameters in the evaluation. We use the hyperparameters recommended by previous works (see the experimental settings in Section 5), and they are not tailored to our selected data. Given that our methods achieve consistently good performance, it is unlikely that our method outperforms the others due to a more favorable set of hyperparameters.

Q1 - intuition behind uniform transport

Thank you for the suggestion. We will add a discussion after L140 to highlight that uniform transport represents the most diverse case where every example receives the same amount of probability mass. Therefore, we use it as a reference point to encourage the distribution to be close to it. We can gain more intuition from Theorem 3.1, where we will add a discussion stating that as we increase the weight of the regularizer, the optimal K also increases.

Q2 - why use optimal transport instead of presenting the method as knn

The key innovation of our framework goes beyond the kNN algorithms themselves. Our framework considers two essential objectives in data selection: distribution matching and diversification [4]. By formulating an optimization problem, we integrate these goals while also taking near-duplicates into account. The formulation also allows us to study the optimality of the problem and have optimal solutions (Theorem 3.1 and 3.2). In our framework, kNN serves as a tool for finding the optimal solutions.

Q3 - other efficient optimal transport algorithms

Thanks for the suggestion. Incorporating algorithms like Sinkhorn for even further efficient computation is an interesting direction. We will add a discussion to related works.

Q4 - effective dataset size

• We sample from the computed distribution with replacement every step. Therefore, the total number of unique examples can be up to 4x the number of examples used per epoch, though the actual number of unique examples is much lower since examples with high probability mass tend to be repeatedly sampled. LESS is a deterministic selection method, and we follow the setting in their paper to select a fixed set and use the selected set for 4 epochs. Then for fair comparison, when evaluating our method, we train for the same number of steps with the same batch size.
• It's important to note that having more unique examples does not necessarily lead to better results, since many candidate examples may not be relevant to the task. For instance, the "Full" method, which has seen all candidate examples, performs worse than both LESS and our method in most cases.
• To demonstrate that the number of unique examples during training is not the primary factor behind our performance gain, we provide additional results that compare our method with LESS when the number of epochs is set to 1 (see Table 1 in the PDF attached to the global response). Specifically, each method selects a set whose size is 4% of the candidates. Then we train the model on the selected set for 1 epoch (the amount of computation is the same as using 1% for 4 epochs). From the results, we can see that our method still outperforms LESS in 5 out of the 6 settings when LESS has a larger “effective dataset size”.

We will add a discussion with the additional experiment to the appendix.

[1] Xia, Mengzhou, et al. "Less: Selecting influential data for targeted instruction tuning." International Conference on Machine Learning. PMLR, 2024.

[2] Xie, Sang Michael, et al. "Data selection for language models via importance resampling." Advances in Neural Information Processing Systems 36 (2023): 34201-34227.

[3] Coleman, Cody, et al. "Selection via proxy: Efficient data selection for deep learning." International Conference on Learning Representations (2019).

[4] Albalak, Alon, et al. "A survey on data selection for language models." arXiv preprint arXiv:2402.16827 (2024).

We appreciate the feedback. Here is our response to the questions and concerns:

W1 - other LLMs

Thank you for your suggestions. We chose llama-7b and mistral-7b since they achieved state-of-the-art performance in various tasks at the time of submission among the 7b-size models and are shown to be effective with LoRA finetuning. Our experiments in Section 5.1 demonstrate that the benefits of our data selection solution hold across both models. A detailed study of other models could be explored in a future, extended version of this work.

W2 & Q1 - runtime

On the same 150M examples (used in the experiment in Section 5.2), DSIR takes 18 hours to initialize. Note that this is a one-time cost for the data repository and the index we build can be reused across tasks. Although our method requires an additional 10 hours to initialize compared to DSIR, it achieves an average gain of 1.92 F1 points when provided with 1K representative examples from the target task. For the task-specific selection after initialization, our method takes 0.11 hours and DSIR takes 0.13 hours when taking 1M examples guided by 10K query examples. The task-specific selection after initialization scales linearly with respect to the number of query examples, and is not affected by the number of examples we want to sample (except for the IO cost). For the experiment in Section 5.1, the initialization time for both LESS and our method is 54 hours, dominated by the gradient computation time for the 270K candidates. For each task-specific selection, both take less than 1 minute. We will add the discussion above to the appendix.