Influential Language Data Selection via Gradient Trajectory Pursuit

## Abstract

Abstract Curating a desirable dataset for training has been the core of building highly capable large language models (Touvron et al., 2023 ; Achiam et al., 2023 ; Team et al., 2024 ) . Gradient influence scores (Pruthi et al., 2020 ; Xia et al., 2024 ) are shown to be correlated with model performance and are commonly used as the criterion for data selection. However, existing methods are built upon either individual sample rankings or inefficient matching process, leading to suboptimal performance or scaling up issues. In this paper, we propose Gradient Trajectory Pursuit (GTP) , an algorithm that performs pursuit of gradient trajectories via jointly selecting data points under an L ⁢ 0 𝐿 0 L0 italic_L 0 -norm regularized objective. The proposed algorithm highlights: (1) joint selection instead of independent top- k 𝑘 k italic_k selection, which automatically de-duplicates samples; (2) higher efficiency with compressive sampling processes, which can be further sped up using a distributed framework. In the experiments, we demonstrate the algorithm in both in-domain and target-domain selection benchmarks and show that it outperforms top- k 𝑘 k italic_k selection and competitive algorithms consistently, for example, our algorithm chooses as low as 0.5% data to achieve full performance on the targeted instruction tuning tasks.

## 1 Introduction

Large language models encompasses an enormous amount of knowledge acquired through pretraining corpus (Brown et al., ; Jiang et al., ; Touvron et al., ). A crucial component that makes language models practically useful is the post-training phase, which empower the model with a variety of extra abilities and skills, such as aligning language models to follow human instructions (Ouyang et al., ; Taori et al., ; Wang et al., ; Xu et al., ), teaching models the tool use (Schick et al., ) or exploring environments (Carta et al., ; Lin et al., ).

Serving as the core of model training, choosing and utilizing desirable datasets are critical for both quickly adapting models given a fixed budget  (Xie et al., ) and maximizing the performance for a targeted task  (Wang et al., ; Xia et al., ). An automatic data selection algorithm that chooses a subset based on certain information is gaining its needs given the massive amount of corpora available online. Gradient information, which reflects the optimization process and has direct correlation with the final model performance, is one of the most common criterion for data selection (Mirzasoleiman et al., ; Killamsetty et al., ; ; Yang et al., ; Pan et al., ; Xia et al., ). However, how to effectively utilize gradient for actual selection process is still a challenging problem.
Top-$k$ selection is fast and the most straightforward way, where data quality is ranked via the similarity scores of individual gradient vectors and the main gradient trajectory (Xia et al., ). Nonetheless, due to its independence assumption, its performance is often suboptimal compared to joint selection (Evans et al., ). Orthogonal Matching Pursuit (Cai & Wang, ; Killamsetty et al., ) can perform joint selection through interatively removing subspace projections, but is highly time-inefficient and requires solving a non-negative least square for as many iterations as the target subset sample size.

Figure: Figure 1: Our method selects a subset of training data through pursuit process of gradient trajectories on a target subspace during warmup training. The algorithm automatically de-duplicates samples instead of simply selecting the top-$k$ data. Text examples in the figure are drawn from instruction tuning datasets.
Refer to caption: x1.png

In this paper, we introduce Gradient Trajectory Pursuit (GTP), an algorithm that selects data samples through matching the trajectory on a gradient subspace. As discussed in previous works, a model’s optimization process can be seen as the Langevin dynamic of an energy-based model  (Welling & Teh, ), and matching the optimization process can lead to model weights with similar performance (Zhao et al., ; Cazenavette et al., ). This sets the base rationale for our algorithm. Further drawing inspiration from studies demonstrating that the optimization process happens in a subspace  (Gur-Ari et al., ; Frankle & Carbin, ; Singhal et al., ), we project the gradients onto a small subspace and greatly reduces the memory cost (both on disk and RAM) during selection. This design also follows and is backed up by a recent work (Xia et al., ). The core of our algorithm is joint data selection, where data weights are solved through compressive sampling matching pursuit process. The pursuit process is equivalent to starting from a top-$k$ intialization, iteratively determines the proper combination of data samples and implicitly performs de-duplication. We also showcase a distributed version of the algorithm which further reduces the selection computation time, indicating its potential for scaling up to larger datasets.

Our contributions are summarized as follows:

- •
We propose an algorithm that jointly selects influential language data given a subset budget, via matching trajectories in a subspace. The algorithm is both effective on select compact dataset and up to $17$x more efficient compared to a vanilla orthogonal matching pursuit-based algorithms and automatically performs de-duplication compared to top-$k$. Our algorithm also support distributed selection across machines.
- •
We experimentally show great improvements over the prior methods on two challenging benchmarks for language models, demonstrating our algorithm’s usage on both in-domain and target-domain data selection. Specifically, we show that our algorithm can select only 0.5% of the whole dataset to achieve full performance on targeted instruction tuning.
- •
We in-depth analyzes our algorithms, including the convergence, the progression of subset quality over iterations, and comparison with top-$k$ selected data on ALFWorld agent dataset, providing understandings on how gradient similarities can lead to duplication and our algorithm’s adjustment process for creating a compact subset.

## 2 Related works

Coreset data selection. Building an effective training dataset has evolved to the cornerstone of foundation model construction. Due to the heavy cost of human annotation and filtering, a flux of works are developed to perform automated data selections. Features of data carry information that discriminates samples and are often used for selecting representative points (Sener & Savarese, ; Kaushal et al., ; Xia et al., ; Wang et al., ; Xia et al., ; Xie et al., ). N-gram features (Xie et al., ) are used for re-weighting and selecting samples for pre-training, demonstrating strong correlations with target domains. Deep learning features (Zhang et al., ; Hanawa et al., ; Xia et al., ) are common choices for effective baselines measuring data similarities. Although frequently adopted, feature-based methods are not guaranteed to have correlation with model optimization. Gradient information is another natural choice for selecting influential data (Yu et al., ; Mirzasoleiman et al., ; Mindermann et al., ; Han et al., ; Xia et al., ). Different from features, gradients often reflect more information about the model optimization process (Mirzasoleiman et al., ; Xia et al., ), but often are used for in-domain selection (Mirzasoleiman et al., ; Killamsetty et al., ) and suffer from the curse of dimensionality, given the current model sizes (Brown et al., ; Jiang et al., ). Instead, our work allows for both in-domain selection and target domain transfer, and focuses low-dimensional representations. Besides the choices of representations, selection algorithms also play a crucial role. Various sampling methods or clustering algorithms are explored (Sener & Savarese, ; Kaushal et al., ; Xia et al., ; Wang et al., ; Mirzasoleiman et al., ; Killamsetty et al., ). Data Models (Ilyas et al., ; Engstrom et al., ), LLM-asking (Sachdeva et al., ), or information optimization (Everaert & Potts, ) have shown their effectiveness. Note that among these methods, gradient-based ones usually have more theoretical supports due to its connection with model optimization.

Data influence. Computing the data influence for model is another line of work for determining the importance of samples. Influence function for deep neural networks (Koh & Liang, ; Koh et al., ) can approximate the influence of dropping training samples. Trajectory influence (Pruthi et al., ; Han et al., ; Xia et al., ) is different definition introduced to measure how each data sample affects the whole model training dynamics. Top-$k$ selection and variants are adopted for selecting data with higher trajectory influences (Han et al., ; Xia et al., ).

Model training dynamics and SDEs. It has been shown that machine learning model trainings are also a diffusion processes (Welling & Teh, ; Hoffman et al., ; Wenzel et al., ) defined through stochastic differential equations. Commonly known in generative models, diffusion processes converge to desired distributions when following the target score vector (gradient) field (Song & Ermon, ; Song et al., ; Ho et al., ). Several works (De Bortoli, ; Zhu et al., ; Benton et al., ) also explore the mixing rate and convergence bounds of diffusion models. Similarly, in optimization processes, when following the gradient vector field, trained models can converge to the same parameter distribution  (Zhao et al., ; Cazenavette et al., ).

## 3 Our method

In this section, we introduce our algorithm Gradient Trajectory Pursuit (GTP), a framework designed to select influential language data through joint data selection. We start from introducing the problem formulation in section 3.1, discussing the base rationale behind trajectoy matching, then explain the main algorithm and the distribtued variants in section 3.2. Finally, we discuss the complexity and computation time of our algorithm in section 3.3.

### 3.1 Problem statement

Problem definition. Given a large collection of training samples $\mathcal{D}_{tr}=\{(\bm{x}_{i},\bm{y}_{i})\}_{i=1}^{N}$, the goal is to select a small subset $\mathcal{D}_{S}=\{(x_{s_{i}},y_{s_{i}})\}_{i=1}^{M}$ that maximizes the generalization performance of models trained on it when evaluated on a test set $\mathcal{D}_{te}$. This test set may either be drawn from the same distribution as the training data or from a different, task-specific distribution, leading to in-domain or targeted-domain data selection.

Main idea. We approach the subset selection problem from considering model training dynamics. With a training set, the gradient descent process converges to the posterior distribution $p(\bm{\theta}|\mathcal{D})\propto p(\mathcal{D}|\bm{\theta})$ over the model parameter $\bm{\theta}$. Assuming the existence, if a selected subset $\mathcal{D}_{S}$ has a gradient vector field w.r.t. $p(\bm{\theta}|\mathcal{D}_{S})$ that closely matches the gradient vector field of the full dataset, a model trained using the subset can converge to the same distribution over $\bm{\theta}$. Considering the intensive computation cost for computing and matching over the whole gradient vector field(^1^11For example, if we store gradient of a full language model over the full training set, the storage is equivalent to saving $N$ copies of language models. Every parameter state we consider will need a set of $N$ copies, leading to $N\times\text{number of states}$ stored on disk or RAM.), we only focus on a part of parameter states obtained from teacher training and perform matching in a subspace. This draws inspiration from findings that the optimization process actually happens in a subspace (Gur-Ari et al., ; Frankle & Carbin, ; Singhal et al., ) where a random guess can already perform decently well. We use parameter states from early trajectories, assuming that the early training provides the most diversity and larger gradient subspace since optimization is an information collapse process. To minimize the matching objective, we collectively solve $|\mathcal{D}_{S}|$ samples, and perform joint selection among all data points. This differs from previous works, where  Killamsetty et al. () focuses on online settings (ours is offline) and uses OMP to select one sample per iteration (leading to scaling issue in offline cases), while compared to Xia et al. () which uses independent selection, our algorithm performs joint selection to achieve better performance. We do not make specific assumptions on in-domain or target-domain and use $\mathcal{D}_{tar}$ to denote the target dataset, which can come from either in-domain data or target domains. Our full algorithm is summarized in algorithm 1. The detailed algorithm components and derivation are discussed in the following section.

### 3.2 Gradient Trajectory Pursuit

Figure: Algorithm 1 GTP Algorithm

We denote the parameter states on the optimization trajectory as $\bm{\tau}=\big{(}\bm{\theta}^{(0)},\bm{\theta}^{(1)},...,\bm{\theta}^{(T-1)}
\big{)}$. $T$ can be manually set. To select the index subset $S$ that matches the gradient of target dataset in a subspace, we define the following objective function:

$$ $\mathcal{L}_{t}(S)=\big{|}\big{|}U^{t}\circ\nabla_{\bm{\theta}^{(t)}}\log p( \bm{\theta}^{(t)}|\mathcal{D}_{tar})-U^{t}\circ\nabla_{\bm{\theta}^{(t)}}\log p (\bm{\theta}^{(t)}|\mathcal{D}_{S})\big{|}\big{|}_{2}$ (1) $$

where $\mathcal{L}_{t}$ is the per-step matching loss, $U^{t}\circ\bm{g}$ denotes the projection of vector on a subspace $U^{t}$, $\mathcal{D}_{tar}$ is the full dataset for in-domain case and target-domain dataset otherwise. We would like to minimize the summation across steps as the final loss: $\mathcal{L}(S)=\sum_{t}\mathcal{L}_{t}(S)$. To formulate the selection set $S$ into the objective, we introduce a set of weights $w_{i}$ for each data point and define the gradient as $\nabla_{\bm{\theta}}\log p(\bm{\theta}|\mathcal{D}_{S})=\sum_{i=1}^{N}w_{i}
\nabla_{\bm{\theta}}\log p((\bm{x}_{i},\bm{y}_{i});\bm{\theta}),\ (\bm{x}_{i},
\bm{y}_{i})\sim\mathcal{D}_{tr}$, where $\bm{w}=(w_{1},...,w_{N})$ is regularized through L0 norm $||\bm{w}||_{0}$. Due to sparsity enforced in $L0$ norm, the indices of final weights $\bm{w}$ that are non-negative is the final selected set $S$. The objective function with index set incorporated is then:

$$ $\mathcal{L}(S)=\sum_{t}\big{|}\big{|}U^{t}\circ\nabla_{\bm{\theta}^{(t)}}\log p (\bm{\theta}^{(t)}|\mathcal{D}_{tar})-\sum_{i=1}^{N}w_{i}U^{t}\circ\nabla_{\bm {\theta}^{(t)}}\log p((\bm{x}_{i},\bm{y}_{i});\bm{\theta}^{(t)})\big{|}\big{|} _{2}+\big{|}\big{|}\bm{w}\big{|}\big{|}_{0}$ (2) $$

Note that directly optimizing the above objective function requires a proper algorithm to minimize the non-differentiable $L0$ norm in equation 2. We will discuss the core solver variants that jointly minimize equation 2 in later this section.

The evolving subspace. We choose to project the model gradients onto a subspace before matching, both to reduce the storage and computation cost and to rule out potential noise signals (Singhal et al., ). As training progresses, parameter state at each step $t$ can have a different meaningful subspace $U^{t}$. We parallelize the subspace computation across and obtain a series of evolving subspace $\big{(}U^{0},\ \ldots,\ U^{T-1}\big{)}$. To compute the subspace, we use the full training set $\mathcal{D}_{tr}$ for in-domain and use target dataset $\mathcal{D}_{target}$ for target-domain case. Note that different subspace analysis algorithms are applicable here. We adopts standard PCA for computing the principal components for the subspace.

Figure: Algorithm 2 Iterative Compressive Sampling

Figure: Algorithm 3 Distributed Compressive Sampling

The non-negative least square problem. To simplify the notations of equation 2, we concatenate the projected gradient vector $U^{t}\circ\nabla_{\bm{\theta}^{(t)}}\log p(\bm{\theta}^{(t)}|\mathcal{D}_{tar})$ across steps as $\bm{b}\in\mathbb{R}^{T\times D}$ where $D$ is the subspace size, and build a matrix $\bm{A}$ where each row is the projected gradient vector for a data point. With non-negative weight vector $\bm{w}\in\mathbb{R}^{1\times N}$, we arrive at a simple non-negative least square problem with $L0$ norm as regularization:

$$ $\min_{\bm{w}}\ ||\bm{w}||_{0}\quad\text{subject to}\ ||\bm{A}\bm{w}-\bm{b}||_{ 2}^{2}=0$ (3) $$

The above is the equation 2 with simplified notation. To solve the objective function, we discuss two algorithms: Iterative Compressive Sampling Pursuit and Distributed Compressive Sampling Pursuit. The two algorithms are summarized in algorithm 2 and 3.

Iterative Compressive sampling pursuit. Equation 3 cannot be directly minimized through differentiation. Instead, a greedy selection algorithm that solves this problem is developed in Needell & Tropp () through iterative greedy selection based on a residual vector. In the selection process, $2M$ top data points are jointly selected all at once, a non-negative least square step is used to adjust the weights (i.e. implicit de-duplication), then the re-ranked top $M$ samples are used to update the residual vector. Compared to orthogonal matching pursuit (OMP) used in  Killamsetty et al. (), the algorithm is much faster and scalable.

Distributed Compressive sampling pursuit. Given the potential larger number of time steps in consideration, and the amount of data samples to be selected, compressive sampling matching pursuit can still be not mostly optimal in terms of computation time. We further design an algorithm that can distribute the pursuit process across multiple machines. We use a star model, where one center machine communicates with multiple machines and adopt column partition to build separate $\bm{A}$s that are distributed and parallelized. Details are shown in algorithm 3.

### 3.3 Computation time analysis

As pointed out in Xia et al. (), the computation time cost for all steps should be considered. Our algorithm consists of warmup model training, gradient subtraction and storing, subspace computation, and data selection, which follows the design in Xia et al. (). The main difference come from our subspace computation and selection algorithm, where we adopt PCA instead of random projection and use GTP instead of top-$k$ for data selection.

## 4 Experiments

**Table 1: Comparison across methods on ALFWorld expert demonstration dataset under various budgets. Top-$k$ selection leads to data points that are more similar (duplicate in terms of information). Joint selection using GTP outperforms all baseline methods consistently.**
| Method | 5% | 10% | 15% | 20% | Full data |
| --- | --- | --- | --- | --- | --- |
| $\#$Samples | 500 | 1000 | 1500 | 2000 | 10000 |
| Random | 42.7 (2.8) | 61.3 (3.5) | 71.7 (3.1) | 80.5 (3.1) | 85.6 (2.6) |
| LESS${}_{\text{SGD}}$ (top-$k$) | 38.9 (1.7) | 41.6 (1.2) | 42.4 (2.2) | 43.2 (0.9) |  |
| G-DIG | 48.8 (4.2) | 58.4 (3.5) | 70.1 (1.9) | 75.6 (2.1) |  |
| RDS (representation-based) | 37.3 (1.8) | 63.0 (2.1) | 74.9 (2.4) | 80.8 (1.2) |  |
| GTP - full (ours) | 52.7 (3.0) | 72.1 (2.8) | 77.2 (2.4) | 82.1 (2.9) |  |
| $\Delta$ (improvement) over random | 10.0% | 10.8% | 5.5% | 1.6% |  |
| $\Delta$ (improvement) over top-$k$ | 13.8% | 30.5% | 34.8% | 38.9% |  |
| GTP - distributed (5 machines) | 50.4 (2.9) | 71.6 (2.8) | 76.2 (1.9) | 81.6 (1.7) |  |
| $\Delta$ (improvement) over random | 6.6% | 10.3% | 4.5% | 0.8% |  |
| $\Delta$ (improvement) over top-$k$ | 10.4% | 30.0% | 33.8% | 38.1% |  |

We evaluate our data selection algorithm on standard benchmarks against gradient-based prior arts. In section 4.1, we perform in-domain data selections on large language model agent tasks. In section 4.2, we compare algorithms on targeted instruction tuning, a benchmark focusing on empowering language model to adeptly follow human instructions.

### 4.1 LLM Agent: ALFWorld

Benchmark setup. ALFWorld is a suite of text environments built upon interactive Textworld  (Côté et al., ) that benchmarks agents on solving sequential tasks. Typical tasks include giving agents the current partial observed environment information and require the agents to achieve a goal (e.g., finding objects).

We focus on offline imitation learning settings. For train and test datasets, we use offline trajectories from the official ALFWorld repository  (Shridhar et al., ) which consists of $3,553$ tasks, each with an average of 19 steps. The training set contains $48k$ state-action pairs, and test set totals at $1,861$ state-action pairs. In practice, we find that using $10k$ training data already achieves similar generalization performance compared to full training set, i.e., dataset reduction beyond $10k$ is not very meaningful since the selection can be easily done through random subset picking, and thus choose to use a subset with $10k$ data points as our main subject of study. In model evaluation, for convenience, we directly compare the model’s per-step action prediction accuracy. Per-step accuracy is known to be reflective on model’s ability and has a strong correlation with final success rates.

Model setup. For evaluating whether a dataset is effective for training, we use supervised fine-tuning on OPT-125M (Zhang et al., ) with LoRA () as the basic training setup. Given a selected subset dataset for training, all models are optimized using Adam  (Kingma & Ba, ) with a learning rate 5e-4 and batch size 80 for 20 epochs to ensure its convergence. We follow the most recent prior method  (Xia et al., ) to generate a warm-up training trajectory for gradient extraction and data selection. The warm-up trajectory is trained using learning rate 1e-4 and batch size 10 for 5 epochs. A model checkpoint is saved every 500 interations, resulting in model states on across timesteps.

Baselines. The most relevant prior method is LESS  Xia et al. (), where we build upon most of its pipeline, including warm-up training and gradient computation. LESS uses top-$k$ selection with scores maxing across pre-defined subtasks, which does not exist in ALFWorld. We instead directly compare individual gradient vectors with gradient trajectories.
G-DIG  (Pan et al., ) is a baseline that uses clustering directly on gradients and samples within each cluster. Following  (Xia et al., ), we also compare with RDS (Representation-based Data Selection)  Zhang et al. (); Hanawa et al. (), which uses the mean of last layer hidden states for selection.

Main results. We summarize the main results in table 1. As demonstrated in previous works  (Xie et al., ; Xia et al., ; Pan et al., ), random sampling is a strong baseline to beat. Gradient-based algorithms tend to pick highly redundant examples when matching gradients. LLM Agent benchmarks contain traces with incremental steps in the training data, and need algorithms to be sensitive to sample duplication. Why does top-$k$ fail? We visualize the selected top three examples based on the scores computed from two algorithm in figure 2, showing that the initial top k selections are highly duplicated, while after our pursuit algorithm the selected samples exhibit more diversity. Representation-based selection (RDS) shows higher performance since it uses the features from pretrained model. Across different data ratios, our algorithm, both standard and distributed versions, consistently outperform the baselines. For the distributed version, we use separate the 10 timesteps into 5 machines.

Figure: Figure 2: Three examples with topmost scores from selection algorithms. Left: Top-$k$ tends to select redundant samples due to lack of de-duplication mechanisms. Right: GTP instead chooses diverse examples.
Refer to caption: x2.png

Figure: Figure 3: Residual norms across iterations.
Refer to caption: extracted/5944882/figs/GTP_residual.png

Residual norms across iterations. We rollout 10 steps of compressive sampling matching pursuit iterations and monitor the L2 norm of residual vectors. The residual norm trends are shown in figure 4. Typically the algorithm converges after 5 steps.

Correlation between residual norm and model performance over iterations. The initial step of GTP is equivalent to a top-$k$ selection, and later iterations of GTP algorithm can be seen as a process of de-duplication through adjusting the data weights $\bm{w}$. In figure 4, we plot the accuracies of subsets of 1000 samples selected after each GTP iteration. Through progressively solving the data weights $\bm{w}$, the performance gradually improves. We only plot 5 steps since the algorithm plateaus (see figure 4).

[Uncaptioned image]: extracted/5944882/figs/computation_time.png

How does the algorithm computation time scale with number of samples? As another algorithm for joint data selection, orthogonal matching pursuit (OMP) can also collectively solve data weights for de-duplication (Killamsetty et al., ). We compare our algorithm GTP, its distributed version, and OMP across various number of budget samples. OMP selects only one sample per iteration while each requires a separate solving of non-negative least squre problem. This rapidly increases the computation time with more samples. Our algorithms, based on compressive sampling process, can scale well and have a much more mild computation time curve.

### 4.2 Targeted instruction tuning

Task setup. We follow the setup in Xia et al. () and evaluate our algorithm on the targeted instruction tuning task. Compared to standard instruction tuning where mixed datasets are used to finetune large language model to follow human instructions, targeted instruction tuning individually selects data samples for different target benchmarks for finetuning language models. The selection algorithm hence plays an important role for improving the final trained model performance. For example, multiple works find that when various data sources are mixed together, it can negatively impact the model performance for different target tasks  Wang et al. (); Xia et al. (), while  Xia et al. () achieves equivalent or better results with only 5% of the full dataset.

Datasets and models. We follow the settings from Xia et al. () and use a diverse and mixed collection of public instruction tuning datasets as the main training set to select from. The datasets include: Flan V2 (Longpre et al., ), CoT (Wei et al., ), Dolly (Conover et al., ), and Open Assistant (Köpf et al., ). The four datasets are commonly used in various models and benchmarks. For the language model being tested upon, we use Mistral-7B  (Jiang et al., )(^2^22We do not evaluate Llama series due to license issue.).

Implementation details. To build the initial parameter states for extracting gradients, we warmup train Mistral-7B (also with LoRA) for 4 epochs using 5% of the full dataset. Data gradient extractions are performed in parallel across machines. We use the selected data from each benchmark as the target dataset and compute subspace and target gradient $\bm{b}$. These are exactly following the standard setups in Xia et al. (). The GTP iterations are performed 5 times to select the final subset.

**Table 2: Comparison of our algorithm with state-of-the-art algorithm (Xia et al., ).**
|  | MMLU | BBH | TydiQA |
| --- | --- | --- | --- |
| Base | 62.4 | 57.7 | 52.2 |
| Rand | 60.0 | 54.5 | 56.9 |
| LESS | 61.8 | 56.0 | 60.3 |
| GTP (ours) | 62.0 | 60.0 | 64.5 |
| $\Delta$ (improv.) | +0.2 | +4.0 | +4.2 |
| #Samples | 1353 | 1353 | 1353 |
| Runtime (seconds) | 800 | 355 | 65 |

Evaluation benchmarks and results. There are three targeted benchmarks considered in the task: MMLU (Hendrycks et al., ), BBH (bench authors, ), and TydiQA (Clark et al., ). MMLU is a de-facto benchmark for evaluating large language model’s capability on 57 subsets across elementary mathematics, humanities, law, social sciences and more. BBH evaluates the reasoning capability of large language models. It is curated as a subset of the BIG-Bench benchmark, focusing on 27 tasks where current models struggle to match human performance. TyDiQA is a multilingual question-answering benchmark covering 11 diverse languages, and features questions from native speakers seekingn answers. Main result. The test accuracies of models trained on MMLU, BBH, and TydiQA benchmarks are summarized in table 2. We observe that our algorithm outperforms the prior method (Xia et al., ) by 0.2% on MMLU, 4.0% on BBH, and 4.2% on TydiQA using only 0.5% of the full dataset, achieving a $10$x reduction on the selected subset size and achieves better results.

## 5 Conclusions

In this paper, we propose an algorithm (GTP) that can both perform joint data selection and has the scalability towards larger target sample sizes. The algorithm automatically de-duplicates samples and can potentially serve as a standard selection technique other than top-$k$. Both in-domain selection and target-domain selection are applicable for GTP. For potential limitations, gradient-based methods still rely on a few steps of model training, which can be more expensive than pure representation-based methods or influence function methods. It will be worth attempting to connect matching pursuit with representation-based or influence function methods and develop an algorithm that does not require additional model training for data selection.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 3
presentation: 3
questions: 1.	How do we select the set of model parameters to compute the gradients? And how large this set should be? 
2.	Will the selected subset be robust to the later choices of model training, e.g., optimizer configuration, learning rate schedule? If so, why is that? 
3.	Will the proposed solution readily applicable to different stages of LLM training?
soundness: 2
strengths: 1.	Data selection for effective model training has become a very important problem for large language models. Hence, this submission touches upon a timely issue.
2.	The descriptions of the proposed solution and related work are largely clear and easy to follow. This helps readers to understand the essence of this submission.
summary: The authors of this paper studied data selection problem for large language models. The basic idea is to select a subset of datapoints from a larger training set, where the gradients of the objective function with respect to a given set of selected model parameters match those computed from the complete dataset. To reduce computation complexity, the matching is performed in the projected gradient space and experiment results were reported on the ALFWorld dataset.
weaknesses: 1.	It is well known that modern large language model training has several stages, e.g., pre-training and post-training. Given the submission is targeted at language model training, it is very unclear which stage the proposed method is designed for. Since different stages have distinct purposes, specific designs are needed, though their loss function might look similar (e.g., next token prediction).
2.	The proposed solution is very similar to the frequently mentioned LESS algorithm (Xia et al. 2024) in the paper, and this submission just changes their cosine metric to L2 norm, with an additional sparsity regularization. And of course, one advantage of the proposed solution over LESS is it can use some less greedy solution to choose data points, which seem indeed brings performance gain to LESS. 
3.	The proposed solution requires to match gradients on a set of given model parameters. It is very unclear in practice how to most effectively select this set of parameters to compute the gradient. The manuscript gives a very heuristic approach: using the model parameters obtained from early stage of training. This comes to my first question: which stage of LLM training it is. The sentence alone seems suggest pre-training (i.e., “a random guess can already perform decently”, “early trajectories”), but the experiment uses pretrained model, OPT-125M, which clearly is not a random guess. 
4.	It is very unclear why the proposed solution can really solve the complicated data selection problem for LLMs. As the authors mentioned, matching/controlling the training dynamics of LLMs is important, which govern where the final model will end up. But the proposed solution did not touch this aspect at all, e.g., the optimizer and learning rate schedule. Matching gradients on some selected points cannot suggest anything about the model training dynamics on the selected subset. And therefore, it is unclear why the selected data points can bring any good properties to the trained model.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: Q1: How do you obtain $\nabla_{\mathbf{\theta}^{(t)}} \log p(\mathbf{\theta}^{(t)} | \mathcal{D}_\text{tar})$ in Equation 1?

Q2: How do you obtain $\nabla_{\mathbf{\theta}^{(t)}} \log p((\mathbf{x}_i, \mathbf{y}_i); \mathbf{\theta}^{(t)})$ in Equation 2?

Q3: What exactly is the result of your computation time analysis in Section 3.3?

Q4: how do you obtain the model parameter trajectory as indicated in your Algorithm 1?

Q5: What exactly do "supp" and "NNLS" mean in Algorithm 2?
soundness: 3
strengths: S1: With the increasing model and dataset size, it is an important field to study effective data selection for LLMs. And the contribution of this paper seems non-trivial to the field.

S2: The overall idea of performing joint selection based on gradient trajectory matching is intuitive and reasonable.

S3: The empirical results based on ALFWorld and three instruction tuning evaluation sets are consistently better than the considered baselines.
summary: The paper introduces Gradient Trajectory Pursuit (GTP), which jointly selects data samples based on gradient trajectory matching, optimizing over an L0-norm regularized objective. It avoids data duplication by considering the entire set of selected data points rather than independently selecting the top-k samples. The algorithm uses compressive sampling techniques, significantly improving efficacy, and further supports a distributed framework for handling large datasets. Empirical results on two benchmark datasets demonstrate consistent performance of GTP over existing top-k selection and other gradient-based baselines.
weaknesses: W1: The writing and presentation of the paper need substantial improvement. It is very unclear how Equation 1 is derived from the main idea discussed in Section 3.1. The three algorithm blocks should also be described with further details in the content sections. In general, I find Section 3 very hard to follow.

W2: The scope of the current experiments seems limited. It would be great to experiment with multiple LLMs from different model families and sizes for each baseline.

W3: It would be great to include some qualitative analysis based on the 5% examples selected by your method. The findings there could be very valuable.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: Please see the weakness points above.
soundness: 2
strengths: **Strengths:**

Incorporating the combinatorial effect into data selection is intuitive, as a top-k selection approach may not fully capture the collaborative potential among data samples.

The method demonstrates considerable performance improvement over the vanilla LESS method on the ALFWorld dataset.
summary: **Summary:**

In this paper, the authors extend the work of Xia et al. (2024) to subset selection settings, where the selection mechanism relies on the collaborative effect of chosen samples rather than a top-k selection paradigm. Experimental results are provided to showcase the effectiveness of the proposed method.
weaknesses: **Weaknesses:**

As noted in subsection 3.3, the main distinction from Xia et al. (2024) lies in the departure from the top-k selection paradigm by incorporating ideas from Needell & Tropp (2009), while most other components remain unchanged.

The computational time analysis in subsection 3.3 is unclear. There is no analysis of computational complexity, nor are there empirical runtime results or GPU memory cost reports. Additionally, PCA may present a computational challenge in large-scale experiments.

The experimental scope is limited. Although multiple datasets are utilized, the main experiment comparing multiple baselines is only conducted on the ALFWorld dataset, with only LESS compared across three instruction-tuning datasets. Specifically, it would be beneficial to see a comparison of performance and computational costs with the Random baseline and G-DIG. Furthermore, the authors could explain why the naïve random selection strategy tends to outperform most baselines in this experimental setup.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 2
questions: -- Can you explain more about why Llama couldn't be compared against? I am not aware of the issue described

-- Are all the methods in Table 2 subject to the same subset size constraint? I assume they are, but it's not totally clear from the wording of the results section, which seems to imply that smaller subsets are chosen
soundness: 3
strengths: -- This type of dataset selection technique is generally important and timely.

-- The paper is clearly written (despite various small types), and has enough detail that what's described should be more-or-less reproducible.

-- The results appear fairly impressive, though seem to be fairly modest compared to the strongest baseline

-- Reasonably strong methodology paper, though the idea is ultimately fairly simple and described in pretty clear terms.
summary: The paper proposes a technique to select influential subsets of language data based on a subset budget. The ideas is based on matching trajectories in a subspace, and is claimed to be 17x more efficient compared to vanilla matching pursuit. For targeted instruction-tuning tasks, whole-dataset performance can be achieved using only 0.5% of the data.
weaknesses: -- The contribution compared to LESS seems not so large, and the performance improvements are also fairly incremental (though they are still significant)

-- Experiments seem overall somewhat less thorough (e.g. in terms of models being compared) than in prior work (such as LESS), though the authors do have some explanations for this

-- Hard to totally make sense of some of the experimental results in a few places