Harnessing Diversity for Important Data Selection in Pretraining Large Language Models

## Abstract

Abstract Data selection is of great significance in pretraining large language models, given the variation in quality within the large-scale available training corpora.
To achieve this, researchers are currently investigating the use of data influence to measure the importance of data instances, i . e . , formulae-sequence 𝑖 𝑒 i.e., italic_i . italic_e . , a high influence score indicates that incorporating this instance to the training set is likely to enhance the model performance. Consequently, they select the top- k 𝑘 k italic_k instances with the highest scores.
However, this approach has several limitations.
(1) Calculating the accurate influence of all available data is time-consuming.
(2) The selected data instances are not diverse enough, which may hinder the pretrained model’s ability to generalize effectively to various downstream tasks.
In this paper, we introduce Quad, a data selection approach that considers both quality and diversity by using data influence to achieve state-of-the-art pretraining results.
To compute the influence ( i . e . , formulae-sequence 𝑖 𝑒 i.e., italic_i . italic_e . , the quality) more accurately and efficiently, we incorporate the attention layers to capture more semantic details, which can be accelerated through the Kronecker product.
For the diversity, Quad clusters the dataset into similar data instances within each cluster and diverse instances across different clusters. For each cluster, if we opt to select data from it, we take some samples to evaluate the influence to prevent processing all instances. Overall, we favor clusters with highly influential instances (ensuring high quality) or clusters that have been selected less frequently (ensuring diversity), thereby well balancing between quality and diversity. Experiments on Slimpajama demonstrate that  Quad significantly outperforms other data selection methods with a low FLOPs consumption. Further analysis also validates the effectiveness of our influence calculation. Our code and data are available at ( https://anonymous.4open.science/r/Quad/ ).

## 1 Introduction

Recently, large language models (LLMs) have significantly advanced the field of artificial intelligence (Zhao et al., ; Hadi et al., ; Minaee et al., ). Due to the unprecedented number of parameters (model size) and the pre-training on huge amount of training data, LLMs are generalizable a broad spectrum of downstream tasks.
However, in practice, the computation resources limit both the model size and the volume of data used in pre-training. In this situation, judiciously selecting train datasets is critical for producing highly performance LLMs (Brown, ; Du et al., ; Gururangan et al., ; Hoffmann et al., ; Raffel et al., ). In particular, the quality of the training datasets vary dramatically, while the LLaMA-3.1 report (Dubey et al., ) shows that the use of high quality data in later training stages can greatly improve model performance.

Typical straightforward data selection approaches include rule-based data filtering (Raffel et al., ; Rae et al., ), querying high-performance models (e.g., GPT-4) (Wettig et al., ; Sachdeva et al., ), surrogate models (Lin et al., ; Shao et al., ), etc. Although these methods have achieved success on some datasets and models, they rely on simple heuristics to select training data. Without explicitly measuring the impact of the selected data on the model, these methods tend to produce sub-optimal pretraining results.
To address this issue, some researchers (Xia et al., ; Yu et al., ) start evaluating each data instance by assigning it a score that reflects its impact on the model. Frequently used scoring methods include the influence function (Xia et al., ), early loss (Albalak et al., ), and perplexity (Chen et al., ). Among these methods, the influence function consistently delivers state-of-the-art results by effectively approximating the impact of adding each instance to the training set.
A higher score signifies a higher priority for selecting a data instance, and hence the top-$k$ (or gumble top-$k$) instances with the highest scores are chosen (Xie et al., ; Wettig et al., ; Yu et al., ).

However, the above methodologies have the following limitations.

Prohibitive Computation Cost.
First, accurately calculating the influence score of one data instance is expensive, because it involves the computation of the Hessian matrix. However, in the LLM pre-training, the number of the candidate data instances is extremely large. It is thus prohibitively expensive to compute the scores for all of the candidates.

Lack of Diversity.
Second, assume that all influence scores have been calculated, as shown in Figure 1. We can see that the top-$k$ instances ($e.g.,$ some high-score instances in $C_{1}$) tend to be closely distributed in the feature space because the influence computation is closely related to the data features.
That is, the training instances selected in this way are lack of diversity ($e.g.,$ other instances in $C_{3}$ with high influence are also worth selecting), while as confirmed by some studies (Abbas et al., ; Tirumala et al., ), diversifying training samples mitigates overfitting, thereby enhancing the generalizability of the model.
Therefore, an effective data training selection method should take both the influence scores and the diversity into consideration.

Figure: (a) Influence scores of data instances
Refer to caption: extracted/5903369/intro-figure-new1.png

We thus propose Quad, a scalabe and effective data selection approach, which successfully addressing above challenges, achieves state-of-the-art pretraining results. Initially, Quad organizes the given dataset into clusters where the data instances within each cluster are similar, and those in different clusters exhibit diversity. Hence, we can sample a data subset from a cluster to estimate the accurate average influence of the cluster, so as to represent the cluster quality $w.r.t$ the model performance.

Next, leveraging the property of the attention-based Transformer architecture which is widely adopted by the LLMs, we design a novel method to accurately compute the influence of an instance on LLM pre-training. More specifically, rather than solely relying on the MLP layers to compute the influence (Koh & Liang, ; Yu et al., ; Grosse et al., ; Engstrom et al., ), we incorporate the attention layers such that the influence computation considers more semantic information. In addition, given that calculating the Hessian matrix is time-consuming, particularly for attention layers with complex interactions, we incorporate the Kronecker product to approximate the Hessian matrix, thereby greatly expediting the computation. This successfully addresses the computation cost challenge.

To improve diversity, we apply the Multi-Arm Bandit (MAB) technique, where each cluster is regarded as an arm of the MAB. Upon selecting an arm, we draw samples from the cluster to calculate influence scores. Subsequently, Quad iteratively samples from clusters, taking into account both the influence score and data diversity, e.g., whether the cluster has already been sampled. Moreover, because this sampling strategy effectively avoids calculating the influence of all instances, it further speeds up the data selection process.

We summarize our main contributions as follows:

- •
To balance the quality and diversity, we incorporate an iterative MAB solution to first cluster the data instances and select data instances from these clusters.
- •
We propose a novel method to compute the influence function in attention-based Transformer architecture, so as to precisely measure the data quality in LLM pre-training.
- •
Experiments on the widely-used dataset Slimpajama and 9 popular downstream tasks demonstrate that  Quad significantly outperforms state-of-art data selection methods by 1.39% in zero-shot accuracy, also with low computation resources consumption.

## 2 Related Work

Rule-based Methods.
Initially, researchers often relied on intuition to design hand-crafted heuristics (Soldaini et al., ) and  (Penedo et al., ), aiming to improve data quality. Deduplication is another typical approach for selecting pretraining data, such as  (Penedo et al., ) and SemDedup  (Abbas et al., ) which use keyword-based and semantic deduplication, respectively. Additionally, certain approaches employ $n$-gram similarity (Gao et al., ; Xie et al., ) to assist in choosing corpora that is semantically aligned with the validation set data.
Although these methods effectively filter out noise and redundant data from web sources, they rely on simple heuristics and cannot be well generalized.

LLM As a Selector.
Although large models such as GPT-4 can effectively assess data quality due to their semantic comprehension capacity, the metrics utilized to rate data ($e.g.,$ writing style, educational value etc.) heavily rely on human intuition (Wettig et al., ; Penedo et al., ; Zhang et al., ; Gunasekar et al., ). This often leads to a mismatch between the selected data and the data desired by the model.

Surrogate Models.
DeepSeekMath  (Shao et al., ) proposes an active learning strategy to train a web data classifier. Similarly, in MATES (Yu et al., ), a surrogate model was developed to estimate the influence scores of the data instances. RHO-1  (Lin et al., ) used a surrogate model trained with high-quality data to perform token-level data filtering.
However, these surrogate models are not trained over large-scale data, and thus their generalizatio ability is limited.

Perplexity serves as a metric for selecting high-probability data in a language model. In  (Chen et al., ; Marion et al., ; Muennighoff et al., ; Wenzek et al., ), perplexity (PPL) is utilized to filter data. As also discussed in Qurating (Wettig et al., ), we observe that this method often incorporates a significant amount of simple and redundant data, because they are easy for the model to predict.

Influence Function
(Grosse et al., ; Choe et al., ) demonstrates that influence function can reveal the impact of training data on the performance of large models. Consequently, LESS (Xia et al., ) and MATES  (Yu et al., ) utilize influence functions for selecting data during the SFT and pretraining phases, respectively. For large models, computing influence functions is computationally expensive.  (Grosse et al., ). Hence, given the large amount of data handled during pretraining, directly using LESS (Xia et al., ) for data selection at this stage poses considerable difficulties. To overcome this, MATES (Yu et al., ) employs a proxy model to approximate the influence score across the full dataset.
However, the limited capacity of this small proxy model hinders its ability to provide accurate influence scores.
Furthermore, relying on the influence to select data solely often leads to a lack of diversity in the chosen data.

## 3 Methods

First, we present our problem statement in §3.1. Next, in §3.2, we explain how our method achieves the balance between quality and diversity in selecting pretraining data. Finally, in §3.3, we
introduce how we compute the influence with attention layers more accurately and efficiently.

Figure: Figure 2: Overview of Quad
Refer to caption: extracted/5903369/mab-new.png

### 3.1 Problem Definition

To enhance the capabilities of large models, it is necessary to retrieve relevant data from a large pool of candidate data and perform further training for the large model.
Formally, given the pool $D_{c}$ and a reference set $D_{r}$, our problem is to select a subset $D_{b}\subset D_{c}$ to fine-tune the large model $M$, with the aim of minimizing the loss of the updated model $M^{\prime}$ in the reference set $D_{r}$.

### 3.2 balance between quality and diversity

As shown in Figure 1, there are significant variations in the distribution of influence scores among different clusters.
To achieve the quality-diversity balance, it is necessary to know the precise average influence score for instances in each cluster.
However, Figure 1 shows that the influence scores for each cluster also fluctuate around the average, indicating a certain level of uncertainty. Estimating the average with a small sample size will not be accurate enough, while taking a large number of samples to compute the average influence is costly.

Hence, we propose to use the MAB (Vermorel & Mohri, ) technique that is capable of making decisions iteratively under uncertainty. At a high level, each cluster represents an arm of the MAB, and during each iteration, a cluster with a high average influence score tends to be selected and sampled. We will then compute the influence of data instances to update the average. Moreover, clusters that are not visited often present significant opportunities for sampling to balance the diversity.

The overall process of this approach is illustrated in Figure 2. Specifically, our method can be divided into the following four steps: First, we sample the top-$k$ clusters with the highest cluster scores (denoted by CS) computed by MAB.
Here, the cluster score is determined by both the influence score and the sample frequency.
Then we calculate the influence scores for the samples in each cluster (Section 3.3). At this point, we select high scoring samples to be added for training and use their scores to update the cluster score for each cluster.
Throughout the iterative process, the MAB algorithm focuses on frequently sampling high-quality clusters that have high influence scores, which also enhances the accuracy of their quality estimation ($i.e.,$ updating the average influence $\overline{I}_{i}$). Simultaneously, it ensures diversity by also sampling less-visited clusters.
Next, we discuss how to compute and update the cluster score in details.

Cluster Score (CS). The Upper Confidence Bound can effectively balance exploration ($i.e.,$ data diversity) and exploitation ($i.e.,$ data quality), so we use it as the cluster score to evaluate each cluster, as shown in Equation (1).
Specifically, the cluster score is determined by the average influence score $\bar{I_{i}}$ and the exploration score $\sqrt{\frac{2\ln{\sum_{j}{T(C_{j})}}}{T(C_{i})}}$, where $T(C_{i})$ denotes the frequency of instances sampled from cluster $C_{i}$, and $\sum_{j}T(C_{j})$ denotes the total times of samples taken from all clusters.

$$ $CS_{i}=\bar{I_{i}}+\alpha\sqrt{\frac{2\ln{\sum_{j}{T(C_{j})}}}{T(C_{i})}}$ (1) $$

Update the cluster score. During each iteration, a subset of data $B_{i}$ is sampled from each cluster with a high cluster score (CS).
The sum of their influence score $I_{i}$ can be used to denote the impact of the samples from the cluster $C_{i}$ on the model.

$$ $R(C_{i})+=\sum_{z\in B_{i}}I_{\theta}(D_{r},z),\quad T(C_{i})+=1$ (2) $$

where $R(C_{j})$ denotes the total reward accumulated by cluster $C_{i}$ over several iterations.Then the average influence score $\bar{I_{i}}$ for cluster $C_{i}$ can be represented as
$\bar{I_{i}}=\frac{R(C_{i})}{T(C_{i})}$.
As the sample size grows, $\bar{I_{i}}$ for each cluster $C_{i}$ steadily approaches the exact average influence of the cluster, which can be used to update the cluster scores for all clusters.

Data selection. During each iteration, we pick a small proportion($\gamma$) of data instances from selected clusters. We also require that these instances have influence scores higher than the threshold $\tau$, otherwise we will not select them, which are then added into the training dataset.

Figure: Algorithm 1 Quad Algorithm

### 3.3 Influence Calculation with attention layers

Instead of retraining the large model with each data sample $z$, the impact of $z$ on the model $M$ can be estimated by calculating the influence function for each instance. In this section, we extend the influence calculation to multi-head attention layers and provide acceleration techniques.

$$ $I_{\theta}(D_{r},z)=-\nabla L(\theta,D_{r})(H+\lambda I)^{-1}\nabla L(\theta,z)$ (3) $$

In the above equation, $I_{\theta}(D_{r},z)$ denote the influence function of data $z$ on model $\theta$. $\nabla L(\theta,D_{r})$ and $\nabla L(\theta,z)$ denote the gradient of reference dataset $D_{r}$ and data $z$, respectively. Since the training of the large model does not often fully converge, resulting in a non-invertible Hessian matrix $H$, a regularization term $\lambda I$ is introduced (Bae et al., ). Equation (3) is typically divided into the following two stages to speed up the computation:

1. Approximate the multiplication of the gradient of the validation set $\nabla L(\theta,D_{r})$ and the inverse Hessian matrix $H^{-1}$ using the inverse Hessian vector product (iHVP).

2. Compute the dot product between the iHVP and the gradient of each training data point $\nabla L(\theta,z)$.

While this framework can accelerate the computation of the influence function, scaling it up to large language models (LLMs) with massive parameters is still expensive. Hence, K-FAC (Martens & Grosse, ; Ueno et al., ) can be used to accelerate the iHVP computation by using the Kronecker product to decompose the Hessian matrix.

The K-FAC approximate the parameters of different MLP layer $\theta_{1}$, $\theta_{2}$ and $\theta_{3}$ as independent. That’s because, during the gradient computation and update process, there are usually only minimal direct dependencies between the gradients of different MLP layers. This is particularly evident during back propagation, where the weight updates for each MLP layer are primarily influenced by the parameters of that specific layer. Therefore, the influence function $I_{\theta_{1},\theta_{2},\theta_{3}}(D_{r},z)$ in K-FAC method can be expressed as:

$$ $I_{\theta_{1},\theta_{2},\theta_{3}}(D_{r},z)=I_{\theta_{1}}(D_{r},z)+I_{ \theta_{2}}(D_{r},z)+I_{\theta_{3}}(D_{r},z)$ (4) $$

In attention mechanisms, there exist complex connections between the Query, Key, and Value layers. As the right-upper corner of Figure 3 shows, separately calculate the hessian matrix of Query, Key and Value layers, will miss massive information of
Consequently, it is essential to consider the QKV layers as a unified layer $\theta_{qkv}$ when computing the influence function. Therefore, the influence function $I_{\theta_{att}}(D_{r},z)$ can be expressed as:

$$ $I_{\theta_{att}}(D_{r},z)=I_{\theta_{qkv}}(D_{r},z)+I_{\theta_{o}}(D_{r},z)$ (5) $$

Figure: Figure 3: Kronecker Product in calculating iHVP
Refer to caption: extracted/5903369/multi-head-attention.png

Then, as the right-lower corner of Figure 3 shows, by decomposing the Hessian matrix into a kronecker product of smaller matrices and computing the inverse of each smaller matrix, we can avoid directly inverting the entire Hessian matrix, significantly reducing computational cost, and accelerate this process:

Forward propagation:

$$ $Attention(Q,K,V)=softmax(\frac{QK^{T}}{\sqrt{d_{k}}})V$ (6) $$

Backward propagation:

$$ $D\theta=vec(DW)=\delta\otimes x$ (7) $$

Here, $\otimes$ denotes the Kronecker product, and $vec()$ represents the vectorization operation. Thus, the gradient of $\theta_{qkv}$ can be written as:

$$ $D\theta_{qkv}=\begin{bmatrix}\mathbf{vec(DW_{Q})}\\ \mathbf{vec(DW_{K})}\\ \mathbf{vec(DW_{V})}\end{bmatrix}\\ =\begin{bmatrix}\mathbf{\delta}_{q}\\ \mathbf{\delta}_{k}\\ \mathbf{\delta}_{v}\end{bmatrix}\otimes x$ (8) $$

Let $\delta_{qkv}$
$=\begin{bmatrix}\mathbf{\delta}_{q}\\
\mathbf{\delta}_{k}\\
\mathbf{\delta}_{v}\end{bmatrix}$. Then, the Hessian matrix $H_{qkv}$ can be estimates by:

$$ $\displaystyle H_{qkv}= E(D\theta_{qkv}{D\theta_{qkv}}^{T})=E(\delta_{qkv} \delta_{qkv}^{T}\otimes x_{qkv}x_{qkv}^{T})$ (9) $\displaystyle\approx E(\delta_{qkv}\delta_{qkv}^{T}) \otimes E(x_{qkv}x_{qkv}^ {T}) = \Delta_{qkv} \otimes X_{qkv}$ $$

Also, $H_{o}=\Delta_{o}\otimes X_{o}$. Thus, the iHVP of the attention layer can be estimated as follows:

$$ $\displaystyle H_{att}^{-1}v_{att}=\begin{bmatrix}\mathbf{H_{qkv}^{-1}v_{qkv}} \\ \mathbf{H_{o}^{-1}v_{o}}\end{bmatrix}=\begin{bmatrix}\mathbf{(\Delta_{qkv} \otimes X_{qkv})^{-1}v_{qkv}}\\ \mathbf{(\Delta_{o}\otimes X_{o})^{-1}v_{o}}\end{bmatrix}$ (10) $\displaystyle=\begin{bmatrix}\mathbf{(\Delta_{qkv}^{-1}\otimes X_{qkv}^{-1})v_ {qkv}}\\ \mathbf{(\Delta_{o}^{-1}\otimes X_{o}^{-1})v_{o}}\end{bmatrix}=\begin{bmatrix} \mathbf{vec(\Delta_{qkv}^{-1}V_{qkv}X_{qkv}^{-1})}\\ \mathbf{vec(\Delta_{o}^{-1}V_{o}X_{o}^{-1})}\end{bmatrix}$ $$

where $v_{att}$, $v_{qkv}$, $v_{o}$ represent the gradient of reference dataset $D_{r}$ on parameters $\theta_{att}$, $\theta_{qkv}$, $\theta_{o}$, respectively. Thus, the influence score of attention layers can be written as: $I_{\theta_{att}}=-\nabla L(\theta_{att},z)H_{att}^{-1}v_{att}$.

To avoid the excessive memory usage of validation set gradients, we apply the Johnson-Lindenstrauss Lemma to reduce the dimensionality of both the iHVP computation results and the training data gradients $\nabla L(\theta,z)$.

## 4 Experiment

### 4.1 Experiment Setup

Dataset Preparation.
We use the entire 627B-token SlimPajama dataset as the candidate pool $D_{c}$. In the clustering process, the BAAI/bge-large-en-v1.5 model is employed to generate embeddings for the input data, and approximately 600 million data points from the candidate pool $D_{c}$ are clustered into 10,000 groups using the $k$-means algorithm. We use LAMBADA  (Paperno et al., ) as our reference set $D_{r}$, which is a widely used language modeling task and often serves as a validation benchmark for language model pretraining.  (Yu et al., ; Xie et al., ; Hoffmann et al., ).

Experimental settings.
We train a transformer-based decoder-only language model that contains 1.3B parameters, uses RoPE embeddings  (Su et al., ), and has a maximum context window of 1024 tokens  (Touvron et al., ).
Following the setting of MATES (Su et al., ), 30B tokens out of the 627B are selected for training using Quad and compare with baselines.
The learning rate is set to $5\times 10^{-5}$, the batch size is set to 4096, and the Adam optimizer is employed with hyperparameters $\beta_{1}=0.9,\beta_{2}=0.95,\epsilon=10^{-8}$. As for Multi-Armed Badit, we set the $\alpha$ = 0.002 , sample proporation $\gamma$ = 0.05 and the sample threshold $\tau$ as 0.0025.

Baselines. We compare our methods with several baselines. (1) Random samples data from the entire candidate dataset randomly. (2) Qurating uses the large language model to select data. (3) DSIR selects data instances that are similar to the LAMBADA dataset. (4) PPL uses perplexity-based data selection, $i.e.,$ selecting data instances with the lowest perplexity scores. (5) MATES trains a surrogate model to evaluate the influence of each data instance on the target model.

Evaluation datasets.
To comprehensively evaluate the capabilities of pretrained models, we conduct experiments on various downstream tasks covering three significant categories:

General Knowledge: ARC-C, ARC-E (Clark et al., ), and SciQ  (Welbl et al., ).

Commonsense Reasoning: HellaSwag  (Zellers et al., ), SIQA  (Sap et al., ), WinoGrande (Sakaguchi et al., ), Logiqa  (Liu et al., ).

Reading Comprehension: OpenbookQA  (Mihaylov et al., ), and BoolQ (Clark et al., ).

Evaluations are conducted using the lm-evaluation-harness (Gao et al., ) framework and the average accuracy ($i.e.,$ Overall Score) is reported for comparison.

### 4.2 Results

Overall Performance.
As demonstrated in Table 1, our method surpasses all the baseline methods in downstream tasks with zero-shot evaluation.
To be specific, we can observe that on General Knowledge and Reading Comprehension tasks, Quad has the improvement of 1.75% and 1.98% respectively compared with Random.
Quad outperforms DSIR and Semdedup because they use rule-based heuristics to select data without considering the model. Although PPL and MATES consider the model, they do not perform well because the former one always selects some simple and duplicated instances, and the surrogate model of the latter one is small and lacking of enough training data. Qurating generally performs the best among other baselines, but still worse than our approach, and it incorporates the highest FLOPS(1e19) because of the usage of LLMs for data selection.
In terms of the FLOPs, we can observe that except the methods ($i.e.,$ DSIR, SemDeDup) that use simple heuristics, we consume minimal computation resources because our MAB solution samples from clusters without considering the entire candidate dataset like PPL, Qurating and MATES.

**Table 1: Overall Performance**
| Selection Method | General<br>Knowledge<br>(3 tasks) | Commonsense<br>Reasoning<br>(4 tasks) | Reading<br>Comprehension<br>(2 tasks) | Overall | FLOPs |
| --- | --- | --- | --- | --- | --- |
| Random | 50.33 | 36.19 | 39.09 | 41.55 | 7.66 |
| DSIR | 50.37^$\uparrow$0.04 | 34.01^$\downarrow$2.18 | 38.80^$\downarrow$1.29 | 40.53^$\downarrow$1.02 | 7.66 |
| PPL | 48.71^$\downarrow$1.62 | 37.72^$\uparrow$1.53 | 38.57^$\downarrow$0.52 | 41.57^$\uparrow$0.02 | 9.51 |
| Semdedup | 50.99^$\uparrow$0.66 | 36.11^$\downarrow$0.08 | 39.44^$\uparrow$0.35 | 41.81^$\uparrow$0.26 | 8.11 |
| Qurating | 51.56^$\uparrow$1.23 | 35.93^$\downarrow$0.26 | 39.70^$\uparrow$0.61 | 42.01^$\uparrow$0.46 | 13.66 |
| MATES | 50.45^$\uparrow$0.12 | 36.06^$\downarrow$0.13 | 39.83^$\uparrow$0.74 | 41.93^$\uparrow$0.38 | 9.81 |
| Quad(ours) | 52.08^$\uparrow$1.75 | 37.03^$\uparrow$0.84 | 41.07^$\uparrow$1.98 | 42.94^$\uparrow$1.39 | 9.15 |

Effectiveness of MAB.
This section evaluates the effectiveness of the MAB approach for data selection in contrast to the straightforward method of choosing the top-$k$ clusters with the highest influence scores for model training. To be specific, we randomly select an equivalent number of data points from the top 150, 500, and 1000 clusters. Figure 4 illustrates the trade-off between data quality and diversity: clusters with higher influence scores do not necessarily enhance model performance on downstream evaluation sets because of their lack of diversity. Hence, the multi-armed bandit method can more effectively capture the trade-off between quality and diversity across clusters, resulting in superior performance on downstream evaluation sets, as opposed to merely choosing the top-$k$ clusters.

Figure: (a) MAB vs. top-$k$ clusters
Refer to caption: extracted/5903369/exp-fig1.png

Effectiveness of Influence Calculation.
This experiment studies the effectiveness of our influence calculation method. In this section, we select the top 500 clusters with the highest scores using three methods: (1) no-Hessian ($i.e.,$ computing the gradient similarity between training data and reference data (Pruthi et al., )) without considering the Hessian matrix; (2) MLP($i.e.,$ calculating influence function on MLP layers) and (3) Ours ($i.e.,$ calculating influence function on both MLP and attention layers).
From each cluster, we uniformly sample data to train the large model. As shown in Figure 4, our solution (MLP+Attention) performs better than MLP because the attention layer considers more semantics. no-Hessian performs the worst because it does not precisely capture the impact of training data instances on the model without the Hessian matrix.

Also, we conduct experiments to verify the relationship between the Query, Key, Value matrices, which is shown in Figure 4.
In this experiment, we compare the Pearson correlation coefficients between the following three methods and the baseline approach, which computes the influence score for the attention layer without any acceleration. (1) No-Hessian($i.e.,$ computing the gradient similarity between training data and reference data) without considering the Hessian matrix; (2) Independent ($i.e.,$ calculating the Hessian matrices of the query, key, and value layers independently) and (3) Ours ($i.e.,$ calculating the Hessian matrices of the query, key, and value layers as a whole).

### 4.3 Ablation Study

This group of experiments performs ablation studies on the hyperparameters of Quad. Figure 4 and Figure 5 show the impact of sample threshold and $\alpha$ respectively.

Sampling Threshold of Influence ($\tau$).
Setting the threshold too high or low will both degrade the model performance. This is because the selected data instances tend to exist in few clusters with high influence scores, resulting in poor diversity. In contrast, when the threshold is set too low, the sampled instances will be from many clusters with low influence scores, which also degrades the model performance.

$\alpha$ for Quality-Diversity Balance.
Our approach employs $\alpha$ to balance the diversity and quality in the MAB framework.
When $\alpha$ is small, we tend to focus on the several clusters with high influence scores without considering diversity much, so the MAB framework is likely to get stuck in a local optimum.
For example, this results in the model enhancing its performance in specific areas (such as Common Sense Reasoning in Figure 5 when $\alpha=1.5e-3$), while the performance in other areas ($i.e.,$ General Knowledge and Reading Comprehension) is not good enough. Thus the overall score is not the optimal.
However, when $\alpha$ is large, the MAB framework focuses too much on diversity without selecting enough high-quality data, which ultimately results in a limited improvement of model performance.

Figure: Figure 5: $\alpha$ for Quality-Diversity Balance.
Refer to caption: extracted/5903369/exp-fig5.png

## 5 conclusion

This paper presents Quad, a method designed to balance both the diversity and quality of data in pretraining data selection. Quad employs the influence function to identify data that benefits the model. First, we group the data into clusters and use a subset from each to represent the influence of the entire cluster. Given that influence scores within a cluster display some uncertainty, we view each cluster as an arm in an MAB framework. This method conducts samplings from high-quality clusters, allowing for more precise estimation of their influence scores and meanwhile maintaining the diversity. Moreover, we extend the influence function to attention layers and enhance the calculation efficiency to better measure the impact of data within each cluster on the model.

## Appendix A Appendix

**Table 2: Performance Comparison**
| Selection Method | General Knowledge | Commonsense Reasoning | Reading Comprehension | Overall |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | arc-e | arc-c | sciq | avg | logiqa | hellaswag | siqa | winogrande | avg | openbookqa | boolq | avg |  |
| Random | 50.27 | 20.31 | 80.40 | 50.33 | 21.20 | 34.11 | 38.49 | 50.99 | 36.20 | 17.60 | 60.58 | 39.09 | 41.55 |
| Semdedup | 51.35 | 20.73 | 80.90 | 50.99 | 19.05 | 34.56 | 39.30 | 51.54 | 36.11 | 18.80 | 60.09 | 39.45 | 41.81 |
| MATES | 50.00 | 21.25 | 80.10 | 50.45 | 21.66 | 33.90 | 38.69 | 52.17 | 36.61 | 19.00 | 60.67 | 39.84 | 41.94 |
| PPL | 45.41 | 20.82 | 79.90 | 48.71 | 20.43 | 35.92 | 39.92 | 54.62 | 37.72 | 18.80 | 58.35 | 38.58 | 41.57 |
| DSIR | 49.28 | 20.14 | 81.70 | 50.37 | 21.20 | 30.89 | 35.98 | 47.99 | 34.02 | 16.20 | 61.41 | 38.81 | 40.53 |
| Qurating | 52.10 | 23.29 | 79.80 | 51.56 | 20.43 | 33.57 | 39.05 | 50.67 | 35.93 | 18.00 | 61.41 | 39.71 | 42.04 |
| Quad(ours) | 52.27 | 21.76 | 82.20 | 52.08 | 22.89 | 34.41 | 38.74 | 52.09 | 37.03 | 20.00 | 62.14 | 41.07 | 42.94 |
| top$k$-cluster |  |  |  |  |  |  |  |  |  |  |  |  |  |
| top-150 | 48.61 | 20.90 | 79.00 | 49.50 | 23.66 | 34.51 | 39.00 | 51.78 | 37.24 | 19.20 | 61.74 | 40.47 | 42.04 |
| top-500 | 51.05 | 21.25 | 79.70 | 50.67 | 22.73 | 34.40 | 39.20 | 52.41 | 37.19 | 18.80 | 62.76 | 40.78 | 42.48 |
| top-1000 | 49.96 | 20.99 | 80.40 | 50.45 | 21.97 | 34.00 | 38.74 | 50.2 | 36.23 | 18.20 | 60.61 | 39.41 | 41.67 |

**Table 3: Ablation Study of Threshold $\tau$**
| Threshold | General Knowledge | Commonsense Reasoning | Reading Comprehension | Overall |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | arc-e | arc-c | sciq | avg | logiqa | hellaswag | siqa | winogrande | avg | openbookqa | boolq | avg |  |
| 0.0015 | 51.26 | 21.16 | 80.20 | 50.87 | 21.51 | 33.92 | 39.00 | 51.07 | 36.38 | 19.60 | 61.74 | 40.67 | 42.16 |
| 0.0020 | 52.23 | 22.27 | 80.70 | 51.73 | 22.89 | 34.77 | 38.33 | 50.20 | 36.55 | 19.20 | 61.50 | 40.35 | 42.45 |
| 0.0025 | 52.27 | 21.76 | 82.20 | 52.08 | 22.89 | 34.41 | 38.74 | 52.09 | 37.03 | 20.00 | 62.14 | 41.07 | 42.94 |
| 0.0030 | 50.25 | 19.62 | 80.80 | 50.22 | 22.27 | 33.96 | 38.96 | 53.28 | 37.12 | 20.60 | 59.20 | 39.90 | 42.10 |

**Table 4: Ablation Study of $\alpha$**
| Alpha | General Knowledge | Commonsense Reasoning | Reading Comprehension | Overall |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | arc-e | arc-c | sciq | avg | logiqa | hellaswag | siqa | winogrande | avg | openbookqa | boolq | avg |  |
| 0.0015 | 50.55 | 20.73 | 79.70 | 50.33 | 23.35 | 34.60 | 40.58 | 52.25 | 37.70 | 19.60 | 60.06 | 39.83 | 42.38 |
| 0.0020 | 52.27 | 21.76 | 82.20 | 52.08 | 22.89 | 34.41 | 38.74 | 52.09 | 37.03 | 20.00 | 62.14 | 41.07 | 42.94 |
| 0.0025 | 51.64 | 22.10 | 78.30 | 50.68 | 21.81 | 34.76 | 38.02 | 52.57 | 36.79 | 19.80 | 62.05 | 40.93 | 42.34 |
| 0.0030 | 50.63 | 21.93 | 78.70 | 50.42 | 22.12 | 34.64 | 39.15 | 51.14 | 36.76 | 18.60 | 61.71 | 40.16 | 42.07 |

**Table 5: Effectiveness of Influence Calculation**
| Method | General Knowledge | Commonsense Reasoning | Reading Comprehension | Overall |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | arc-e | arc-c | sciq | avg | logiqa | hellaswag | siqa | winogrande | avg | openbookqa | boolq | avg |  |
| Random | 50.27 | 20.31 | 80.40 | 50.33 | 21.20 | 34.11 | 38.49 | 50.99 | 36.20 | 17.60 | 60.58 | 39.09 | 41.55 |
| No-Hessian | 49.03 | 20.99 | 80.50 | 50.20 | 22.58 | 33.40 | 38.89 | 52.41 | 36.82 | 19.20 | 61.50 | 40.35 | 42.06 |
| MLP | 50.63 | 21.50 | 78.90 | 50.34 | 22.89 | 33.32 | 38.74 | 52.57 | 36.88 | 19.60 | 61.77 | 40.69 | 42.21 |
| Ours | 51.05 | 21.25 | 79.70 | 50.67 | 22.73 | 34.40 | 39.20 | 52.41 | 37.19 | 18.80 | 62.76 | 40.78 | 42.48 |

**Table 6: Model Architecture**
| Hyperparameter | Value |
| --- | --- |
| Vocabulary Size | 32,000 |
| MLP Ratio | 8/3 |
| Hidden Dimension Size | 2048 |
| Number of Layers | 24 |
| Number of Attention Heads | 16 |
| Number of KV Attention Heads | 16 |
| RoPE Base | 10,000 |
| Maximum Context Window Length | 1024 |
| Number of Parameters | 1,345,423,360(1.3B) |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Clustering Quality and Sensitivity: How sensitive is the clustering quality to the choice of the number of clusters (10,000 in this case)? What criteria were used to determine the optimal number of clusters, and how might different numbers of clusters affect the performance of the MAB approach?

2. Generalizability and Robustness: The experiments were conducted on a specific set of downstream tasks and a particular reference set (LAMBADA). How well does the method generalize to other types of tasks or datasets? Have you tested the method on a broader range of tasks or datasets to ensure its robustness?
soundness: 3
strengths: 1. Clear Problem Definition: The paper clearly defines the problem of enhancing large models by selecting a subset of data from a large pool to fine-tune the model, aiming to minimize the loss on a reference set. This provides a concrete and well-defined goal for the research.

2. Balanced Quality and Diversity: The authors use Multi-Armed Bandit (MAB) techniques to balance quality and diversity in data selection. By treating each cluster as an arm and using the Upper Confidence Bound (UCB) to determine cluster scores, the method ensures frequent sampling of high-quality clusters while also exploring less-visited clusters to maintain diversity.

3. Efficient Influence Calculation: The paper introduces an efficient method to estimate the influence of individual data points on the model without full retraining. By extending the influence function to multi-head attention layers and using K-FAC for acceleration, the authors provide a scalable solution for large language models, significantly reducing computational costs.
summary: This paper explores the significance of data selection in the pretraining of large language models (LLMs) and the limitations of existing methods, including high computational costs and a lack of diversity. To address these issues, the authors propose a new method called Quad, which aims to balance data quality and diversity. Quad clusters the data and samples from each cluster to assess its influence, reducing computational burden while ensuring diversity. Additionally, Quad leverages attention mechanisms and the Kronecker product to efficiently compute data influence, further enhancing computational efficiency. Experimental results show that Quad significantly outperforms existing methods on the Slimpajama dataset and nine downstream tasks, achieving a 1.39% improvement in zero-shot accuracy with lower computational resource consumption. Overall, Quad provides an effective and scalable data selection approach, offering a significant breakthrough for LLM pretraining.
weaknesses: 1. Trade-off Between Sample Size and Accuracy: Although the MAB technique can make decisions under uncertainty, there is still a trade-off between sample size and accuracy when estimating the average influence score. Small sample sizes may lead to inaccurate estimates, while large sample sizes increase computational costs. Finding the optimal sample size in practical applications remains a challenge.

2. Complexity and Computational Resource Requirements: Despite the introduction of K-FAC and other acceleration techniques, the computational complexity and resource requirements remain high when dealing with large-scale language models. In environments with limited resources, this method may be difficult to implement or may require significant computation time.

4. Assumptions and Model Dependence: The method relies on certain assumptions when calculating influence scores, such as the approximation of the Hessian matrix and the independence of gradients. These assumptions may not hold in some cases, affecting the effectiveness and reliability of the method. Additionally, different model architectures may require different adjustments and optimization strategies.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: Could the authors report the results of statistical significance tests?

How does the selection of clustering algorithm and number of cluster affect the performance of the proposed method?

Can the authors clarify how they compute the FLOPS column in Table 1?
soundness: 3
strengths: The paper addresses the important challenge of data selection via a novel method and practical acceleration techniques. Influence function is an expensive operation and the proposed techniques make it more accessible for LLMs where computational overhead is an important factor. 

The paper presents a comprehensive evaluation on downstream tasks as well as ablations and statistics to help understand the proposed method. 

Despite the method's complexity, the paper is overall clearly written and easy to follow. The authors also attach the code, which helps the community to adapt this method for custom usecases.
summary: The paper presents a method for data selection, namely Quad, which (i) uses the influence function with multi-head attention layers to estimate the potential impact of individual samples on the model's performance; (ii) uses clustering and multi-armed bandits (where each cluster is an arm) to balance between quality and diversity. The paper further proposes acceleration techniques to reduce computational cost and memory usage. The proposed method is applied on the SlimPajama dataset by training a 1.3B model, evaluated on a set of common text benchmarks, and compared with a set of data selection baselines.
weaknesses: The proposed method requires a complex set of steps and hyperparameters (e.g., $\alpha$,  $\gamma$, and $\tau$ for multi-armed bandit) and, even though there are ablations for some of these, it 
might be hard and resource-intensive to develop and tune in practice for diverse applications.

The experimental findings show that all evaluated methods have limited difference in downstream performance and it is not clear whether the observed improvements are statistically significant. It is also concerning that baselines like DSIR show no benefit compared to random selection, which might be contradicting with findings from previous work. Therefore, I would recommend a more thorough discussion behind the empirical benefits of the proposed method and the connection of the findings compared to the previous literature. To better demonstrate the benefits, it might be worth (i) considering more experimental settings, for example with models of varying sizes; (ii) plotting the training and validation losses for each data selection method.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 2
questions: 1. Explain figure 1(a) better, what is the x- and y-axis?
2. How do you apply the “JohnsonLindenstrauss Lemma” (line 343)
3. Why only select “30B tokens out of the 627B” (line 360), what if we just do a coarse-grained filtering of the original 627B to leave around 400-500B data (e.g., via near-duplicate removal) and use all of them to do LLM pretraining, will the model performance better?
4. Is the Quad method robust to the clustering algorithms? What features are used in the clulstering?
soundness: 3
strengths: 1. The paper presents a framework that is overall reasonable.
2. The influence calculation with attention layers technique is interesting and could potentially be useful in other settings.
3. Experiments (of the current limited settings) are complete and the authors do release the code for better reproducibility
summary: This paper presents a new method for selecting LLM pretraining data. The core idea is to measure data quality with diversity. The authors introduce a new approach named Quad, which utilizes attention layers for influence calculation, and clusters data to maintain diversity. Experimental results on the Slimpajama dataset show that Quad surpasses other data selection methods in performance.
weaknesses: The main limitation of this study is its experiment scale is too small (in terms of claiming the method for LLM pretraining). Only 30B out of 627B data are selected for training a small 1.3B decoder-only LLM. In practice, most LLMs will start with at least 7B so how the technique can generalize to reasonably sized LLMs is questionable. Also related to this limitation is the diversity of testing datasets, little (long-text) generation tasks are evaluated.

The second limitation of this study is that the authors do not fully motivate the problem setting. While cleaning noisy data and removing near-duplicate examples are useful for LLM pretraining, I am not fully convinced why you have to select a subset of "quality data" from the full dataset. For example, if we do a coarse-grained filtering of the original 627B to leave around 400-500B data and use all of them to do LLM pretraining, will the model performance better?

Finally, there are a few presentation issues such as the unclear method descriptions (see the below questions) and typos/grammatical issues (just list a few below).

1. Line 042: ”LLMs are generalizable a broad spectrum of downstream tasks”
2. Line 095: “..., which successfully addressing above challenges, achieves” 
3. Line 289, “.. will miss massive information of Consequently, …“

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: * How do you choose the sampling threshold hyper-parameter in the main experiment? Have you tuned it in a validation set?
soundness: 2
strengths: * The introduction of MAB to LLMs' pre-training data selection is reasonable and interesting, addressing a practical problem about diversity and efficiency.
* The experiment results demonstrate that their proposed method significantly improves the pre-trained LLM's performance while maintaining efficiency.
summary: This work proposes to balance the diversity and quality in data selection for LLMs pre-training. The author adopts the MAB method to dynamically estimate and sample data points from various clusters. Additionally, they introduce a novel influential calculation approach to estimate the data point more accurately while improving processing efficiency. Their experimental results on extensive benchmarks validate the effectiveness and efficiency of the proposed method.
weaknesses: * The experiments are not thorough. The experiment and analysis only take up 2 pages, which fails to thoroughly discuss many important components and designs of the proposed method. For example, whether the CS function could be designed in another way? Also, while the evaluation is conducted in various benchmarks, the authors only adopt experiments in a single LLM. This also makes people doubt the generalization of their method. I suggest the author to fully utilize the 10-page space and conduct further experiments.

* According to ablation study results, the model's performance is quite sensitive to the sampling threshold of influence. And there are no hyper-parameter selection strategies introduced in the paper. All these make me worry about the robustness of the proposed method to apply in real-world scenarios.

Meta-review: (a) Summary of Scientific Claims and Findings

The paper proposes Quad, a novel method for optimizing pretraining of large language models (LLMs) through dynamic data selection, leveraging Multi-Armed Bandit (MAB) algorithms. Quad aims to improve model performance and computational efficiency by selecting the most relevant data during pretraining, rather than using raw, unsorted datasets. The authors show, through experiments with 1.3B and 7B parameter models, that Quad enhances model quality and reduces computational costs by focusing on data that aligns with the model's current state. The scalability of Quad to larger models and datasets is discussed but not fully demonstrated.

(b) Strengths of the Paper

1. Innovative Approach: The use of MAB for dynamic data selection is a novel contribution to pretraining optimization. This method provides a computationally efficient way to enhance LLM pretraining by focusing on the most relevant data.
  
2. Clear Experimental Setup: The experiments clearly validate Quad’s effectiveness, showing improved performance in comparison to baseline methods. This solidly supports the claims made in the paper.

3. Practical Impact: The method has practical applications for resource-constrained environments, addressing the growing need for more efficient LLM pretraining, particularly in the context of large-scale models.

4. Clarity and Accessibility: The paper is well-written, with clear explanations of the approach and experimental results, making it accessible to both academic and industry audiences.

(c) Weaknesses of the Paper and Missing Aspects

1. Limited Scalability Testing: Experiments are limited to 1.3B and 7B parameter models. The paper would benefit from more extensive experiments, particularly with larger models and datasets, to substantiate claims of scalability.

2. Lack of Hyperparameter Tuning Details: The paper does not provide sufficient details on how hyperparameters, such as the sampling threshold, were chosen. This omission affects reproducibility and a better understanding of the method’s tuning.

3. Continuous Pretraining Not Addressed: The paper focuses on the initial stages of pretraining but does not discuss how Quad performs in continuous pretraining scenarios, which are common in real-world model training and fine-tuning.

4. Model-Task Mismatch: The paper's experiments focus on a limited set of tasks. A broader range of downstream tasks would help demonstrate the generalizability of Quad across various domains.

5. Integration with Existing Frameworks: The paper does not explore how Quad integrates with existing pretraining frameworks used in practice. More discussion on how Quad could be adopted in real-world systems would enhance its relevance.

(d) Reasons for Decision to Accept

1. Novelty and Contribution: Quad introduces a new, efficient approach to LLM pretraining, contributing both to the theoretical understanding of data selection and the practical needs of model training. The use of MAB algorithms offers an innovative way to optimize large-scale pretraining.

2. Positive Reviewer Feedback: The reviewers have generally praised the novel approach and the soundness of the experimental methodology. While they suggested further exploration in certain areas, the core contribution is considered valuable.

3. Potential for Scalability: Though scalability is not fully tested, the authors make a convincing case for Quad’s potential to scale with larger models, which could make a significant impact on LLM training.

4. Clarity and Accessibility: The paper is well-written and accessible, making it an effective contribution to the community. The methodology and results are presented in a clear and understandable manner, benefiting both researchers and practitioners.

Decision: Accept (Spotlight)