TSDS: Data Selection for Task-Specific Model Finetuning
Authors: Zifan Liu
Amin Karbasi
Theodoros Rekatsinas
Venue: NeurIPS 2024 poster (2024)
Keywords: data selection
finetuning
foundation model
large language model
optimal transport
Primary area: generative_models
TL;DR: We present a framework to select data for task-specific model finetuning, guided by a small but representative set of examples from the target task.

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 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.

TSDS: Data Selection for Task-Specific Model Finetuning
Zifan Liu
Amin Karbasi Yale University New Haven, CT
University of Risconsin-Madison
Madison, WI
21iu676wwisc.edi
amin. amin.karassi karbasi
hattasiaaaaai@lleee..mm @yale. edu
Theodoros Rekatsinas
Apple
Zürich, Switzerland trekatsinas@apple.com
Abstract
Finetuning foundation models for specific tasks is emerging paradigm in mod- ern 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 small but representative set of examples from the target task. To do so, formulate data selection for task-specific finetuning as an optimization problem with distribution alignment loss based optimal transport capture the dis- crepancy between the selected data and the target distribution. In addition, we add regularizer encourage the diversity of the selected data and incorporate kernel density estimation into the regularizer to reduce the negative effects 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 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 1 % selection ratio often outperforms using the full dataset and beats the baseline selection methods by 1.5 points in F1 score on average. Our code available at https://gthuub.co////aa///////
Introduction
Finetuning foundation models 3 is the de-facto paradigm for building machine learning applications that focus specific tasks. Models such BERT (10] and I.LMMA (43) are large-scale models pretrained on massive unlabeled data across a wide range of domains. Those models can be specialized downstream tasks through finetuning. Finetuning can take a variety of forms depending on the target task. For instance, continued pretraining [17] extends the pretraining stage of a model on dataset that is more closely related to a target domain. As another setting, instruction tuning [51] trains generative foundation model on instruction-respone pairs to improve its performance in responding to spec-spec instructions.
Finetuning foundation models can lead to significant improvement in downstream tasks, but the effectiveness heavily relies the right choice of training data 117 repositories that one considers during training of generative models tend to be conside for
47 However, the data
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

example the use of Common Craw which contains 250 billion web pages, or The Pile [14 and hence, is impractical manually select the data that are distributed like the use cases in the target task. Therefore, automated task-specific data selection becomes critical.
In this paper, we propose TSDS (Task-Specific Data Selection), framework to select data for task- specific model finetuning. We consider the scenario of finetuning a foundation model to customize it for specific task characterized by a few representative examples. The input to our framework the representative examples and a massive repository of candidate data. Guided by the representative examples, we select training data from the repository for task-specific finetuning. We identify the following requirements for our framework.
(Distribution Alignment) First, the distribution of the selected data should match the distribution of the representative examples from the target task. Distribution alignment is essential for a model to learn the target distribution and enable data-efficient finetuning for the target task 40 Many works 38 examples. Such heuristics do not ensure distribution alignment between the selected data and the representative examples. A recent work [48] selects data by importance resampling to match the target distribution but is limited to an n-gram feature space, which cannot capture high-level semantics.
50 447 retrieve candidate examples that are most similar to the representative
(Diversity) Second, the selected data should be diverse so that the model can learn a wide range of related knowledge rather than overfitting specific examples. In practice, data repositories created by web crawling may contain a large portion of near- -uplicicss
[33 28 that can compromise diversity
and negatively impact model performance [28 19 For example, study 13) several snapshots of 2 and Common Crawl shows that 14% to 52% of the documents are near-duplicates.
Previous works (3 17 50 48 on task-specific data selection overlook near-duplicates, leading to the over-representation of such examples in the selected data. We require our framework to ensure diversity in selection even when a large fraction of the candidate examples are near-duplicates.
(Scalability) Finally, the selection algorithm should be efficient, considering the increasing scale of modern data repositories. The high volume of candidate data (e.g., 250 billion pages in Common Crawl) poses a great challenge to efficient selection.
Our framework formulates task-specific data selection as an optimization problem that allows a smooth trade-off between two crucial objectives: distribution alignment and diversity. The solution to the optimization problem is categorical distribution assigned to the candidates which we will sample from. In the optimization objective, we use optimal transport to measure the discrepancy between the distribution assigned to the candidates and the target distribution, encouraging the alignment between them. We show that the optimization problem admits efficient algorithms to compute the optimal solution. In addition, our framework supports distribution alignment in any metric space that supports efficient nearest- neiehst search, including model-agnosten semantic embedding and -specfip features such gradients.
Our contributions: 1) We formulate data selection for task task-specific finetuning as an optimization problem based on optimal transport for distribution alignment, with a regularization term that encourages diversity. 2) We make our framework robust to near-duplicates by incorporating kernel density estimation [36] into the regularization term. We show the connection between the optimal solution to the optimization problem and nearest neighbor search, which allows us to develop efficient algorithms employing approximate nearest-neighhor search techniques 23
We conduct extensive experiments validate the effectiveness of our framework. We focus on natural language processing tasks where foundation models have shown great advancements. We show that our framework beats the tate-of-the-au baseline 47 by 1.5 points in Fl score on average with a selection ratio of 1% on instruction tuning for two modern large language models on three tasks. In addition, continued pretraining using domain-specific data selected by our framework outperforms the other selection methods by up to 3 F1 points on four classification tasks from various domains. We also demonstrate that our framework is robust to near-duplicates in the data repository, maintaining consistent performance when 1% of the candidate examples are duplicate for up to 1,000 times. Our method is efficient, taking 28 hours to preprocess a corpus of 150M examples and less than hour for each task-specific selection.
httpe://commonrra///
attpe://lemuprrojeet.org

---

2 Background and Overview
In this section, we provide background information that is essential for the problem, followed by formal statement of the problem and an overview of our proposed framework.
2.1 Background
We introduce the notations that will be used throughout the paper and the optimal transport problem.
Notation We use Rz0 to represent the set of non-negative real numbers, and I2zu to represent the set positive real numbers. Let N be positive integer and we use [N] to denote the set of integers from 1 to N. We use bold letters to denote matrices and the corresponding plain letters with subscripts to denote the entries in the matrix. For example, 7 R xMM is matrix with size M N, and 7ay or Ji, is the entry in the 7hh row and the jth column (1-indexed).
Optimal Transport between Discrete Distributions We introduce the optimal transport problem, which forms the basis of our data selection framework. Let (A, f) be a metric space where A is finite set and f A x A + Rzu is distance function. Consider two discrete distributions j on U S A and on V [ A, where both U and V are finite sets, Let be the 1h example in U and H4 = u(u;)) be the probability of u,. Similarly, let vy be the jth example V and vj ((v)) ) be the probability of uj. Yet R =20
[U]x|V| be transport of probability mass between and v, where Wij is amount of probability mass transported from uy to vj. Assume that the cost of transporting one unit of probability mass from u, v, f(,,,)) the distance between u, and Vy. Optimal transport the problem of transporting all the probability mass from U to V with a minimal cost:
JU M
M Σ
min YER 20
Yof(u, 0y) subject to
Tis - Ma:
Σ
=
[[V]]
DIIIV
2.2 Task-Specific Data Selection Problem Statement
We now introduce the problem of data selection for task-specific finetuning. We assume access to set of M representative examples Q = J from the target task, which we call query examples. Consider a data repository = { j=1 containing N candidate examples. Note that Q and D are multisets that may contain duplicates. We aim to select B examples from the repository guided by the query examples. The selected examples will be used to finetune a model to tailor it to the target task. We adopt the model-agnostic formulation above for the generality of the solution. However, our framework can be applied to model-specific selection by using model-specific data representations; an example evaluation for model-specifie instruction tuning is presented in Section 5.1
N
2.3 Framework Overview
Our framework takes the candidate examples and the query examples as inputs and outputs a set of task-specific examples by the following workflow. 1. (Encoding) We first encode the query examples and the candidate examples into the same metric space with specified distance function. (Probability Assignment) We determine the probability mass assigned each candidate example by solving an optimization problem. 3. (Sampling) We take random sample with replacement from the candidate examples following a categorical distribution where the probability is determined by the assignment in the previous step.
3 Data Selection and Optimal Transport
Data selection for task-specific finetuning can be expressed as an optimization problem for proba- bility assignment to the candidates in the data repository. First, we discuss the formulation of the optimization problem and then show the existence of closed-form solutions. In addition, we propose a regularization term that addresses the problem of near-duplicates among the candidates. The proofs of the theorems in this section are provided in Appendix

---

3.1 Optimization Problem
Consider the metric space (2, where 2 = QU D contains all the examples and f: : 2 x 2 -
is a distance function. Let d E 18-00 MxN be the distance matrix, where dij = fags 2z) is the distance between the th query example and the j candidate example.
We propose an optimization problem that transports probability mass from the query examples to the candidates. The objectíve is a linear combination of a probability transport cost for distribution alignment and a regularization term to encourage diversity. Formally, given d E 20 RRMN,, we consider the following optimization problem, which we refer to as Problem RT (regularized transport):
M
N
N Σ
min N
a'
mydds +1 - -))(G))))))))))))))))) subject to =1
ViE [M]
YER
C
M
20
where C' 0 is scaling constant, E [0, 1] is hyper-ppane hyper-parameter that controls the trade-off between distribution alignment and diversity, and G is regularization function. The first term in Problem RT the cost of probability transport where Yey is the mass transported from the 7th query example the 5th candidate. Each query example has is probability mass to transport, stated in the constraint. The probability transport cost measures the cost of transforming one distribution to another by moving probability mass between them, providing method to quantify probability alignment. The second is a regularization term that encourages the diversity of probability transport.
Let > be an optimal solution to Problem RT example žy. which is the sum of the probability mass it receives from all the query examples. When we sample from the candidate examples in the subsequent step, 2j has probability pj.
We assign pj = - Mi probability to candidate
We propose two instantiations of the regularization term that encourage the diversity of probability transport by penalizing its discrepancy to the uniform transport:
Go (7) = M max hig the uniform transport.
captures the largest probability gap between y and
MN
GTv(7) Σ Σx Ia uniform transport.
is the total variation distance between y and the MN
We uniform transport reference point encourage diversity it represents the most diverse way of transporting the probability mass from one query example to all the candidates, assuming the candidates are distinct.
3.2 Closed-Form Solution
When G = Goc. Problem RT can be solved by standard linear programming techniques, but they run in ()(MN)) time, which is prohibitively expensive. Instead, we show the existence of closed-form solution that can be computed in O(M N log N) time (see Section for the algorithm).
Using Goo as the regularization function, we get an optimal solution by transporting the probability of each query example evenly to its K-nearest neighbors among the candidates, where K is determined by the tradeoff between distribution alignment and diversity:
Theorem 3.1. Given d E MR-M where N > 1, consider Problem RT with G(y) = Goo(T)- =
M MAX Mig - MN For all i [M], let ......i%x be a reordering of [N] such that dist S S dijz. Consider x E MRMN whose entries are KM if E .ii....,.h K& and 0 otherwise. where K = max(k [N11 EM
G-n Et E dof)) <1 a) Assume K < N/2, and then 7" is the unique minimizer if a MM 1 des))
M
" is minimizer of Problem RT
- (1-a)M and Di E [M] such that digk dijk
Similarly, there exists a closed-form solution that can be computed in O(N log N) time when G = GTv (see Appendix A for the solution and the algorithm).

---

3.3 Addressing Near-Duplicates via Kernel Density Estimation
When there exists a large fraction of near-duplicates among the candidates, Go fails to charac- terize the diversity of probability assignment since it treats duplicates as distinct examples. Consequently, the contents the near-duplicates will be over-sampled. For example, 100 of the K-nearest neighbors of a query example are duplicates and the others are distinct, the content in the duplicates will receive 100 times as much probability mass as any other example.
To address the near hear-duplicate problem, we propose a regularization function incorporating kernel density estimation (KDE) [36], which a non-parametric method estimate the probability density function from finite examples. We determine the duplication level of a point by the kernel density estimate at its position. We use the Epanechnikov kernel such that given D, the density estimate at
point is ExrED max(1 f(z,²'² O), where h > 0 is the kernel size and f is the distance function. For example, for a point z: in D whose distance to any other point is larger than h. the density estimate at is 1. If we create two duplicates of z and add them to D, the density estimate at r increases to 3.
h2
Our KDE-based regularization function is GKDE(r) M (MX Pa M ME 1/py'
1/0:
where P3 EI x")) is the density estimate at z.z- G compares y to the probability assignment that is proportional to the inverse of the density, and penalizes the largest gap weighted by the density. Note that Goo a special case of GKDE(") with P; for all j E [N]
The optimal solution to Problem RT when G = GRDE can be obtained by assigning the probability mass of each query example to the nearest neighbors among the candidates, weighted by the inverse of their density estimate, as is shown by the following theorem.
N and PI.....PN E Rzo, consider Problem RT with G(r) = Theorem 3.2. Given d E R220 MxN
GXDE(T) = MM maXiE(M))E[N] Polaks - M EyEEN 1/p9" For all i E [M], let I .N be a reordering of [N] such that disi S.... S...S &igj. Let 2 = E 1/pz3: and be a discrete variable that takes value from S Es [M].k [N]] [0). Les c(8)
1/p;
EM Di=i (((s), where (s) = if S any k 2 2. Let max(s (8 E (lgc(() < S s and (x)) = -1-1 4ss Aus Toh_
--1
Pof
(1 a)M), and K; max(k (0, N -1)1s S Assume 5 S Σ5 1/pj. and then
Y is minimizer of Problem RT where Vi E [M], [N]
1/(Ms* Pst); 1
if K;
Ton
1/(Ms* Pst)) M -
if = +1
0,
otherwise
" is the unique minimizer if Ass E S such that ge(s) -1- M and 7i E [M] such that disk, = =dijk.ts or diskk,.1
Intuitively, we count candidate v. as 1/p, examples. For each query example, the optimal solution assigns probability mass to the candidates in its neighborhood proportional to their adjusted counts. The size of the neighborhood is determined by the limit s* on the sum of the adjusted counts.
In Figure D we show an example comparing the optimal transport with G and GKDE- When G = Goc. the probability is transported uniformly to the candidates regardless of their relative positions. When G = GKDE, the clustered candidates receive less probability due to their high density, and they will be less over-represented when take samples according to the assigned probability.
4 Efficient Probability Assignment Algorithms for Data Selection
We propose efficient algorithms to assign probability mass to the candidates according to the optimal solutions to Problem RT] For G = Go and G GRUE- the corresponding algorithms are KNN- Uniform (Algorithm D and KNN- KDE (Algorithm 2 Each algorithm takes the query examples and the candidates as input and outputs the probability assigned to each candidate.
Both algorithms prefetch the nearest neighbors of each query example from the candidates as the first step, where L is limit the neighborhood size. Specifically, GETKNN(O, Q, D) returns the
5

---

25
222
(a) Coo
(b) CKDE
Figure 1: An example of the optimal probability transports under different regularization terms. We consider query example y and 5 candidates Z1,-... , I5 embedded a 2-dimensional space. Assume that the candidates that form cluster (i.e., 23, 24, 25) have density estimate of each and the others have a density estimate of
Algorithm 1: KNN-Uniform.
Input: query examples Q {4 == candidates D {0, )- number of nearest neighbors
to prefetch L, a [0, 1]. C > Output: ,I....,P'i
j. GETKNN D, L):K + 1;
3 while K < L and C E Ii - dal < (1 - M do
K K 1;
5 foreach j E [N] do
Pj 0:
7 foreach i E [M] do
foreach k E [K] do
Pisk f Pjs KM
indices j E NMXx of the nearest neighbors and the corresponding distances d RM where Ja is the index of the htth nearest neighbor of q: in D, and da is the distance between qi and F Retrieving nearest neighbors exactly requires computing the distance between every query example and all the candidates, which is inefficient when the candidate size N is in the order of millions and billions. Alternatively, we can employ approximate nearest search techniques [23 [6] to improve efficiency at the cost of lower accuracy.
Then the algorithms assign probability mass to the nearest neighbors of each example. KNN-Uniform determines K based the tradeoff between distribution alignment and diversity. Then the algorithm assigns the probability mass of each query example evenly to its K -ncarest neighbors. KNN-KDD assigns probability mass to the nearest neighbors proportional to the inverse of their kernel density estimates (Line 18). The sizes of the neighborhoods are determined by Line 7 where we increase the limit on the sum of the inverse of the density estimates the neighborhood until the condition on Line 9 is satisfied. We use a priority queue to store the possible values s can take and retrieve the smallest one in each iteration.
In KDE-KNN, we also precompute the kernel density estimate for the I--naarst neighbors of each query example. To estimate the kernel density of each candidate example, we need to compute the distance between and all the other candidate examples. To reduce the computational cost, we use the -nnarestt neighbors among the prefetched examples the set compute KDE for each candidate example. Let D' be the set containing the L-nearest neighbors all the query points and Nr, be the
1-nearest neighbors of z in D'. We compute the KDE of example
ΣzEN, (1 f(x,x))
KNN-Uniform runs in O(ML + T1) time, and KNN-KDE runs in O(ML log M + TT) time, where Ti is runtime GETKNN, and T2 is the runtime of DEMPUTEKDE. With exact nearest neighbor search, Ti O(MN log N) and T2 O(M2 log( M L)). If we employ approximate nearest neighbor search techniques such as HNSW 34 for real vectors and l distance, we have
T = O(M N) log .) and T2 O(M log ML).

---

Algorithm 2: KNN-KDE.
1 Input: query examples @ = {q: (i)) candidate examples D {zy]} number of nearest
neighbors to prefetch L > I, a E [], C 0; Output: pl....,PN:
2 J, GETKNN 2, D, L);
3 p COMPUTEKDE(j].. D) /* p RxxL and Pik is the density zj
4 H < EmptyPriorityueeul))
5 for E [M] do
K; 4 0; + 0: (ppush((1PP..)))
7 while H. is not empry do
8, < H.pop(:: K; +- + + 1; + Ki (k=i(d, - da.)/Po Pas
8
if & Σ ===1 & 2 (1-a)M a)M then
M
10
8* - 8: break;
if K; +
L then
12
H. L/PL..+1,
13
for J [M] do
Py f 0;
15
for i [M] do
for k É [Ki] do
n7
Piu - Pin + 1/(M pa));
Ek=1 1/(Ms* pa)):
P1. Pi
Table 1: Information of the target datasets for instruction tuning.
Dataset
Task
# Test Instances
# Query Examples
# Shots*
Metric
TydiQA 7
Multilingual QA
1,713
9
Fl score
MMLU 18
Multiple choice
18,721
285
Accuracy
BBH
Reasoning
920
81
Accuracy
# shots is the number of QA examples provided in the prompt when querying the model.
Experiments
We evaluate our framework data selection for task-specific instruction tuning and domain-specific continued pretraining, using different encodings as needed. We show that 1) our framework out- performs the ttate-of-the-art methods on data selection for task speciffe instruction tuning and domain-specinn continued pretraining by up to points and 3 points in F1 score respectively: 2) our framework robust duplicates, exhibiting consistent performance when 1% of the candidate examples are duplicated up to 1000 times, while baseline methods show a drop of 2 points in F1 score (see Appendix [.. 3) our method is efficient, requiring 28 hours to preprocess 150 million candidate examples and less than hour for each task-specific selection (see Appendix E.2
5.1 Evaluation on Task-Specific Instruction Tuning
We select training data to perform instruction tuning to tailor model to specific downstream tasks. We assume access to several query examples that represent the use cases of the target task and repository of instruction-respons pairs to select from. The detailed setting as follows.
Target Tasks, Model, and Data Repository We consider three tasks from standard benchmarks for language model evaluation. The properties are shown in Table T We use two models: LLAMA-2-
7B [3] and MISTRRII-7BB [22]. We use combination of FL.AN V2 [1. CoT 45, DOLLY [8] and OPEN ASSISTANT 226] as the data repository for selection, which contains 270K examples.
Encoding encode the examples using rescaled and randomly projected gradients from a LLAMA- 2-7B model finetuned on random 5% of the data repository. The encoding process follows Xia et al. [47, who show that gradient-bar encoding is essential capture the utility of training examples in instruction tuning. We use la distance the distance function. See Appendix(Cite the details.
7

---

Methods 1) Rand selects a random subset from the data repository; 2) LESS 47 of-the-art method on data selection for task specifciii instruction tuning) selects training data from the data repository based on their gradient similarity to the query examples; 3) Ours is the KNN- KDE instantiation of our framework with C 5, a = 0.075 and h 0.2. We discuss how we The implementation details of our method can also be found [ Note that our method not sensitive to the hyperparameters, as shown by the
(the state- choose the parameters in Appendix [
in Appendix microbenchmarks in Appendix B
Evaluation Protocol Following Xia et al. 47 we finetune the base model on the selected data for 4 epochs. The dataset size 0.5% 1.0% 5% of the data repository. Since our method is based probabilistic sampling, sample randomly from the data repository following the assigned probability. The hyperparameters for finetuning also follow Xia et al. E7 (see Appendix D). We repeat each experiment for three runs with different random seeds and report the mean and standard deviation.
do not select fixed training set. Instead, in each epoch
Table Performance of instruction tuning with dataset selected by our method compared with the baselines. The subscripts represent the standard deviations.
Model
LLAMA-2-78
MISTRAL.-7B
Dataset
TydiQA
MMLU
BBH
TydiQA
MMLU
BBH
Base Full
40.6
45.7
39.1
49.6
62.4
56.5
52.7
51.4
41.4
44.7
58.9
48.0
Rand 49.82.4
45.00.4
38.3p.s
57.01.3
59.5g.3
49.70.1
Ratio 0.5%
LESS
52.31.4
46.20.7
39.00.6
55.03.0
60.6o.5
53.0p.9
Ours
53.71.5
47.20.2
40.60.2
61.60.9
60.3g.9
55.01.7
Rand 47.81.7
45.90.5
38.20.7
57.80.4
59.40.2
53.71.0
Ratio 1.0%
LESS
54.01.0
48.30.2
40.20.6
59.00.s
61.10.1
53.71.8
Ours
55.40.5
47.90.2
42.01.1
63.61.4
60.5g.s
56.32.1
Rand 49.51.4
46.0o.s
40.80.6
57.60.7
60.2o.3
54.81.1
Ratio 5.0%
LESS 54.30.7
50.66.0
40.21.8
60.41.3
61.3p.s
53.70.6
Ours
54.31.0
50.9o.4
42.70.2
60.91.s
59.90.4
56.00.5
Results The results shown in Table 2 where "Base" is the base model without finetuning and "Full" is the model finetuned on the full data repository. Our method consistently outperforms the baselines TydiQA and BBH across different selection ratios, beating the state-of-the method (LESS) by up to 6 points. With selection ratio 1,, our method outperforms the full data repository on TydiQA and BBH. On MMLU, our methods show comparable results LESS. Note that for MISTRAL-7BB finetuning on the full repository leads to worse performance than no finetuning, which highlights the importance of careful data selection for task-specific instruction tuning. We also notice that finetuning MISTRAL-7B on any selected set does not increase its accuracy on MMLU. The reason could be that the base MISTRAL-7B model has already been well-tuned for multiple-choice questions using high-quality data. We observe drop in the performance of our method on TydiQA when the selection ratio increases from 1% to 5%, which may be caused by overfitting. We can early stop the training process to avoid overfitting in practice.
5.2 Evaluation on Domain-Specific Continued Pretraining
In this experiment, we select data for domain-specifer continued pretraining to adapt a model to a specific domain. We assume access to a set of annotated data for domain-specifer task that serves as query examples and a repository of unlabeled data to select from. We continue pretraining the base model on the selected data and then perform supervised finetuning using the annotated data.
Target Tasks and Data Repository We consider four datasets focused on classification tasks across diverse domains. The properties are provided Table 3 We select data for continued pertaining from data repository consisting of 150M sequences crafted by Xie et al. [48 from The Pile [4]
8

---

Table Training, validation, test sizes and the number of classes in the datasets.
Dataset
Domain
Train
Validation
Test
Classes
Metric
ChemProt 25
Biomedical
4,169
2,427
3,469
13
micro-F] score
IMDB 33
Movie review
20,000
5,000
25,000
2
macro-FI -- score
SCIERC 332
Computer science
3,219
455
974
macro-Fl score
AGNews 152
News
114,947
4,999
7,596
4
macro-FI score
Table 4: Fl scores of the downstream tasks. Standard deviations are shown in the subscripts.
IK Annotated Data
3K Annotated Data-
10K Annotated Data AGNews
ChemP. IMDB
SCL.
AGNews ChemP.
IMDB
SCL.
AGNews
IMDB
Base 69.61. 88.0p.4 60.12.3
87.10.2
77.11.1 88.7p.4 75.8
87.70.3
90.0g.u
89.1o.1
Rand 69.71.s 87.3 62.72.9
87.20.3
78.60.2
88.5..
77.51.s
88.2p.1
90.2p.1
90.2u.1
DSIR 74.8u.7
87.70.6
68.5g.1
87.40.2
82.20.1
89.4..2
78.90.7
89.1o.3
90.8p.1
90.1o.1
Ours 76.7o.6 89.80.1 72.1o.6
87.30.2
81.90.4
90.70.0 79.20.9
89.3p.1
91.60.1
90.70.1
Target-Domain Data Accessibility To simulate different levels of access to target-domain annotated data, we consider three settings with varying sizes of annotated data (1K, 3K, and 10K). When the size is set to M and the original target-domain training set is larger than M, we sub-sample it by choosing M examples uniformly at random without replacement.
Encoding We encode the examples into TR512 using the Universal Sentence Encoder to capture
semantic meanings and use 12 distance as the distance function.
Methods 1) Rand selects a random subset from the data repository; 2) DSIR 47
(the otate-of-
the-art art method on data selection for cemain--peccffe continued pretraining) selects examples by importance resampling to match the unigram and bigram distribution of the query examples.; 3) Ours the KNN-KDE instantiation of our framework with 5, a 0.6 and h 0.1.
Evaluation Protocol For each domain task, we provide the annotated set to the selection methods as the query examples to guide the selection. We perform continued pretraining on 1M examples selected by each method from the data repository for one epoch (see Appendix E.44 for different selection sizes), starting from the base ALBERT (27] model. Then we finetune the model on the domain-specific annotated set and evaluate on the test set. The hyperparameters for training follow previous works [17 49 48) (see Appendix D) The experiments are repeated five times with varying random seeds. We remove the best and the worst among the five runs to rule out outlier runs and report the mean and standard deviation.
The test Fl scores of the downstream classification tasks are reported in Table 4 As reference point, we provide the performance of finetuning the model directly without continued pretraining (Base). Our method outperforms the baselines in most cases except ChemProt (3K) and AGNews (IK), with a gap of up to 3 points in Fl scores. On ChemProt (3K) and AGNews (IK), our method is comparable DSIR. We also notice that method shows average improvement of 1.92 points over DSIR with an annotated set size of IK and 0.38 points with an annotated set size of 3K. This indicates that method is particularly effective with small annotated sets.
Results
6 Related Works
Task-Specific Data Selection Similarity basedd methods 39 the candidates, ranked by their similarity to the representative examples from the target task. The features used for similarity computation can be embeddings or ngrams for texts. Another line of works B5 48 use two generative models where one leams the distribution of the target-task data and the other learns the ceneral-purpose data. Model-specific data selection methods [2 47 choose data to maximize the model performance on the target task. Given the high cost of actually training model and evaluating it on the target task, these methods often estimate the model performance by approximation. DSDM [2] approximate the model performance using datamodels 221] a function
50
retrieves the top ones from

---

that maps the training data membership (whether each candidate is included in the training set or not) to the model performance. LESS [47 employs the influence function 244 to approximate the marginal gain on the model performance when including a candidate into the training set. 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.
Diversity Measurement for Data Selection Measuring diversity is a critical aspect of data selection, as it ensures that the chosen dataset represents a wide range of examples rather than being overly concentrated on similar or redundant instances. DEITA ( selects data in iterative manner, where the contribution of a new example to the overall diversity is measured by the clipped cosine distance between the new example and the closest examples that have been selected. QDIT 44) measures the diversity of the selected data using the facility location function that quantifies how well each example in the full set is represented by the selected set. Wang et al. 4 measure the diversity using the log determinant distance between the selected set and a reference set that is maximally diverse.
Data Deduplication Data deduplication removes duplicates or near-duplicates from a dataset. Exact duplicates can be detected using hash functions [1] 46), while the detection of near-duplicates is more challenging. Some works [37] identify near- duplicates utilizing
hashing [15]. Others 28 6 compute edit distances between examples find -upr-ddulica Another line of works relies on learned embeddings of the examples to detect near-duplicates.
7 Conclusion
In this paper, we proposed a framework for data selection for task task-specific model finetuning, based on optimal transport, which allows smooth tradeoff between distribution alignment and diversity. We incorporated kernel density estimation to make the selection robust to near-duplicates. Experimentally we showed that our method is effective in both task-specific instruction tuning and domain-specifi continued pretraining. A potential direction for future work is to incorporate more efficient variants of optimal transport, such Sinkhorn distances [9], to further improve the computational efficiency. One limitation of our framework is the reliance on a set of representative examples to guide the selection, which may not be easy to craft. The representative examples may also contain biases that can be exaggerated through the selection process, leading to negative social impacts. In practice, additional effort must be allocated to ensure the quality of the representative examples and the size of the representative examples needs to be decided according to the budget of human effort.
References
[1] ABBAS, A., TIRUMALA, K., SIMIG, D.,, GANGULI, S., AND MORCOS, A. S. Semdedup: Data-efficient learning at web-scale through semantic deduplication, 2023.
[2] AHARONI, R., AND GOLDBERG, Y. Unsupervised domain clusters in pretrained language models. In Proceedings of the 58th Annual Meeting of the Association for Computational Lin- guistics (Online, July 2020), D. Jurafsky, J. Chai, N. Schluter, and J. Tetreault, Eds., Association for Computational Linguistics, pp. 7747-7763.
[3] BOMMASANI, R., HUDSON, D. A., ADELI, E., ALTMAN, R., ARORA, S., VON ARX, S., BERNSTEIN, M. S., BOIGG, J., BOSSELUT, A., BRUNSKILL, E., ET AL. On the opportunities and risks of foundation models. arXiv preprint arXiv:2 07258 (2021).
[4] BUKHARIN, A., AND ZHAO, T. Data diversity matters for robust instruction tuning, 2024.
[5] CER. D., YANG, Y., YI KONG, S., HUA, N., LIMTIACO, N., JOHI,, R. S., CONSTANT, M., YUAN, S., TAR, C., SUNG, Y.-H., STROPE, B., AND
N., SHANARDO-CESPEDEE KURZWEIL, R. Universal sentence encoder, 2018.
[6] CHOWDHERY, A., NARANG, S., DEVLIN, J., BOSMA, M., MISHRA, G., ROBERTS,
A., BARHAM, P., CHUNG, H. W., SUTTON, C., GEHRMANN, S., SCHUH, P., SHI, K.,
TSVYASHCHENKO, S., MAYNEZ, J., RAO, A., BARNES, P., TAY, Y., SHAZEER, N., PRAB-
HAKARAN, V., REIF, E., DU, N., HUTCHINSON, B., POPE, R., BRADBURY, J., AUSTIN,
J., ISARD, M., GUR-ARI, G., YIN, P., DUKE, T., LEVSKAYA, A., GHEMAWAT, S., DEV,
10

---

S.. MICHALEWSKI, H.. GARCIA, X., MISRA, V., ROBINSON, K., FEDUS, L., ZHOU, D.,
IPPOLITO. D., LUAN, D., LIM, H.,, ZOPH, B., SPIRIDONOV, A.. SEPASSI, R., DOHAN, D..
AGRAWAL, S., OMERNICK, M., DAI, A. M., PILLAT,, T. S., PELLAT,, M., LEWKOWYC7, A., MOREIRA, E.,, CHILD, R., POLOZOV, O., LEE, K., ZHOU, Z., WANG, X., SAFTA, B., DIAZ, M., FIRAT, O., CATASTA, M., WEI, J.,
K., ECK, D., DEAN, J. PETROV, S., AND FIEDEL, N. Palm: scaling language modeling with pathways. J. Mach. Learn. Res. 24, (mar 2024).
[7] CLARK, J. H., CHOI, E.,, COLLINS, M., GARRETTE, D., KWIATKOWSKI, T., NIKOLARV, V., AND PALOMAKI, J. TyDi QA: A benchmark for information-seelie question answering in typologically diverse languages. Transactions of the Association for Computational Linguistics
8 (2020), 454-470.
[8] CONOVER, M., HAYES, M., MATHUR, A., XIE, J., WAN, J., SHAH, S., GHODSI, A., WENDELL, P., ZAHARIA, M., AND XIN, R. Free dolly: Introducing the world's first truly open instruction-tuned Ilm, 2023.
[9] CUTURI, M. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems (2013), C. Burges, L. Bottou, M. Welling, Z. Ghahra- mani, and K. Weinberger, Eds., vol. 26, Curran Associates, Inc.
[10] DEVLIN, J., CHANG, M.-W., LEE, K., AND TOUTANOVA, K. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume I (Long and Short Papers) (Minneapolis, Minnesota, June 2019), Association for Computational Linguistics, pp. 4171-4186.
[11] ELAZAR, Y., BIIAGIA, A., MAGNUSSON, I. H., RAVICUANDER A., SCIIWEN, D., SUIRR, A., WALSH, E. P., GROENEVELD, D.. SOLDAINI, L., SINGH, S., HAлSHIRZI, H., SMITH, N. A., AND DODGE, J. What's in my big data? In The Twelfth International Conference Learning Representations (2024).
[12] ENGSTROM, L., FELDMANN, A., AND MADRY, A. Dsdm: Model-aware dataset selection
with datamodels, 2024.
[13] FRÖBE, M., BEVENDORFF, J., GIENAPP, L., VÖLSKE, M., STEIN, B., POTTHAST, M., AND HAGEN. M. Copycat: Near-duplicates within and between the clueweb and the common crawl. In Proceedings of the 44th International ACM SIGIR Conference on Research and Development in Information Retrieval (New York, NY, USA, 2021), SIGIR '21, Association for Computing Machinery, p. 3398-24044
[14] GAO, L., BIDERMAN, S., BLACK, S., GOLDING, L., HOPPE, T., FOSTER, C., PHANG, J., HE, H., THITE, A., NABESHIMA, N., PRESSER, S., AND LEAHY, C. The pile: An 800gb dataset of diverse text for language modeling, 2020.
[15] GIONIS, A., INDYK, P., MOTWANI, R., ET AL. Similarity search in high dimensions via hashing. In Vidb (1999), vol. 99, pp. 518-529.
[16] GUO, R., SUN, P., LINDGREN, E., GENG, Q., SIMCHA, D., CHERN, F., AND KUMAR,
S. Accelerating large-scale inference with anisotropic vector quantization. In International Conference Machine Learning (2020), PMLR, pp. 3887-3896.
[17] GURURANGAN, S., MARASOVIĆ, A., SWAYAMDIPTA, S., Lo, K., BELTAGY, I., DOWNEY, D.. AND SMITH, N. Don't stop pretraining: Adapt language models domains and tasks. In Proceedings of the 58th Anmual Meeting of the Association for Computational Linguistics (Online, July 2020), Association for Computational Linguistics, pp. 8342-8360.
[18) HENDRYCKS, D., BURNS, C., BASART, S., ZOU, A., MAZEIKA, M., SONG, D., AND STEINHARDT, J. Measuring massive multitask language understanding. In International Conference Learning Representations (2021).
[19] HERNANDEZ, D., BROWN, T., CONERLY, T., DASSARMA, N., DRAIN, D., EL-SHOWK, S., ELHAGE, N., HATFIEL--DODDN Z., HENIGHAN, T., HUME, T., JOHNSTON, S., MANN, B., OLAH, C., OLSSON, C., AMODEL. D., JOSEPH, N., KAPLAN, J., AND MCCANDLISH, S. Scaling laws and interpretability of learning from repeated data, 2022.
[20] HU, E. J., SHEN, Y., WALLIS, P., ALLEN-ZHU, Z., LI, Y., WANG, S., WANG, L.., AND CHEN, W. LoRA: Low-rank adaptation of large language models, In International Conference on Learning Representations (2022).

---

[21] ILYAS, A., PARK, S. M., ENGSTROM, L., LECLERC, G., AND MADRY, A. Datamodels: Understanding predictions with data and data with predictions. In Proceedings of the 39th International Conference on Machine Learning (17-23 Jul 2022), K. Chaudhuri, S. Jegelka, 1.. Song, C. Szepesvari, G. Niu, and S. Sabato, Eds., vol. 162 of Proceedings of Machine Learning Research, PMLR, pp. 9525-9588
[22] JIANG, A. Q. SABLAYROLLES, A., MENSCH, A., BAMFORD, C., CHAPL,OT, D. S., DE LAS CASAS, D., BRESSAND, F., LENGYEL, G., LAMPLE, G., SAULNIER, L., LAVAUD, L. R., LACHAUX, M.-A., STOCK, P., SCAO, T. L., LAVRIL, T., WANG, T., LACROIX, T., AND SAYED, W. E. Mistral 7b, 2023.
[23] JOHNSON, J., DOUZE, M., AND JÉGOU, H. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data 7,3 3 (2019), 535-547.
[24] KOII, P. W. AND LIANG, P. Understanding black-box predictions via influence functions. In Proceedings of the 34th International Conference on Machine Learning- Volume 70 (2017),
ICML'17, 17, JMLR.org, p. 1895-
[25] KRINGELUM, J., KJAERULFF, S. K., BRUNAK, S., LUND, O., OPREA, T. I., AND TABOUREAU, O. Chemprot-3.0: global chemical biology diseases mapping. Database 2016 (2016), bav
[26] KÖPP, A., KILCHER, Y., VON ROTTE, D., ANAGNOSTIDIS, S., TAM, Z.-R., STEVENS, K., BARHOUM, A., DUC, N. M., STANLEY, O., NAGYFI, R., ES, S., SURI, S., GLUSHKOV, D., DANTULURI, A., MAGUIRE, A., SCHUHMANN, C., NGUYEN, H., AND MATTICK, A. Openassistant conversations democratizing large language model alignment, 2023.
[27] LAN, Z.., CHEN, M., GOODMAN, S., GIMPEI, K., SHARMA, P., AND SORICUT, R. ALBERT: lite BERT for self-supervised learning of language representations. CoRR abs//909. / 1942 (2019).
(28) LEE, K., IPPOLITO, D., NYSTROM, A., ZHANG, C.,, ECK, D., CALLISON-BURGH C.,, AND CARLINI, N. Deduplicating training data makes language models better. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers) (Dublin, Ireland, May 2022), Association for Computational Linguistics, pp.
[29] LIU, W., ZENG, W., HE, K., JIANG, Y., AND HE, J. What makes good data for alignment? comprehensive study of automatic data selection in instruction tuning. In The Twelfth International Conference on Learning Representations (2024).
[30] LIU, Y., OTT, M., GOYAL, N., DU, J., JOSIII, M., CIEN, D., LEVY, O., LEWIS, M., ZETTLEMOYER, L., AND STOYANOV, V. Roberta: A robustly optimized BERT pretraining approach. CoRR abs/1907. (2019).
[31] LONGPRE, S., HOU, L., VU, T., WEBSON, A., CHUNG, H. W., TAY, Y., ZhOU, D., LE, Q. V., ZOPH, B., WEL, J., ET AL. The flan collection: Designing data and methods for effective instruction tuning. arXiv preprint arXiv:2301. 13688 (2023).
[32] LUAN, Y., HE, L.,, OSTENDORF, M., AND HAJISHIRZI, H. Multi-task identification of entities, relations, and coreference for scientific knowledge graph construction. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing (Brussels, Belgium, Oct.-Nov. 2018). E. Riloff, D. Chiang, J. Hockenmaier, and J. Tsujü,, Eds., Association for Computational Linguistics, pp. 3219-3232.
[33] MAAS, A. L., DALY, R. E.,, PHAM, T., HUANG, D., NG, A. Y., AND POTTS, C. Learn- ing word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies (Portland, Oregon, USA, June 2011), D. Lin, Y. Matsumoto, and R. Mihalcea, Eds., Association for Computational Linguistics, pp. 142-150.
[34] MALKOV, Y. A., AND YASHUNIN, D. A. Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs. IEEE transactions patlern analysis and machine intelligence 42, 4 (2018), 824-836.
[35] MOORE, R. C., AND LEWIS, W. Intelligent selection of language model training data. In Proceedings of the ACL 2010 Conference Short Pupers (Uppsala, Sweden, July 2010), J. Hajič, S. Carberry, S. Clark, and Nivre, Eds., Association for Computational Linguisties, pp. 220- 224.
12

---

[36] PARZEN, E. On estimation of a probability density function and mode. The annals of mathe-
matical statistics 33, (1962), 1065-1076.
[37] RAE, J. W., BORGEAUD, S., CAI, T., MILLICAN, K., HOFFMAAN, J., SONG, F., ASLANIDES,
J., HENDERSON, S., RING, R., YOUNG, S., RUTHERFORD, E., HENNIGAN, T., MENICK,
J., CASSIRER, A., POWELL, R., VAN DEN DRIESSCHE, G., HENDRICKS, L. A., RAUH, M.,
HUANG, P.-S., GLAESE, A., WELBI,, J., DATHATHRI, S., HUANG, S., UESATO, J., MELLOR,
J., HIGGINS, I., CRESWELL, A., MCALEESE, N., WU, A., ELSEN, E.,, JAYAKUMAR, S., BUCHATSKAYA, E., BUDDEN, D., SUTHERLAND, E., SIMONYAN, K., PAGANINI, M., SIPRE, L., MARTENS, L., LI, X. L., KUNCORO, A., NEMATZADEIL, A., GRIBOVSKAYA, E., DONATO, D., LAZARIDOU, A., MENSCII, A., LESPIAU, J.-B., TSIMPOUKELLI M., GRIGOREV, N., FRITZ, D., SOTTIAUX, T., PAJARSKAS, M., POHLEN, T., GONG, Z.,
TOYAMA, D., DE MASSON D' AUTUME, C., LI, Y., TERZI, T., MIKULIK, V., BABUSCHKIN, I., CLARK, A., DE LAS CASAS, D... GUY, A., JONES, C., BRADBURY, J., JOHNSON, M., HECHTMAN, B., WEIDINGER, L.., GABRIEL, I., ISAAC, W., LOCKHART, E., OSINDERO, S., RIMELL, L.,, DYER, C., VINYALS, O., AYOUB, K., STANWAY, J., BENNETT, L.., HASSABIS,
D., KAVUKCUOGIU, K., AND IRVING, G. Scaling language models: Methods, analysis & insights from training gopher, 2022.
(38) RUDER, S., AND PLANK, B. Learning to select data for transfer learning with Bayesian optimization. In Proceedings of the 2017 Conference on Empirical Methods in Natural Language Processing (Copenhagen, Denmark, Sept. 2017), Association for Computational Linguistics. pp. 372-382.
[39] RUDER, S., AND PLANK, B. Learning to select data for transfer lcarning with Bayesian optimization. In Proceedings of the 2017 Conference Empirical Methods in Natural Language Processing (Copenhagen, Demmark, Sept. 2017). M. Palmer, R. Hwa, and S. Riedel, Eds., Association for Computational Linguistics, pp. 372-382.
[40] SHACHAF, G., BRUTZKUS, A., AND GLOBERSON, A. A theoretical analysis of fine-tuning with linear teachers. In Advances in Neural Information Processing Systems (2021), M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, Eds., vol. 34, Curran Associates, Inc., pp. 15382-15321 15394.
[41] SUZGUN, M., SCALES, N., SCHÄRLI, N., GEHRMANN, S., TAY, Y., CHUNG, H. W., CHOWDMERY, A., LE, Q. V., CHI, E. H., ZHOU, D... AND WEI, J. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprint arXiv:2210.0926/ (2022).
[42] TIRUMALA, K., SIMIG, D., AGIIAJANYAN, A., AND MORCOS, A. S. D4: Improving Ilm pretraining via document de-duplication and diversification, 2023.
[43] TOUVRON, H., LAVRIL, T., IZACARD, G., MARTINET, X., LACHAUX, M.-A., LACROIX, T.,
ROZIÈRE, B., GOYAL, N., HAMBRO, E., AZHAR, F., RODRIGUEZ, A., JOULIN, A., GRAVE, E., AND LAMPLE, G. Llama: Open and efficient foundation language models, 2023.
[44] WANG, P., SHEEN, Y., GUO, Z., STALLONE, M., KIM, Y., GOLLAND, P., AND PANDA, R. Diversity measurement and subset selection for instruction tuning datasets, 2024.
[45] WEL, J., WANG, X., SCHUURMANS, D., BOSMA, M., BRIAN ICITERR XIA, F., CIII, E. H., LE, Q. V., AND ZHOU, D. Chain of thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems (2022), A. H. Oh, A. Agarwal, D. Belgrave, and K. Cho, Eds.
[46] WENZEK, G., LACHAUX, M.-A., CONNEAU, A., CHAUDHARY, V., GUZMÁN, F., JOULIN, A., AND GRAVE, E. CCNet: Extracting high quality monolingual datasets from web crawl data. In Proceedings of the Twelfth Language Resources and Evaluation Conference (Marscille, France, May 2020), N. Calzolari, F. Béchet P. Blache, K. Choukri, C. Cieri, T. Declerck, S. S.Ggggi, H. Isahara, B. Maegaard, J. Mariani, H. Mazo, A. Moreno, J. Odijk, and S. Piperidis, Eds., European Language Resources Association, pp. 4003-4012.
[47] XIA, M., MALLADI, S., GURURANGAN, S., ARORA, S., AND CHEN, D. Less: Selecting influential data for instruction tuning.
[48] XIF, S. M., SANTURKAR, S., MA, T., AND LIANG, P. Data selection for language models via importance resampling. arXiv preprint arXiv:2302.03 (2023).
[49] YAO, X., ZHENG, Y., YANG, X., AND YANG, Z. NL.P from scratch without large-scale pretraining: A simple and efficient framework. CoRR abs/2111.04130 (202)).
13

---

[50] YAO, X., ZHENG, Y., YANG, X., AND YANG, Z. NLP from scratch without large-scale pre- training: A simple and efficient framework. In Proceedings of the 39th International Conference on Machine Learning (17-23 Jul 2022), K. Chaudhuri, S. Jegelka, I.. Song, C. Szepesvari, G. Niu, and S. Sabato, Eds., vol. 162 of Proceedings of Machine Learning Research, PMLR,
pp. 25438-25451
[51] ZHANG, S.. DONG, L., LI, X., ZHANG, S., SUN, X., WANG, S., LI, J., HU, R., ZHANG, T.,
WU, F., AND WANG, G. Instruction tuning for large language models: A survey, 2024.
[52] ZHANG, X., ZHAO, J. J., AND LECUN, Y. Character-level convolutional networks for text classification. CoRR abs// 0bs///09.00/26 (2015).
14

---

A Closed-Form Solution and Algorithm for Grv
When GTv. for each query example, we transport MN probability mass to any candidate example whose distance to the query example is less than (l-aCC plus the distance between the query example and its I-nearest neighbor. Then we transport all the remaining probability mass to the l-nearest neighbor of each query example.
Zcx
Theorem A.I. Given d E R-- where N > 1, consider Problem RT with G( -Gr(T))) =
20
DM Ej=1 higg MN| For all E [M], let Ii..... be a reordering of [N] such that dif .... 2 disk.- Consider 7" R where Vi [M]
if - -iis (1-a)C
Wk (2,... N}. Rog
MN
0,
otherwise
and
1
Yig
Ti
M
k=2
Then Yf is a minimizer of Problem RT (1-a)C and dajz * digg-
Y is the unique minimizer if E [M]v e [N], dist -di +
The corresponding algorithm is KNN-T (Algorithm 3p. KNN-TV assigns MN unit of probability mass the nearest neighbors that satisfy the distance condition in Line and Lnd the rest to the 1-nearest neighbor. KNN-TV has the same time complexity as KNN-Uniform.
Algorithm 3: KNN-TV.
1 Input: query examples = { candidate examples D { number of nearest neighbors to prefetch L, a E C > 0;
Output: Pl.-..;PN: pN;
3 j. +- GETKNN(O, D, L);
for j [N] do
5
Pj + O;
for [M] do
Pin Pii + M
k- 2:
while l < L and & [dik du)) < 51 a) do
10
Pju & Pja MN.
Pr. P/. MN
12
k + k + +1;
B Proofs
B.1 Proof of Theorem A.1
Proof. Let [(r) = 8 I/ Σnn nydy (al-a)Cae GTV((())) be the optimization objective, We prove the theorem by showing that for any o R M that satisfy the constraint (vi [M] ). L(r') 2 ((Y*).
Let n MXN be a probability transport such that 20
Keg 0,
dfg dusi
(1-a)C
Ak E (2....,N]
otherwise
15

---

We show that [(>) 2 [(r"). For any i E [M], let h; max f E [N]
-dist < [laa)C)
Then we have
M
N
M N
1 -
1
(())) - (6")"
(gg - kijdd +
(h/)
-h
0
2
MN
MN
1=1
i= j=
M
N
de - -7/3) + 2
a
(h1)
- h/
DI
MN
MN
M
N
des (oest - -m6)
(11bst
)1
MN
MN
N
(dish
- -dit 2692
+
1-
a
1
2
MN
MN
+1
T2
T2
N
0
(h69
- Wist
+
Pett
1
DI
2
MN
MN
P3
The last equation is due to the fact that M/s
0 for & hs and N/3 1et + Ex Yyy
Since dejt - duf 2 (1-a)C for any k > ks, we have Ti 2 (1 - a) Ex=k+11 Yag By the triangle equality, we have T2 - E - (-bbg) and T3 2 I-a N
2k=k (-1/g) Therefore, we have
Ti T2 + 73 2 0 and consequently L(>/) 2 ((r").
Let Khigh {2 < k < kal ha > MN and Kinn = (2 125k < k < M Let y E M
20
be a probability transport such that
if k E Kiig otherwise
vk E (2...., NJ....,,N.... N., !1
Then we show that L(7") 2 f(7").
Since
MN for l E Kiigh and " 7if
EnEKiNO (7652 MN
1 ), we have
M
((7")) ((T =
(dos) dost) (rat
1 MN) )
1-a a
MN |+
2
1=1
k Tigh
KE
Tss
1 a (h75
Σ
DI
2
MN
MN
MN
kE
To
Again by the triangle inequality, we have Ts 2 1-a
2 EEEK
Wigh
MN l. and therefore Ts To 2
1
O. Since we also have T4 2 0, it follows that L(r")) 2 L(r"").
Finally, we show that L(r"') [()*). Since Yog
- + EKERS 10
(1/9)) MN). ) we have
L(7")) -(r*))
(duy
dist) 17
1
- Σ h
1
+
C
MN
MN
K
kEK
Tr
Ts
a
-11
+
Σ
2
MNN
MN
MN
16

---

Since digh dist < [L-aلC
for any & E Kow we have T7 2 (1 a) Σ ZkEKi Mitt MN ). Notice
that Ta5 2 MN since Tafi M
اk=z
!! and for k {2..... N). Vaft I + We also Lnj
have 71 + ΣKERS (iigi MN z MN since Pijj MN. Therefore, we have Ts + Tg
(1 (1) a) ENERE, MN - 2/11 and T7 + Tg + Tg 2 O. The it follows that L()") 2 (-**). -ast
Thus, we have L(r)) 2 ((7") 2 _()") ) 2 -(7").
Next we show that if Vi E [M]Vk E [M], dest -ayi + [1-aIC and dif + dugs. Y is the unique solution. We consider two cases. In the first case where Vi [M] E [NV. dist
dist < (1-a)C we have Vi [M]vj E [M],n
MN. For any o # y", there must exist i E [M], k E [2. such that Ma > MN in which case T1 > 0 or Yegt MN in which
a
case T7 > (1 a-a) ELERE, (og)
1 ). In the second case where 3i E [M]k [N] such that MN
dyj d > (1-a]C, we have Ti > (1 a) Ex Yoe for that both cases, L(7)) > L(7))
and thus 7* is the unique solution.
B.2 Proof of Theorem 3.1 and Theorem 3.2
We show that Theorem 3.1 states special case of Theorem 3.2 Then we prove Theorem
3.2 and
follows that Theorem 3. holds as well.
We first show the connection between Theorem 3.1 and Theorem 3.2
In Theorem 3.2 when Pj 1
3.1 Then the optimal solution in Theorem 3.2 is also the same as the one in Theorem 3.1 if we substitute s* by and all the pj 's by 1.
for all j E []]. st = k and s* is the same as the K in Theorem
Let [(>)
Σ ΣSn Middy + (1 a GxDEBD())) be the optimization objective. We prove 3.2 by showing that for any 7 E R MRMMx that satisfy the constraint (Vi E [M] N Ej
Theorem
M) <(7 C(r").
For conciseness, we let d(a, = dist and Y(s,)) = Voj
We first show that c(s) a necddecrggi function. Since cs(s) is step function and
Σ 1==1
desetn-dea Ek-l k-1 dawn-dann >t= decat1) 2 0 for any k E [N - 1], cs()
Psj
Psf
non-decreasing. Therefore, (;(s) = D M
(s)s non-decreasing.
Let = max maxke]K,] Pst M(,,y consider the following two cases.
In the first case when r' S
1 M. we have
M
-(7) - -(7))
Σ
dann) - -((.)) + (1 (1-10000 ( aa)((kr(((7) - GKDE(Y"))
C
1=1
M
K.
a
d(c (7(.,) -((as) + + di.. -(0.k.11)
k=
n
(1 a)(GRDE(T') - GKDE(TT))
T2
17

---

Since vk 2 K; 22 2 d(,K,,)) and Ek=Ki+1 Mids - -o,K+1)) =- Ex=i( (76..) - -M(cs)
we have
M
K.
7
2
a
(1 -(a) + d(.,itti Σ
Has - ((kKt)))
k=
a
M
K2
+()) -((it)) (Piz M(.k) - Pas Pa)))
C
Pst
1=1
k=1
K2
+1) - -d(skk)
-r')
Pst
Ms
Let å = [M]] $ K1+1'' Then we have > s* and E K; k=1
dec.k, 11-4(m)) = c(3). Since c(s)
Psf
is non-decreasing, we have tc( 2 (1 M. Then it follows that Ti 2 (1 - a ( r').
Let 5 = -E 1/py. Given the assumption that s* S 3, we have Ms* 2 2 Ma' For Forany E [M]] for any & Kg we have p3 (a.) = Ms and for = K; + we have PJj (e.) = Ms (s* skx Pr S
Mr. ( sk M.r Therefore, MaX 1pst ((.,)) ME. Ms
Then
we have
To (1 - M) max kE[M],KE(N]
max (E[M
1pSE((N) ((A)
M8
M3
1
2(1 a.) max max Pit kE[M] K.
M8
M
M8
1
1
21 a)M(('
))
M5
Ms*
M5
2(1 - (1-a)
Ms
The last inequality follows the triangle inequality. Then it follows that Ti T2 2 0 and L(>) [(>") 2 0.
In the second case when r' > Mar. let K; max(K E (OJ S M) When M. Therefore, R2 s Kg Consider another probability transport
K > K; r'pp > 8'r > 1
7" E MxN where 20
TlPp M
if < K; if = R +1 otherwise
R. ER.
r' p'ipp k=
0,
Note that by the definition of K, we have (6,N))33 < for K, +
Then we have
N
- C (aa - -((m) + (1 a)(GkDE(7')) GKDE(>"))
M
d(s,as (esss((( - -(rx)) +
- (()R,1))) Σ
R=k++2
T3
a)(GRDE (T)) GKDE(TT)))
18

---

Since bk 2 K; 2, d(a 2 2d (1.Ri41) and Ek=kit1 (am) - -aakit1) - EK==n( - -r(a))
we have
M
R.
Ts
0
daas)laa M((s)) + do.Ritii Σ
Hals - ((,R
C
1=1
1m
k=R++1
a
M
R2
(s,
- -d(i))) (P5E7 - Psh. P(,)))
Pjt
20
In addition, since r' > Ms 2
1
2
and Ps ((,)) S r' for any i E [M] and k E [N] we
Ma
have maxiE(M])ME(() (4.K)
1 2 [ [ax Ipy Kam)) ME = r' Ms and
ME
MaXic(M].AE(N] 1Pt (f,N) Me
r'
ME. Therefore,
1 Ms
T4 (1 - M
max ke(M)
1ps
max 10 (E[M],E[[]]
M8
20
Then it follows that L(T') L(7") 2
Let S' = {s S| 1
< s S }*} and s ,....
..... ((S') be the elements in S' in the ascending order. Mrr and K:O) K, For t E [[S']] let K(") max E [M]]s} I s(1).
Mr
Let Le) = 7", (00) 1
consider the probability transport 7(1) E 20
MxN where
M 1/Pst
if S K(!)
(1)
-k=1 L/p Ma if K!") +
k!
(1,b)
0,
otherwise
Then we have
M
N
a
-7(,)) (1-1) ) (1-1)) (xxx((()(()))1)) - GKDE(T(()) (t)
= C
(T(
1=1
k=
T6
Ts
By the definition of k(!)) and s(1), either K.) = k((-1) or k(() k((-1) + 1. For any E [M] (t)
such that K-10) K- (-(t-1) we have ((,k) = 0 for h: > k((-1) + 1. For any i E [M] such that
K(") K (t-1) + 1, we have st k(t) s(t), in which case we also have 7(0) ((,2)) for K(--"+ +1 1.
Therefore, we have (t-1)
[L.K
(i.K(t-)44 111 +1) ((. = EK
(1)
x((-1) (t-1) (t)
((()))) (2,A))
Then it follows that
()
Ts
(du..) d0.ke (44))(((()))) ((te))
(e-1) (1)
((--))
a Σ Σ (d (dxx((-)44) d(a,))//ss
1
1
(
M s(t)
s(t-1)
Let g(0) min 8 k 41 and then we have T3- a 3
1
(1))
(((8))). Since
g00 < 8* and c(s) is non-decreasing, we have ge(((()) < gc(s") that T3 (1 a)( L0 ).r1)))
< (1 a:) M and then it follows
In addition, since s(1-1) < s(() ss we have Pzt p,,,7'' (t-1) > p3) 2
1 2
2 and further
To (1-
max Ips, (1-1) [E(M)
((.A))
max <E(M)
lp;; ((.A)) M8
(1) )
M3
1 1
=(1 a)( (t--1) (){((()))
s(t)
19

---

Therefore, we have ((y(t-1)) [(7((t) = T3 + To 2 0. Since o 2 7" 7(1) 2
y(15'1) = =y", we have y 2 7
If il 8 such that Ec(s) (1 a:) M and di [MM] such that dij,k.) - (1) or d(i,K,+1)
+2), we have Tj > (1 a)M N... r)) and T3 > O, and therefore 7 > y i.e. 7 the
unique solution.
C Implementation Details
In this section, we provide details of the implementations.
Implementation of Our Method For experiments in Section 5.2 GETKNN is implemented as two-stage retrieval. We first build a coarse Faiss [33 index for the data repository D and use it to retrieve the 2000 nearest neighbors of each query example. The retrieved examples form a new set D'. 'hen we build fine-grained index for D' and use to retrieve and return the 2000 nearest neighbors of each query example. COMPUTEKDE in KNN-KDE computes the kernel density of each example in T' by retrieving its 1000 nearest neighbors using the fine- fine-grained index. The coarse index LPQ06_11111111 112, IwF66536__HHHC _______NN,3222 PQ7+56, and the fine-grained index is IndexIVFFlat. We refer the readers to the Faiss documentation for the details of those indexes.
For experiments in Section 5.1 we use exact search for GETKNN to retrieve 5000 nearest neighbors of each query example and IndexIVFFlat for COMPUTEKDE.
Encoding Process for Instruction Selection We encode the examples following 47 using rescaled and randomly projected gradients from LLAMA-2-7B model finetuned on a random 5% of the data repository. Specifically, we finetune the base model on the randomly selected dataset for epochs and use the gradients from the checkpoint at the end of each epoch as the example encoding. The dimension of the projected gradient from each epoch is 8,192. We refer the readers to [47] for more details. Then for each example, we multiply the gradients from the 4 checkpoints by the corresponding learning rate and concatenate them to get the final encoding, which a 7766-6dmmmesiinll vector.
Parameter Selection Note that our framework only has two effective parameters (C a constant to make sure that the transport cost and G(y) are on the same scale). The way we set the hyperparameters as follows:
We set C' to 5 when the embeddings normalized.
We set h to the maximum distance between 10 hand-crafted near-duplicates, The intuition is that the points within the distance of will be considered as near-duplicates and the probability assigned to them will be reduced.
a can be any value between 0.05 and 0.95, and the performance is not sensitive to it as long as is not too small or too large (see Appendix E.
In practice,
validation set and small surrogate model guide the parameter selection.
D Hyperparameters of Finetuning
We apply LoRA 20 for parameter-efficien instruction tuning for the experiments in Section 5.1 The hyperparameters are shown in Table 5 We use an NVIDIA 1100 Tensor Core GPU with 40G memory for instruction tuning.
For the experiments in Section the hyperparameters for continued pretraining are provided in Table and those for supervised finetuning are in Table The hardware for continued pretraining and supervised finetuning an NVIDIA Tesla V100 GPU with 32GB memory.
E
Additional Experimental Results
20

---

Table 5: Hyperparameters for instruction tuning.
maximum token length
2048
batch size
128
epochs optimizer
AdamW
weight decay
0.0
Adam 31
0.9
Adam B2
0.999
Adam
1e-8
warmup ratio
0.03
learning rate scheduler
cosine
learning rate
2e-5
LoRA rank
128
LoRA
512
LoRA dropout rate
0.1
Table 6: Hyperparameters for continued pretraining.
maximum token length
256
batch size optimizer weight decay
128
AdamW
0.01
Adam 81
0.9
Adam 32
0.999
Adam
1e-6
warmup ratio
0.1
learning rate scheduler
linear
learning rate
5e-4
Table 7: Hyperparameters for finetuning. We set patience for early stopping to epochs so that finetuning stops when the validation F1 score does increase for 3 epochs.
maximum token length
256
batch size
16
epochs patience for early stopping optimizer weight decay
10
epochs AdamW
0.1
Adam 31 Adam B2
0.9
0.999
Adam
1e-6
warmup ratio
0.1
learning rate scheduler
linear
learning rate
5e-5
21

---

Fraction for Duplication: D.1
Fraction Duplication: 0.1%
Fraction Duplication: 0.1%
78.00
89.50
91.50
77.00
F1 Scare
89.00
Fl Score
76.00
91.00
Fl
75.00
88.00
90.50
87.50
10
100
1000
100
1000
100
1000
Duplication Factor
Duplication Factor
Dupli on Factor
Fraction Duplication: 1.0%
Fraction Duplication: 1.0%
Fraction Duplication: 1.0%
78.00
76,00
S.Srrre
89.00
Fl S.Ccore
91.00
88.50
FL
74.00
88.00
90.00
72.00
KNN-Uniform KNN KDE
87.50
89.00
10 Duplication Factor
100
1000
100
1000
100
1000
Duplication Factor
Daplicatt on Factor
Figure 2: Fl scores of the downstream tasks under different duplication settings.
E.1 Robustness to Near-Duplicates
We evaluate the robustness of the selection methods against near-duplicates in the candidate examples. We follow the same evaluation protocol described in Section[5.2 while injecting duplicates to the candidate examples. We set different levels of duplication by varying the fraction of examples chosen for duplication and the duplication factor (number of duplicates per example). The fraction for duplication is set to 0.1 1 / 1% and the duplication factor is set 10 / 100 1000. For example, the fraction for duplication is 0.1 % and the duplication factor is 10, we randomly choose 0.1% 11% of the examples from the data repository and duplicate each 10 times. We use ChemProt (IK), AGNews (3K), and IMDB (10K) to perform the analysis, where the numbers in the parentheses represent the sizes of the annotated data. We include KNN-Uniform with the same parameters as KNN-KDE to show the effectiveness of the KDE-based regularization.
The results show that KNN-KDE is the only method that is robust to all the duplication settings. We observe that under low duplication levels, specifically when (fraction for duplication, duplication factor) is (0.1%) %. 10), (0.1% %. 100), or (1%. 10), all the methods perform similarly to the case without duplication. Given that the injected duplicates constitute less than 10% of the data repository in those settings, it is not surprising that they do not have much effect the downstream performance. However, when the duplication factor is increased to 1000 with the fraction for duplication set to 0.1%, the performance of DSIR drops by 0.7 points on average, whereas KNN-KDE and KNN-Uniform retain their performance. Moreover, when the duplication factor is increased to 1000 with the fraction set to 1 1,, all the methods except KNN-KDE show notable decline (more than 2 points average) in their performance.
E.2 Runtime and Scalability
We report the runtime of our method that can be split into a pre-processing stage and a selection stage. We use a machine with an Intel(R) Xeon(R) Gold 5115 CPU @ 2.40GHz (40 cores) and 250GB RAM. The example embedding is computed using an NVIDIA Tesla V100 GPU with 32GB memory, while the other computations are on the CPU. In the pre- pre-processing stage, our method embeds the candidate examples in the data repository and further builds indexes for the embeddings. This stage takes 28.38 hours for the data repository Section 5.2 that contains 150M examples. In the selection stage, our method embeds the query examples, computes the probability assignment, and takes random samples according to the probability. This stage takes 0.7 hours for 10K query examples. The runtime of the selection stage scales linearly with the number of query examples and remains unaffected by the number of examples to be sampled except for the I/O cost. Note that while method takes substantial amount of time in the pre- pre-processing stage, the cost is one-time and
22

---

the index can be reused for a variety of tasks that require similarity search. In general, our methods are practical in terms of runtime.
E.3 Task-Specific Instruction Tuning for One epoch
In Section 5.1 we perform instruction tuning for 4 epochs, and our method takes a random sample in each epoch instead of using fixed set. 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. 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 I. Specifically, each method selects a set whose size is 4% of the candidates. Then we train the model on the selected set for epoch (the amount of computation is the same as using 1% for 4 epochs). The results are shown in Table 8 From the results, we can see that our method still outperforms LESS in out of the 6 settings when LESS has access more unique examples.
Table 8: Performance of instruction tuning with dataset selected by our method compared with the LESS. The dataset size is 4% of the candidate data repository and we train each model for one epoch on the selected set. The subscripts represent the standard deviations.
Model
LLAMA-2-7B
MISTRAL-7B
Dataset
TydiQA
MMLU
BBH
TydiQA
MMLU
BBH
LESS
54.40.0
46.50.9
40.41.3
60.51.6
60.8p.4 60.5g.8
55.81.5 56.32.1
Ours
55.40.5
47.90.2
42.01.1
63.61.4
E.4
Domain-Specific Continued Pretraining with Different Selection Sizes
We compare with the baselines when the size of the selected data is 100K and 300K for domain- specific continued pretraining while the size of the annotated dataset is fixed to 3K. The other settings are the same as in Section 5.2 The results are in Table D which show that our method is either better than or comparable the baselines.
Table 9: F1 scores of the downstream tasks when the sample size varies. The size of the annotated data is set to 3K. Standard deviations are shown in the subscripts.
100K Sample Size
300K Sample Size
ChemP. IMDB
SCI.
AGNews ChemP. IMDB
SCL. AGNews
Base 77.11.1 88.7 75.8
87.70.3
77.1 88.70.4 75.8
87.70.3
Rand 77.80.4
88.90.2 78.7
88.5p.a
78.2n.a
89.2o.2
78.8p.2 88.5p.2
DSIR 80.90.9 89.0n.a
79.91.n
89.0g.1
82.1o.3
89.40.3
78.2g.4
88.9n.2
Ours 80.4p.s 90.1o.1
80.5p.4
89.30.1
81.601 90.2o.s
79.80.2 89.2o.1
E.5 Miico-Bench
In this section, we provide micro- berrrobbencharkss that study the effects of the hyperparameters in our framework. In addition, we show the performance of KNN-TV, an instantiation of our framework that is not covered in the main experiments. The experiments focus on domain-specific pretraining, following the same settings as Section 5.2 The datasets and sizes used in the bicro-benchmmrks are ChemP (1K), AG (3K), and IMDB (10K).
E.5.1 Tradeoff between Distribution Alignment and Diversity
We study the effects a, the hyperparameter that controls the tradeoff between distribution alignment and diversity in our framework. We vary the value of a in KNN-Uniform and KNN-KDE and report the Fl scores of the downstream tasks in Figure all three datasets, we observe notable drop in Fl scores when 0 or 1, and consistent performance when the value of a: is set to other
23

---

ChemProt (1K)
AGNews (3K)
IMDB (10K)
91.75
78-
89.5
91.50
FL SLScrre
89.0
90.75
90.50
70-
88.0
90.25
ENN-Uniform
68
90.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.50
0.75
1.00
Figure Performance KNN-KDE when a varies. The error bar shows the standard deviation.
Table 10: Performance of KNN-KDE when the kernel size varies. Fl scores of the downstream tasks are reported with standard deviations shown in the subscripts.
Kernel Size
0.1
0.3
0.5
ChemProt (1K)
76.70.5
76.81.1
78.0o.3
AGNews (3K)
89.20.1
89.30.1
89. 29.2
IMDB (10K)
91.5p.1
91.40.2
91.80.0
values. Note that KNN-Uniform KNN-KDE is equivalent Uniform when O, and transports all the probability mass of each query example to its 1-nearest-neiihhor in the data repository when 1. The former does consider distribution alignment, while the latter results in overfitting to the 1-nearest-neighhors For the other values of a, we report the corresponding neighborhood size (the final K in KNN-Uniform and the average of the final K; in KNN-KDE) in Table 12 The consistent performance with a E (0.2, 0.4, 0.6, 0.8} shows that our framework is not sensitive to the choice of
E.5.2 Effects of Kernel Size in KNN-KDE
We vary the kernel size for the kernel density estimation in KNN-KDE. The performance is shown in Table The Fl scores of all three downstream tasks are consistent across different choices kernel size. The results show that the performance of KNN-KDE is not sensitive to the choice of kernel size.
E.5.3 Performance of KNN-TV
We evaluate KNN-TV (C 0.25, a 0.6) and show the results in Table I1 KNN-TV performs similarly to KNN-KDE (a = 1), and significantly worse than KNN-KDE (a = 0.6). The reason is that KNN-TV assigns almost all the probability mass (more than 99.99% the 1-nearest neighbor of each query example and causes overfitting to them, behavior similar to KNN-KDE (a 1).
Table 11: The performance of KNN-TV compared with KNN-KDE (a 1) and KNN-KDE (a 0.6). F1 scores of the downstream tasks are reported with standard deviations shown in the subscripts.
Dataset
ChemProt (1K) AGNews (3K)
IMDB (10K)
KNN-TV KNN-KDE (a = 1)
65.8
88.0p.7
91.0p.1
67.7p.1
88, 1o.3
91.1o.1
KNN-KDE (a 0.6)
76.7p.5
89.2p.1
91.5p.1

---

Table 12:
The neighborhood size of KNN- KNN--nnifom KNN-KDE for different values of a. The numbers before the slashes are for KNN-Uniform and those after are for KNN-KDE.
Dataset ChemProt (IK) AGNews (3K)
IMDB (10K)
0.2
959 / 993
813 / 830
918 / 928
0.4
398 / 408
342 / 346
372 373
= 0.6
189 / 191
165 / 164
177 / 174
a 0.8
76 / 74
69 / 65
72/
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We clearly state our contributions and scope in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations in the conclusion.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they seale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
25

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Each theorem comes with a clear statement of the assumptions and proofs are provided in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The details and hyperparameters of the experiments provided in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
26

---

(c) If the contribution is a new model (e.g.. a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: All the data used are open sourced.... We have included the code in the Supple- mentary material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
cabbic/guides. httes//oooSuumiiiiooniiie ooiiooicoolylloooooo for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
7 //nips http:c/ppppppc..deceeeoeeeeeepee htlpli/c/guddes//oeeeeiee
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The training and test details are in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
27

---

Justification:
Error bars are reported for every experimental result.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The hardware we use is reported in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: Yes
Justification: We have reviewed the NeurlPS Code of Ethics and confirm that our research conforms it.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
28

---

Answer: [Yes]
Justification:
We have such discussion in the introduction and the conclusion.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer: [NA]
Justification:
This paper does not introduce new models or datasets.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets g.g code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We use open-source models and data, which have been properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
29

---

The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, palasets.papersstt paperswi thcode. com con/dataset has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors are encouraged reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The code is well documented.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We do not have such experiments.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not have this type of studies.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
30

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
31


Reviews:
[{"review_id": "OA81V3eCbo", "reviewer": "Reviewer_mg1e", "summary": "This paper proposes task-specific training data selection for language model fine-tuning. Given a (small) set of representative examples for a task and a large set $D$ of possible training examples, the proposed method uses (regularized) optimal transport to assign a probability distribution over $D$ that matches the distribution of representative examples while also encouraging diversity among the elements of $D$ assigned a nonzero probability.\nThe authors prove that with a certain choice of regularization function, this is equivalent to (an adaptive version of) $k$-nearest neighbor selection of candidate data similar to the representative examples. Since $k$NN treats near-duplicates as distinct examples (which would decrease diversity of the selected data), the paper additionally introduces another regularization term based on kernel density estimation; the optimal transport with this regularization is a weighted $k$NN that intuitively accounts for the frequency of near-duplicates for each example.", "questions": "- It might be helpful to have some (simple) intuition after L140 explaining why regularizing distance to the uniform transport encourages diversity.\n\n- If the optimal transport formulation is equivalent to a certain type of $k$NN, why not just present the method as a type of $k$NN? $k$NN has a long history in data selection going back to at least Wilson (1972). It's not clear what the optimal transport discussion buys other than added complexity.\n\n- Given the optimal transport framing, I think there should be some discussion of other (efficient) regularized optimal transport algorithms, such as Sinkhorn? [2]\n\n- If I understand L252--L254 correctly, the effective dataset size for the proposed method is actually up to 4x the reported size, because the data are resampled from the computed distribution each epoch. Does LESS (the baseline) get the same advantage? I.e., does LESS get to use 4x the data or do some kind of resampling?\n\n[1] Wilson, Dennis L. \"Asymptotic properties of nearest neighbor rules using edited data.\" IEEE Transactions on Systems, Man, and Cybernetics 3 (1972): 408-421.'\n\n[2] Cuturi, Marco. \"Sinkhorn distances: Lightspeed computation of optimal transport.\" Advances in neural information processing systems 26 (2013).", "limitations": "The paper includes a reasonable discussion of its limitations, but it could be improved by discussing some of the additional limitations with the empirical results mentioned above.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Good data selection is an important problem given that today's models are both expensive to fine-tune and very sample-efficient *if* they are given the \"correct\" high-quality fine-tuning data [1]. Most high-performing efforts still tweak the composition of these small task-specific datasets by hand. This paper has an interesting new take on framing task-specific data selection as an optimal transport problem between representative task examples and a large pool of candidate training data.\n\n- Theorems 3.1 and 3.2 shows that with certain regularization terms, the optimal transport selection procedure is equivalent to certain variations of $k$-nearest-neighbor. This allows for efficient computation of the optimal data selection under this objective.\n\n- The proposed approach can naturally be combined with approximate nearest neighbor search methods for efficiency.\n\n- Strong empirical results showing that the proposed selection procedure can even outperform tuning with the full candidate dataset.\n\n- The experiments include standard deviations across three runs, giving a sense of how big the gains are compared to noise.\n\n[1] Zhou, Chunting, Pengfei Liu, Puxin Xu, Srini Iyer, Jiao Sun, Yuning Mao, Xuezhe Ma et al. \"LIMA: less is more for alignment.\" In Proceedings of the 37th International Conference on Neural Information Processing Systems, pp. 55006-55021. 2023.", "weaknesses": "- Missing an ablation using embeddings instead of gradients, or any other distance function for the examples.\n\n- Missing several relevant data selection baselines that also encourage diversity, e.g. Deita [1], QDIT [2], and methods based on DPPs [3].\n\n- Changing the data mix changes the optimal learning rate (e.g., since it changes the scale of the loss function at initialization). The paper compares models trained on different data mixes with the same learning rate, but the fair comparison is optimal : optimal. It's not clear from the experiments whether the reported gains are due to the learning rate being more optimal for the selected data mix, especially since the metric used to select the data is based on the gradients of a model.\n\n[1] Liu, W., Zeng, W., He, K., Jiang, Y., & He, J. What Makes Good Data for Alignment? A Comprehensive Study of Automatic Data Selection in Instruction Tuning. In The Twelfth International Conference on Learning Representations.\n\n[2] Bukharin, Alexander, and Tuo Zhao. \"Data diversity matters for robust instruction tuning.\" arXiv preprint arXiv:2311.14736 (2023).\n\n[3] Wang, P., Shen, Y., Guo, Z., Stallone, M., Kim, Y., Golland, P., & Panda, R. (2024). Diversity Measurement and Subset Selection for Instruction Tuning Datasets. arXiv preprint arXiv:2402.02318.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback. We address the concerns and clarify the questions as follows:\n\n>W1 - ablation study using embeddings\n\nWe 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. \n\n>W2 - additional methods\n\nThanks 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.\n\n>W3 - hyperparameter setting\n\nWe 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.\n\n>Q1 - intuition behind uniform transport\n\nThank 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.\n\n>Q2 - why use optimal transport instead of presenting the method as knn\n\nThe 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.\n\n>Q3 - other efficient optimal transport algorithms\n\nThanks for the suggestion. Incorporating algorithms like Sinkhorn for even further efficient computation is an interesting direction. We will add a discussion to related works.\n\n>Q4 - effective dataset size\n\n- 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. \nLESS 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. \n- 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.\n- 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”. \n\nWe will add a discussion with the additional experiment to the appendix.\n\n[1] Xia, Mengzhou, et al. \"Less: Selecting influential data for targeted instruction tuning.\" International Conference on Machine Learning. PMLR, 2024.\n\n[2] Xie, Sang Michael, et al. \"Data selection for language models via importance resampling.\" Advances in Neural Information Processing Systems 36 (2023): 34201-34227.\n\n[3] Coleman, Cody, et al. \"Selection via proxy: Efficient data selection for deep learning.\" International Conference on Learning Representations (2019).\n\n[4] Albalak, Alon, et al. \"A survey on data selection for language models.\" arXiv preprint arXiv:2402.16827 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BpXxpjjV1U", "reviewer": "Reviewer_zKiE", "summary": "This paper presents a method for data selection for task-specific model finetuning. The method relies on a small, representative sample of data from the target task to select matching, relevant data from a corresponding corpus. The method relies on framing this task as an optimization problem, utilizing an optimal-transport-based distribution alignment loss and a KDE-based regularizer to avoid oversampling (near-)duplicates. \nThe authors show this method to be highly scalable and data efficient, being competitive with, and often outperforming state-of-the-art methods for domain-specific continued pretraining and task-specific instruction tuning.", "questions": "Again, comparing efficiency with other SOTA methods on the datasets used in this paper would be helpful in better contextualizing the performance presented.", "limitations": "The authors address limitations.", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper rigorously presents and tests the proposed method, with a detailed theoretical motivation. \n- Sections 2-4 are well structured and introduce the method in a clear, progressive way. \n- Performance results, especially for very small sample sizes, strongly support the utility of this method.", "weaknesses": "- Section 5.1: Given how different some of the performances are between llama and mistral, including other LLMs may give a more complete picture of the efficacy of this method.\n- Section 5: Efficiency claims would benefit from context. How does the 28 hour initialization time compare to other SOTA methods on this dataset? How does it scale after initialization is done when repeatedly drawing task-specific samples compared to other methods?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the feedback. Here is our response to the questions and concerns:\n\n>W1 - other LLMs\n\nThank 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.\n\n> W2 & Q1 - runtime\n \nOn 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.\nFor 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).\nFor 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.\nWe will add the discussion above to the appendix.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "p1fzqZDlPz", "reviewer": "Reviewer_rjVY", "summary": "This paper formulates data selection for task-specific fine-tuning as an optimization problem based on optimal transport for distribution alignment. It proposes two KNN-based implementation methods and evaluates them on datasets for task-specific instruction fine-tuning and domain-specific continued pretraining. The experimental results demonstrate that their methods are more effective than the baseline systems (LESS and DSIR).", "questions": "See weakness.", "limitations": "They acknowledge one limitation in conclusion: the method still relies on M query samples for a specific task. It is also raised in my comments.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "1. The paper formulates data selection as an optimal transport problem, providing a detailed problem definition and a closed-form solution. Additionally, it proposes using Kernel Density Estimation to address the issue of near-duplicates.\n\n2. The authors introduce KNN-Uniform and KNN-KDE algorithms for data selection, showing that their performance is superior to the baseline systems in both task-specific instruction fine-tuning and domain-specific continued pretraining experimental setups.", "weaknesses": "**Regarding the methodology:**\n\n1. The connection between data selection and the optimal transport problem is not clearly established. Despite mentioning it in lines 114-115 of Section 3, it remains unclear why data selection can be considered an optimal transport problem.\n\n2. Much of the paper is based on LESS, including the representation of samples and the task definition. However, there is minimal mention of LESS, making it challenging to understand without prior knowledge of LESS.\n\n3. The method still relies on M query samples for a specific task, which poses certain limitations.\n\n**Regarding the experimental section:**\n\n4. The experimental section contains too many specific settings. For instance, special settings mentioned in lines 248 and 286 make it difficult to determine how these parameters were chosen, even after reviewing the appendix.\n\n5. The two experimental setups are inconsistent in task-specific instruction fine-tuning and domain-specific continued pretraining. In Table 2, the Ratio of 0.5%-5% can be understood as the number of data samples selected. However, in Table 4, the 1K, 3K, and 10K seem to refer to the number of query samples, but there is no comparison of the number of selected samples.\n\n6. There is a lack of comparison with various baseline systems. Only one baseline system is used for comparison, and its performance differs from that reported in the original paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We address the concerns and answer the questions below:\n\n>W1 - connection to optimal transport\n\nAs 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.\n\n>W2 - discussion of LESS\n\nWe 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.\n\n>W3 - need M query examples\n\nIn 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.\n\n>W4 - how to set parameters\n\nNote 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: \n- We set $C$ to 5 when the embeddings are normalized.\n- 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.\n- $\\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.\n\nWe will add the details to the appendix.\n\n>W5 - evaluation using different sample sizes for continued pretraining\n\nWe 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.\n\n>W6 - baseline systems\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PMLLSprfoj", "reviewer": "Reviewer_CE8N", "summary": "This paper proposes a method for data selection in foundation model fine-tuning. The proposal contains a distribution alignment loss based on optimal transport to capture the discrepancy between the selected data and the target distribution, a regularizer to encourage the diversity of the selected data, and kernel density estimation to reduce the negative effects of near-duplicates among the candidate data. Experimental on fine-tuning the language model are reported.", "questions": "As discussed in the Weakness part.", "limitations": "Yes", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The proposal studied in this paper is interesting since how to select data, and how to improve the data quality is important for the training and fine-tuning of the foundation model.\n2. The proposed method which considers distribution discrepancy minimization, diversity, and near-duplicates, is technically sound.", "weaknesses": "1. The novelty and contribution of the proposed method are limited. Data selection is important and well-studied in the machine learning community. For example, in active learning, we need to select examples to label according to some metrics; in domain adaptation, we need to select data to help the model reuse. Some widely adopted methods can be applied to the problem of data selection for the foundation model and the authors didn't provide a comprehensive study and comparison. Moreover, the techniques adopted in the proposal are also widely used techniques. \n\n2. For the experiments, the authors only conduct experiments on the language model, can the proposal be applied to other foundation models, such as the vision-language model?\n\n3. It seems that the random selection method can also achieve good performance. So I am wondering about the difficulty of the problem,  maybe we can improve the performance using some trivial techniques.\n\n4. The time cost between fine-tuning with the full dataset and the selected dataset should be reported.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> W1 - novelty\n\nTo 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. \n\n> W2 - extension to VLMs\n\nWe 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.\n\n>W3 - performance of random selection\n\nThis 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.\n\n>W4 - finetuning cost\n\nFinetuning 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.\n\n[1] Xia, Mengzhou, et al. \"Less: Selecting influential data for targeted instruction tuning.\" International Conference on Machine Learning. PMLR, 2024.\n\n[2] Xie, Sang Michael, et al. \"Data selection for language models via importance resampling.\" Advances in Neural Information Processing Systems 36 (2023): 34201-34227.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Following the reviews, we would like to expand on our choice of using optimal transport as a means to solve the problem of task-specific selection. We use optimal transport to capture the discrepancy between the distribution we will sample from and the target distribution. We include probability transport cost in the optimization objective to encourage alignment between the two distributions.
We formulate our framework as an optimization problem to provide guidance for task-specific data selection. By framing the problem this way, we are able to integrate and balance two crucial objectives of task-specific data selection: distribution matching and diversification.
We find this new formulation also allows us to obtain practical solutions that achieve state-of-the-art results for a diverse array of settings. We will update our introduction to clarify the above regarding the novelty of our work.

Meta-review: Summary: 
The paper presents a method for selecting data for fine-tuning LLMs. Starting with a small representative set of examples for fine-tuning, the method formulates the data selection problem as an optimization problem with a distribution alignment loss from optimal transport to capture the discrepancy between the selected data and the target distribution. Empirical evidence shows that the method is scalable, data efficient, and competitive or better than the state-of-the-art baselines for this task.   

Strengths: 
* The paper presents a novel technique for task-specific data selection for finetuning LLMs, an important problem in machine learning.
* The proposed approach is theoretically well-grounded and motivated
* The empirical evidence strongly supports the effectiveness of the proposed approach
* The experimental validation is quite rigorous 
* The paper is well-written and easy to understand

Weaknesses 
* Discussion and comparison (wherever possible) with the previously proposed data selection methods is lagging and could be improved. 
* Given the huge difference in performance between the two LLMs testing, the paper would benefit from showing the performance across some more open-source LLMs

Overall Recommendation: 
The paper presents a novel approach based on optimal transport to select data efficiently for fine-tuning LLMs on a specific task. The approach is well-motivated, the experimental results are convincing, and the paper is clearly written. The overall recommendation is to accept this paper for publication, subject to the proposed revisions, as it will benefit the wider ML community.
Decision: Accept (poster)