Combatting Dimensional Collapse in LLM Pre-Training Data via Submodular File Selection
Authors: Ziqing Fan
Siyuan Du
Shengchao Hu
Pingjie Wang
Li Shen
Ya Zhang
Dacheng Tao
Yanfeng Wang
Venue: ICLR 2025 Oral (2025)
Keywords: file selection
large language model
pre-training
submodular optimization
Primary area: other topics in machine learning (i.e., none of the above)

Abstract: Selecting high-quality pre-training data for large language models (LLMs) is crucial for enhancing their overall performance under limited computation budget, improving both training and sample efficiency. Recent advancements in file selection primarily rely on using an existing or trained proxy model to assess the similarity of samples to a target domain, such as high quality sources BookCorpus and Wikipedia. However, upon revisiting these methods, the domain-similarity selection criteria demonstrates a diversity dilemma, i.e. dimensional collapse in the feature space, improving performance on the domain-related tasks but causing severe degradation on generic performance.To prevent collapse and enhance diversity, we propose a DiverSified File selection algorithm (DiSF), which selects the most decorrelated text files in the feature space. We approach this with a classical greedy algorithm to achieve more uniform eigenvalues in the feature covariance matrix of the selected texts, analyzing its approximation to the optimal solution under a formulation of $\gamma$-weakly submodular optimization problem. Empirically, we establish a benchmark and conduct extensive experiments on the TinyLlama architecture with models from 120M to 1.1B parameters. Evaluating across nine tasks from the Harness framework, DiSF demonstrates a significant improvement on overall performance. Specifically, DiSF saves 98.5\% of 590M training files in SlimPajama, outperforming the full-data pre-training within a 50B training budget, and achieving about 1.5x training efficiency and 5x data efficiency. Source code
is available at: https://github.com/MediaBrain-SJTU/DiSF.git.

Published as a conference paper at ICLR 2025
COMBATTING DIMENSIONAL COLLAPSE IN LLM PRE- TRAINING DATA VIA DIVERSIFIED FILE SELECTION
Ziqing Fan 122, Siyuan "e.,",,henghaaa Hu Pingjie Wang!? Li Shen Ya Zhang! .hacheg Tao", Yanfeng Wangh2 Shanghai Jiao Tong University, China; SShanghai AI Laboratory, China: FFudan University, China *Shenzhen Campus of Sun Yat-sen University, China; 'Nanyang Technological University, Singapore edu. en, dusi yuan@pjlab.b...o org. en, charles-hu@e jtu, .edu.ct en, prrgiiewagges jtu. edu. on, waattshheni@@@maaa com, ya ya_zhangés jtu edu. cn, dacheng. taolntu. edu, ag, wangyanfengee edw. on
ABSTRACT
Selecting high-quality pre-training data for large language models (LLMs) is cru- cial for enhancing their overall performance under limited computation budget, improving both training and sample efficiency. Recent advancements in file se- lection primarily rely on using an existing or trained proxy model to assess the similarity of samples to target domain, such as high quality sources BookCorpus and Wikipedia. However, upon revisiting these methods, the domain-similarity selection criteria demonstrates diversity dilemma, i.e. dimensional collapse in the feature space, improving performance on the domain-related tasks but causing severe degradation on generic performance. To prevent collapse and enhance diver- sity, we propose DiverSified File selection algorithm (DiSF), which selects the most decorrelated text files in the feature space. We approach this with classical greedy algorithm to achieve more uniform eigenvalues in the feature covariance matrix of the selected texts, analyzing its approximation to the optimal solution under a formulation of 7-weakly submodular optimization problem. Empirically, we establish a benchmark and conduct extensive experiments on the TinyLlama architecture with models from 120M to 1.IB parameters. Evaluating across nine tasks from the Harness framework, DiSF demonstrates a significant improvement overall performance. Specifically, DiSF saves 98.5% of 590M training files in SlimPajama, outperforming the full-data pre-training' within a 50B training budget, and achieving about 1.5x training efficiency and 5x data efficiency. Source code is available at: https://github.. www.eoooooee/aeieee -/TTUintTSSTTDSSEEEE
INTRODUCTION
Pre-trained Large Language Models (LLMs) have demonstrated remarkable capabilities (Brown, 2020a; Chowdhery et al., 2023a; Touvron et al., 2023b), but their training is computationally ex- pensive, with costs increasing et al., 2021; Thoppilan et al., 2022). For instance, training GPT-3 with 175 billion parameters is estimated produce 552 tons of CO2 emissions and consume 1,287 MWh of energy (Patterson et al., 2021). In practice of commercial use and common academic research, training budgets such as the number of pre-trained tokens are typically predefined, determined by available devices and training time constraints (Hoffmann et al., 2022). To optimize the performance of LLMs within the budget, selecting high-quality pre-training data from large text corpora essential, boosting both training and sample efficiency.
model size and training data grow (Rac et al., 2021; Patterson
Recent innovations of selecting files for pre-training LLMs mostly rely on using an existing or trained proxy model and designing proxy function to access the similarity to target domain, which is regarded as high-quality data. Heuristic classification (Brown, 2020b; Chowdhery et al., 2023b) trains binary classifier and select similar content to text domains like Books and Wikipedia (Computer,
Full-data pre-training in this paper refers to pre-training the L.MMs all training files in SlimPajama until the specified training budget is reached.

---

Published as a conference paper at ICLR 2025
(a) Heuristic
(b) DSIR
(c) QuRating-W W
(d) D4
(e) DISF
Figure I: The t-SNE (Van der Maaten & Hinton, 2008) visualization of text features (normalized to the unit sphere) selected by different methods on SlimPajama. We use Contriever (Izacard et al., 2021) to extract features. (a) and (b) show Heuristic classification and DSIR based on the Wikipedia and Book domains, while (c) depicts QuRating based on writing judgments. We visualize top 500 text features selected by their criterion, which forms long narrow band, indicating dimensional collapse. (d) and (e) represent D4 and our DISE. For D4, we display 500 random samples after reducing redundancy, while for DiSF, we select samples with the highest values based on equation 6. Both methods, especially DiSF, show more uniformly scattered features, indicating improved diversity.
2023). DSIR (Xie et al., 2023b) also targets these domains, using a hashed n-gram extractor to measure similarity. QuRating (Wettig al., 2024) leverages GPT--35-ttubo to train judge model that evaluates the quality of domains like writing and education. However, as shown in Figure/ 2, these methods based on specific domains lead a diversity dilemma, known dimensional collapse in representation learning (Jing et al., 2022; Zbontar et al., 2021; Bardes et al., 2022; Shi et al., 2022; Fan al., 2024). Their feature vectors of samples span a lower-dimensional subspace, indicating less diversity. improving performance in domain related tasks, such as reading comprehension (e.g.. ARC (Clark al.. 2018) and OBQA (Mihaylov al., 2018)). but causing severe degradation in overall performance across diverse domains, particularly in physical world knowledge tasks like PIQA (Bisk et al., 2020) and HellaSwag (Zellers et al., 2019).
In our paper, we revisit these algorithms by visualizing the feature representations of their selected text samples as shown in Figure 1, and performing eigenvalue analy- sis on the features' covariance matrix (see Section for details). As depicted in Figure 1 (a). (b), and (c), ob- serve dimensional collapse, where text samples selected based specific domain show dominant top eigenval- ues, indicating long narrow feature spaces. In contrast, text samples selected using our diversified method exhibit less dominant eigenvalues, leading to greater uniformity across feature dimensions (Figure (e)). Recent methods, such as D4 (Tirumala et al., 2023) and INGENIOUS (Renduch- intala et al., 1991), recognize the importance of diversity and select informative samples by leveraging feature dis- tances and similarity kernels, respectively. However, they fall short in achieving the level of uniform representations attained by our approach as shown in Figure I (d) and (c) and Figure 3.
Random
DSIR
QuRaaiing W
DiSP(Ours)
WinoGrande
HellaSwag
52.8
33.5
BoolQ
60.8
27
59
47.4 PIQA
16.8
1955
OBQA
20.9
45
47.3
23:4
ARC-e
ARC-c
Figure 2: Commonsense reasoning abil- ity of pre-trained TinyLlama 1B us- ing various selection methods evalu- ated on seven tasks of Harness. DSIR uses Wikipedia and BookCorpus as high quality source and QuRating-W selects samples with writing style score. All methods select 1.5% of training files in SlimPajama and pre-train 50B tokens.
To prevent collapse and enhance diversity, we propose DiSF, that selects files by minimizing the Frobenius norm of the features' covariance matrix, achieving more uniform eigenvalues. We approach this with the classical greedy algorithm and analyze its approximation to the optimal solution under the formulation of 7-weakly submodular optimization problem (DAS, 2011). Empirically, we establish benchmark on the newly released and popular TinyLlama (Zhang et al., 2024a) architecture with models of 120M, 560M, and 1.1B parameters. Extensive experiments and ablation studies, conducted across nine tasks on the Harness framework, demonstrate the superior general performance of our method compared to baselines. Specifically, out of 590M training files in SlimPajama (Touvron al., 2023a; Computer, 2023), our DiSF selects just 1.5% (about 9B training tokens), outperforming full-data pre-training under 50B training budget, and achieving approximately 1.5x training efficiency and 5x data efficiency. In summary, our research makes three significant contributions to the field:
2

---

Published as a conference paper at ICLR 2025
Table 1:
Summary of recent innovations in file selection for LLM pre-training, categorized by their proxy model, proxy function to estimate sample importance, and whether requiring a target domain.
Selection Method
Proxy Model (M)
Proxy Function (Fxx
Domain Independent
Heuristic Cls.
Trained binary classifier
Probability to target domain
DSIR 2N)
Hashed n-gram extractor
Similarity target distribution
QuRating (2CM)
Trained judge model via GPT-3.5
Judgement score target ability
INGENIOUS (Em)
Pre-trained LLM with warming-up
Facility Location slimilarity matrix
D4 (NeurIPS'25) 23)
Existing feature extractor
Distance feature space
DiSF (Ours)
Existing text feature extractor
Decorrelation of feature dimensions
We rethink recent file selection innovations in pre-training LLMs, and identify diversity dilemma known dimensional collapse in feature representation learning, improving performance in domain-specific tasks but causing severe degradation overall performance (Section 2).
To prevent collapse and enhance diversity, we propose DiSF, which selects the most decorrelated text files using a classical greedy algorithm, and analyze its approximation to the optimal solution under the formulation of 7-weakly submodular optimization problem (Section 3).
We established benchmark on TinyLlama architectures and SlimPajama text corpus with evaluation on nine tasks from Harness. Extensive experiments and ablations demonstrate the superior performance of our method, with improved training and sample efficiency (Section 4).
RETHINKING FILE SELECTION FOR LLM PRE-TRAINING UNDER BUDGET
In this section, we first introduce the definition of file selection objective, and then revisit recent file selection innovations for pre-training LLMs, with a particular focus on the diversity dilemma.
2.1 PROBLEM STATEMENT
Given training and selection budgets T and S, our objective is to select the most valuable samples V from a text corpus D = [D1, Di.... Dx) } collected from various domains D; (e.g., Wikipedia or samples with high writing qualities) to optimize model weights W on pre-training task L, thereby maximizing general performance A. While A is challenging to verify given LLM, it can be inferred through diverse abilitics such as commonsense and problem-solving, with evaluation tools like Harness (Gao et al., 2024). Mathematically, V can be obtained through selection objective as
arg max UCD
A(arg min CWW,T,WW))
s.t. S.
There are various options to define T, A, and this work, we define the training budget T as the number of pre-trained tokens, the selection budget S as the number of files, the pre-training objective C as next-word prediction on SlimPajama (Computer, 2023; Touvron et al., 2023a), and the LLM's generic performance A as the evaluation across different tasks using the Harness framework.
2.2 RECENT SELECTION METHODS
Analyzing the objective defined in equation 1, the inner part minimizes the pre-training objective on LLM with predefined training budget and the selected samples, while the upper level focuses on selecting the most valuable samples to maximize the LLM's generic performance within the selection budgets. Directly searching for valuable samples V in the full corpus extremely time-conuuiing and expensive. To reduce this cost, recent innovations in file selection for LLM pre-traee pre-training mostly rely on using an existing or trained proxy model M and designing proxy function FM based on a target domain. Through linking text samples E T to the generic performance A, they transform equation into choosing the samples with the highest values of (x))), as follows:
V Top M D).
(2)
As summarized in Table 1, typical Heuristic classification (Brown, 2020b; Chowdhery et al., 2023b) trains a binary classifier to filter web data, selecting files with probabilities to target format above
3

---

Published as a conference paper at ICLR 2025
noisy, such BookCorpus and Wikipedia (Computer, 2023). Similarly, DSIR (Xie al., 2023b) improves and also treats these two domains as high-quality domains to select files for general purpose, with hashed n-gram feature extractor measure the similarity between the text features and the target distribution. QuRating (Wettig et al., 2024) queries GPT--35-turbo and trains judge model assess the text samples' quality of target style, such writing and mathematics. However, as shown in Figure 2, selection methods based on specific domains lead to a diversity dilemma, known dimensional collapse in representation learning (Jing et al., 2022: Zbontar et al., 2021; Bardes et al., 2022; Chen et al., 2020), improving performance in the domain related task but causing severe degradation in overall performance across diverse tasks. The recent method D4 (Tirumala et al., 2023) and INGENIOUS (Renduchintala et al., 1991) reduce file redundancy by leveraging feature distances and selecting informative samples based on similarity kernels respectively, which improves diversity but can not effectively mitigate dimensional collapse as ours (see both Figure 1 and Figure 3).
2.3 DIVERSITY DILEMMA: DIMENSIONAL COLLAPSE
As shown in Figure 1, dimensional collapse occurs in the embedding space when samples are selected based on the target domains. Their embedding vectors extracted by the Contriever (Izacard et al., 2021) span a lower-dimensionll subspace, indicating less diversity. To quantify this, we conduct an eigenvalue analysis on the covariance matrix of selected text features under different selection methods, and visualize the dominance score of the top,, eigenvalues
0.6
Dominance score
0.5
0.4
0.3
DiSFrourai NGENIIUS
QuRatin D4
calculated as E1= Aj where À is the i-th large eigen-
0.2
gulating-W
value of the covariance matrix, and d is the dimension of the feature space. Smaller dominance score suggests more uniform feature dimensions and richer information (Chen al., 2020; Zbontar et al., 2021; Shi al., 2022; Jing et al., 2022; Bardes et al., 2022). As demonstrated in Figure 3, Heuristic classification, DSIR, and QuRating with a focus writing style (denoted as QuRating-W -) exhibit significantly higher dominance scores compared to D4, INGENIOUS, and QuRating with all styles (de- noted as QuRating-A), highlighting the reduced diversity caused by target domains selection criteria. To address this and enhance diversity, we propose a novel diversified file selection algorithm (DiSF) to select text samples, span- ning more uniform feature dimensions in the embedding space. As red line shown in Figure 3, compared and QuRating-A A, our DiSF further reduces the dominance score, demonstrating its effectiveness in uniforming dimensions and mitigating dimensional collapse.
0.1
Hearicic
20
40
60
80
100
k
Figure 3: The dominance score for
recent methods calculated k A Σ jam Ay where A; represents the i-th largest eigen- value the feature covariance matrix, and d is the dimension of the feature space. We use Contriever model to ex- tract features. We select the top 500 text samples based on their respective selec- tion criteria. For D4, we select 500 ran- dom samples after reducing redundancy.
D4
3 DISF: DIVERSIFIED FILE SELECTION
In this section, we define the diversified selection criterion of DiSF, and the selection procedures with classical greedy algorithm in batch level, and then analyze the function from view under 7-weakly submodular optimization theories, followed with time complexity analysis.
3.1
METHOD
Selection criterion. Given set of text samples U = (x1, ..., N.)) their text features with
standard normalization Z = {Z1,..., Zi, .. 2n3 are obtained by a text feature extractor M as
f(x;, M) p
2;
(3)
4

---

Published as a conference paper at ICLR 2025
where f calculates the feature representations of text samples,and and the mean and variance of M), f (xx,M), of M)), respectively. Then, the covariance matrix C is defined as
C(U,
1
2 2 24-
As discussed in Section 2, our goal is to prevent dimensional collapse by ensuring more uniform eigenvalues in the covariance matrix the selected samples. Directly calculating and selecting based on eigenvalues are costly, but it is feasible to optimize the Frobenius norm of the covariance matrix || (Zbontar et al., 2021; Bardes et al., 2022; Shi et al., 2022), as described in Lemma I:
Lemma 1. Assuming a covariance matrix C E Raxd computed from the features with the standard normalization, and its eigenvalues {A1; A2, .... >a), we will have the following equality that satisfied
d
Σ d j=1
1
2,)²
||C1] d.
We provide a detailed proof in Appendix A.6 for clarity. From Lemma 1, it is evident that ensuring the uniformity of the eigenvalues of the covariance matrix can be translated into minimizing the Frobenius norm of the covariance matrix. Thus, given any text set U C D, we evaluate its importance as M)]C( and define the selection objective as
arg max -ICC(U
S.L. S
(5)
where S is the predefined selection budget. Differing from typical selection algorithms that directly define the importance of individual text samples through proxy function shown in equation 1, our selection objective defined in equation is calculated on a subset of samples.
Selection Procedure. To satisfy the non-negative requirement in the later analysis, we first reformulate our proxy function into a non-negative form as
Algorithm Selection procedure of DiSF
V/ + (,,
(6) where U C D. We apply the classical greedy al- gorithm to select the most valuable samples based on our proxy function. This allows to iteratively choose the most valuable samples as follows:
FOISF (U) e IIC(U.M))l.
Divide D into batches of b, with scale b.
for 1,..., [ do
randomly select E and U, + {x*}
while |U:| S M do
27
b. b. {x'}
U U (arg max DiSF (UU (x})}-
(7)
x = arg max REb
D (U,U
XEDŞU
U, - fx*)
However, directly applying the greedy algorithm as shown in equation the entire text corpus com- putationally expensive, we perform the selection at the batch scale. In Section 4.4, we provide an abla- Output:V tion study of selection scale. Given text corpus D, selection scale b, selection budget S, and proxy model M, the selection process is outlined Algorithm 1. We first divide the text corpus into [D] batches, where |-] is the round down command. In each batch d,, we initialize U; with randomly selected sample and remove from b,. Then, we iteratively apply (L [[5)S] 1) times the greedy algorithm, removing the most valuable sample from by and adding it to U, based on equation Finally, the selected samples are the combination of all U,.
end while
V V-U
end for
3.2 ANALYSIS
Selection Analysis. Diversity is often effectively modeled by submodular functions (Nemhauser al., 1978; Fujishige, 2005; Balakrishnan al., 2022; Hong al., 2024; Zhang et al., 2025; 2024b; 2023), which exhibit a property of diminishing returns, i.e., the marginal gain an element brings to subset decreases more elements are added that subset. Mathematically, given set Ω, set function I: 2 + R is --weakly submodular and only if, for any subsets A C B C Ω and a element E Ω \ B. the following inequality holds:
[AAU (x)) ((A) ([BBU (x]) I(B),,
(8)
where y E (0,1).
5

---

Published as a conference paper at ICLR 2025
A non-negative monotone y-weakly submodular maxi- mization problem can be solved using the classical greedy algorithm, which guarantees (1 (1-e -0---=)-ppppp ">-approximation to the optimal solution (function value on selected samples compared to optimal value) (DAS, 2011). As demon- strated in Figure 4, we empirically verify the monotonicity of our proxy function by increasing the number of samples randomly chosen from the selected text files of different methods. Additionally, the results in Figure indicate a significantly higher proxy value for our selection method compared to others, demonstrating the superior ability to mitigate dimensional collapse. Since our proxy function also aims to capture the diversity within a given set, we algorithm ap-
Value of proxy function(e gcaur
0.8
0.7
0.6
DiSF[Ours]
D4
QuRating-A
Random
QuRating W
0.3
200
400 600
800
1000
Number sample
try to provide some insights into how proaches the optimal solution under this formulation. To provide theoretical guarantee of the selection, aim to establish lower bound for the submodular ratio under any given set where the gain is positive, as most empirical verifications in Figure 4 suggest. Since the original formu- lation challenging to analyze directly, we introduce two bounds on the function value. One is the upper bound on gain difference, as demonstrated in Assumption and the other one is the upper bound on the average utility of function value as shown in Assumption 2. Therefore, we prove lower bound of the submodular ratio of e 2p
Figure 4: Proxy value, as defined in equation 6, calculated on different meth- ods. For each method, we randomly choose samples from their selected text files with 1.5% selection budget. All cases generally demonstrate the property of monotonicity. Moreover, our selec- tion method achieves significantly larger proxy values, indicating much better uni- formity of feature dimensions.
- which means given assumptions 1,
6
2, and positive gain, our function is at least e 2p g-u-s_1 -weakly submodular function. The detailed proof is provided in the Appendix A.5.3
eAp
Time Complexity Analysis. Given the feature dimension d, the time complexity of DiSF is less than (dd)) and we provide detailed calculation of the results. All terms in the time complexity most quadratic and independent of the overall dataset size, which we consider acceptable for applying to file selection for LLM pre-training. Additionally, in Figure of Section 4.4, we report the exact time required to select data on the benchmark dataset. In Section 4.3 and Appendix A.5.2, we compare the sampling and training efficiency of our method against the baselines.
4 EXPERIMENT
This part introduces the experimental setup, including dataset, evaluations, model architecture, baselines, and training details in Section 1, the main results and efficiency analysis in Section 4.2 and Section 4.3, and extensive ablation studies for better understanding of our DiSF in Section 4.4.
4.1 SETUP
Dataset and evaluation. Following many prior works (Touvron et al., 2023a; Zhang al., 2024a; Wettig et al., 2024; Xie et al., 2023a), we employ SlimPajama (Touvron et al., 2023a; Computer, 2023) the text corpus, which is specifically curated for pre-training LLMs. All selections performed on about 590M training files of SlimPajama, processed with Llama tokenizer (Touvron al., 2023a). Notably, QuRating (Wettig et al.. 2024) provides judgments on various properties of text samples in SlimPajama using a judge model trained based on GPT-3.5-turbo, which can be directly utilized in our implement. To capture the diversity dilemma and evaluate generic performance of pre- trained LLMs, we use seven commonsense reasoning tasks from the popular framework Harness (Gao et al., 2024), including four reading comprehension tasks (ARC-e, ARC-e (Clark et al., 2018), OBQA (Mihaylov et al., 2018), and BoolQ (Clark et al., 2019)), and three physical world knowledge tasks (PIQA (Bisk et al., 2020), HellaSwag (Zellers et al., 2019), and WinoGrande (Sakaguchi also employ another two tasks MMLU (Hendrycks et al., 202 1), and BBH (Suzgun et al., 2022) from Harness to evaluate the problem-solving capabilities for further clarify of our effectiveness. See Section A. for detailed introduction of the dataset and tasks.
et al., 2021)). Besides,

---

Published as conference paper at ICLR 2025
Table 2:
Performance seven tasks and their average with the Harness framework. In the upper part, we pre-train TinyLlama with 120M. 560M, and I.IB parameters from scratch, with training budget of 10B tokens and selection budget of 1.5% of total training files. In the lower part, the training budget is increased to 50B. Results impacted by the diversity dilemma, best individual task results, and the best average performance respectively highlighted in bold blue, black, and red.
Method
Pre-training Llama from scratch with 10B training budget 1.5% selected files
Model Size
#Metric
ARC-e
ARC-c
OBQA
BoolQ
WinoGrande
HellaSwag
PIQA
Avg.
Random
37.1
18.2
12.8
61.1
49.9
27.0
58.7
37.8
DSIR
36.6
18.2
13.8
57.6
51.7
25.1
54.8
36.9
D4
38.0
18.1
13.6
60.1
51.8
27.8
57.8
38.2
QuRating-W
37.7
18.6
15.0
60.9
51.3
24.3
55.6
37.7
120M
QuRating-A
39.1
18.3
14.2
60.7
50.8
27.4
Doremi
36.8
18.2
13.8
60.9
52.7
27.3
58.8
38.5
INGENIOUS
39.5
18.6
13.4
60.2
50.4
27.4
58.7
38.3
DiSF (Ours)
39.9
61.7
51.9
27.6
Random
43.2
19.9
15.2
59.3
$2.9
30.4
62.0
40.4
DSIR
41.8
19.3
16.8
60.5
50.3
26.0
D4
46.5
19.5
16.0
60.5
53,2
29.5
61.4
40.9
QuRating-W
42.0
21.3
17.4
59.5
49.8
28.4
58.2
39.5
560M
QuRating-A
46.7
19.1
17.2
61.3
51.1
28.9
41.1
INGENIOUS
45.8
20.6
16.0
58.3
51.4
30.8
62.5
40.8
Doremi
42.7
19.3
15.8
60.9
50.4
30.0
63.7
40.4
DiSF (Ours)
47.5
21.2
16.2
58.9
51.0
31.1
64.2
Random
44.8
19.0
16.4
59.9
51.3
30.8
64.1
40.9
DSIR
45.7
20.3
18.6
59.8
50.4
27.6
D4
46.2
19.3
18.8
60.2
51.3
30.9
65.4
41.7
QuRating-W W
44.4
21.4
17.0
59.6
51.4
31.0
60.1
40.7
1.1B
QuRating-A
47.4
20.9
59.1
50.2
30.1
INGENIOUS
45.3
19.6
19.8
60.0
51.2
31.2
64.9
41.7
Doremi
44.7
19.7
17.6
61.0
51.2
31.1
64.6
41.4
DiSF (Ours)
47.7
19.5
18.2
59.7
52.2
32.3
65.6
42.2
Method
Pre-training Timy from scratch with 50B training budget on 1.5% selected files
Model Size
#fMetric
ARC-e
ARC-c
OBQA
BoolQ
WinoGrande
HellaSwag
PIQA
Avg.
Random
41.5
17.6
16.6
56.3
52.2
28.4
59.5
38.9
DSIR
44.2
18.4
17.2
53.9
49.6
25.2
56.3
37.8
D4
42.1
19.2
17.0
58.3
52.6
28.1
60.9
39.7
QuRating-W W
41.6
18.9
16.2
59.5
51.6
27.6
56.8
38.9
120M
QuRating-A
46.8
19.7
15.8
60.5
51.1
27.9
58.4
40.0
INGENTOUS
42.2
19.1
16.3
58.7
51.2
28.5
61.0
39.6
Doremi
40.7
18.9
16.2
60.2
52.7
28.1
61.4
39.7
DiSF (Ours)
44.3
18.3
17.8
61.1
53.3
28.7
61.0
40.6
Random
46.0
20.9
16.8
59,0
52.5
31.6
64.7
41.6
DSIR
45.9
22.2
18.5
58.1
50.7
27.8
59.3
40.4
D4
47.2
21.7
18.2
58.7
52.2
32.4
65.3
42.2
QuRating-W W
46.9
23.4
19.5
56.5
52.2
28.5
61.3
41.2
560M
QuRating-A
48.3
19.2
18.2
58.9
52.1
34.1
67.4
42.6
INGENIOUS
46.8
21.6
17.0
59.3
52.2
32,8
66.5
42.3
Doremi
47.2
22.4
59.0
52.6
33.1
DiSF (Ours)
47.3
22.0
18.8
60.8
52,8
33.5
67.4
43.2
Random
51.4
20.9
18.2
56.2
51.3
34.7
67.3
42.9
DSIR
50.2
20.6
20.0
54.6
52,4
30.4
64.2
41.8
D4
53.6
22.1
20.9
57.3
52.9
35.2
67.7
44.2
QuRating-W
53.9
23.1
20.8
55.0
52.7
35.0
63.3
43.4
1.1B
QuRating-A
53.6
23.7
21.2
60.1
51.6
36.5
67.0
44,8
INGENIOUS
51.7
21.6
22.2
56.9
51.7
36.8
67.9
44.1
Doremi
51.2
21.9
21.8
58.1
54.3
36.6
44.5
DiSF (Ours)
51.8
22.7
20.0
62.0
53.5
37.3
69.0
45.2
Model architecture. As for the model, we adopt Tinyllama architecture (Zhang et al., 2024a) with 120M, 560M, and 1.1B parameters. Thanks to FlashAttention (Dao et al., 2022) and Lit- GPT (LightningAA, 2023), all experiments can be conducted on NVIDIA GeForce RTX 4090 GPUs with 24GB memory, which is feasible for general academic research. All experiments and selection implemented by PyTorch (Paszke al., 2019) platforms with 8 GPUs and 64 CPUs. Except for Tinyllama, we also adopt OPT (Zhang et al., 2022) and Pythia (Biderman et al., 2023) to verify the scalability of our method model architecture as ablation study shown in Section 4.4. For feature extraction, we utilize the Contriever (Izacard et al., 2021) with approximately 110M parameters to calculate feature representations of the text samples defined in equation 3. We also experiment
architeeo

---

Published as a conference paper at ICLR 2025
with other pre-trained models as feature extractors, including the text encoder of CLIP (Radford et al., 2021) and GPT-2 (Radford et al., 2019), as discussed in the ablation study of Section 4.4. See Appendix A.4 for detailed introduction of model structures, well as their training times.
Baselines. We compare our DiSF with Random selection and existing file selection methods for LLM pre-training, including DSIR (Xie et al., 2023b), QuRating (Wettig et al., 2024), INGE- NIOUS (Renduchintala et al., 1991), and D4 (Tirumala et al., 2023). Since DSIR improves on Heuristic classification (Brown, 2020b; Chowdhery et al., 2023b), we present only DSIR results based on Wikipedia and Books. For QuRating, based on the top judgment values, we select text samples for writing style (denoted as QuRating-W) and uniformly select across all styles (denoted as QuRating-A). For INGENIOUS, we utilize Contriever to extract features, which is more efficient than warmed-up model, as analyzed in Appendix A.3.4, Additionally, for a comprehensive compari- son, we include Doremi (Xie et al., 2023a), a recently proposed method that produces weights for pre-training on multiple text domains. We use the weights as the selection ratio of text samples in different domains in our experiment. Notably, in our ablation studies, we also compare our method with Full Data pre-training, which refers to pre-training the LLMs on all training files in SlimPajama until the specified training budget is reached. See Appendix A.3 for detailed implementation of the baselines.
Pre-training details. We follow all settings in TinyLlama (Zhang et al., 2024a). The optimizer AdamW (Loshchilov& & Hutter, 2019), setting parameters 31 at 0.9 and {3 at 0.95. We adopt the cosine learning rate schedule with maximum learning rate of 4e-4 and the minimum of 4e-5, the batch size of 2M tokens, the weight decay of 0.1, and the gradient clipping threshold of 1. The training budgets are 10B and 50B tokens, with 1.5% selection budget of SlimPajama's 's training files. Note that, we choose to report performance with a 1 1.5% selection budget, since it achieves comparable performance compared to Full Data pre-training under 50B pre- pre-training budget TinyLlama I.IB. Unless otherwise specified, the selection scale b is set to 1024. See Appendix A.4 for more details.
4.2
MAIN RESULTS
Performance on commonsense reasoning tasks. As shown in Table 2, selection methods based on a target domain such as DSIR, D4, and QuRating-WW W improve performance on reading comprehension tasks like ARC-c and OBQA, but suffer significant declines on physical world knowledge tasks especially HellaSwag and PIQA (highlighted in blue), revealing the diversity dilemma. contrast, Doremi, D4, QuRating-A, and our DiSF, which select samples from multiple text domains, achieve competitive results on OBQA and ARC-c while significantly improving overall performance across the remaining tasks. This highlights the critical importance of diversity in pre-training LLMs. Notably, our DiSF outperforms all baselines in terms of average performance across the seven tasks, with an average improvement of 2.5% compared to DSIR. Furthermore, with increasing training budget, the improvement on DSIR becomes more pronounced, rising from 2.1 % to 3.4% on TinyLlama I.
Scalability model size. To verify the scalability of selection methods across different model sizes, we use the TinyL architecture with models of 120M, 560M, and 1.IB parameters. As shown in Table 2, our DiSF con- sistently outperforms all baselines, demonstrating strong scalability with increasing model size. Notably, as the model scales from 120M to 1.1B, DiSF's improvement over DSIR grows from 2.8% to 3.4%, suggesting even greater effectiveness for larger L.LMs.
Table Problem-solving performance on MMLU (5 shot) and BBH (3 shot) of TimyLlama 1.1B with 1.5% selection budget and 50B pre- pre-training budget.
Method
MMLU(S shot) BBH(3 shot)
DSIR
22.9
18.5
QuRating-VV W QuRating-A
23.1
18.710.4
24.5-0.4
20.2+0.4
D4
24.3
20.8
Performance on problem-solving tasks. As shown in Table 3, we further verify the performance of pre-trained LinyLlama I.IB on two additional problem-solving tasks,
Doremi
23.1
20.3
DISF (Ours)
25.4
21.2+0.4
MMLU and BBH, using various selection methods. Except for better commonsense abilities, results
of our DiSF shown in Table 3, demonstrate better problem-solving capability, outperforming all other baselines. Specifically. on MMLU and BBH, our DiSF respectively achieves 2.5% and 2.7% improvement compared to DSIR
8

---

Published as a conference paper at ICLR 2025
Performance (96)
41.0
TinyLlama 120M. 508
Performance (%)
45.0
TinyLlama 560M, 508
Performance (%)
46.5
TinyLlama 1.1B, 508
Foll Data
Full
Fill Data
39.5
41.5
44.0
0.1
0.2
|
|
|S/IT
Figure 5: Average performance of TinyLlama pre-trained after 50B tokens on files selected by our
DiSF with varying selection budgets compared to Full Data pre- pre-training SlimPajama.
Performance (96)
TinyLlama 120M
Performancc (%6)
TinyLlama 560M
Performance (%)
46
Tiny 1.1B
DISF (Ours)
DISF (Ours)
DisF (Ours)
340
37
36
25 Pre-trained tokens (B)
50
Pre-trained tokens (B)
Pre-trained tokens (B)
Figure 6: Performance of TinyLlama pre- pre-trained under DiSF with 1.5% selection budget compared
to Full Data pre-training on Slim Pajama, as the pre-trained tokens (training budget) increases.
4.3 EFFICIENCY ANALYSIS
In this section, we try to analyze how many samples and computations we can save within the training budget, compared to Full Data pre-traininn pre-training, that is data efficiency and training efficiency. Please note that, Full Data pre- me-ttrining refers pre-training the LLMs on all SlimPajama's training files until the specified training budget is reached. In Figure 5, we demonstrate the average commonsense performance of TinyLlama with 120M, 560M, and 1.1B param- eters, pre-trained after 50B tokens on files selected by DiSF with varying selection budgets compared to Full Data. It is evident that to achieve the equivalent or superior performance compared to Full Data with 50B training budget, we only need to select about 20% of the pre-trained tokens (1.5% of SlimPajama' s training files), achieving at least 5x data efficiency. In Figure 6, we show the curves of averaged common- sense reasoning performance during the pre-training under DiSF with 1.5% selection budget compared to Full Data. It can be seen that, when achieving the performance Full Data pre-trained with 50B tokens, our methods only need 27B, 36B, and 40B tokens, respectively on TinyLlama with 120M, 560M, and 1.IB parameters, achieving an average of 1.5x training efficiency.
46
hours
hours
26 hours
41 hours
hours
A
Performance (%)
TicyLlama 1.1
Lama 560M
Thay Laama 120M
10 11 12
Selection Scale (log2b)
Figure 7: Performance of pre-trained TinyL- lama under different selection scales (log, b). With 1.5% selection budget and 50B training budget, we identify ideal point with ac- ceptable computational cost and near-optimal performance, marked by a red star.
4.4 ABLATION STUDY
Selection budget. This part analyzes the perfor- mance of pre-trained LLM with different selection budgets when using our DiSF. We conduct the abla- tion on TinyLlama 120M with 50B training budget and present its average commonsense reasoning per- formance. As shown in Figure 8, the performance initially rises rapidly with increasing selection ratios, reaching a peak 41.2 with 3% of total training files (about 20B tokens). Once the selected samples exceed 3%, the performance begins to decline and converge to the performance of Full Data pre-trate training. This insight may help choose selection ratios and en- hance understanding of selection algorithm.
Table 4: Ablation study of our DiSF with dif- ferent model architectures compared to other selection methods, respectively using TinyL- lama with 1.1B parameters, Pythia with 1B parameters and OPT with 1.3B parameters.
Method
TinyLlama Pythia
OPT
DSIR
41.8
41.4
43.1
D4
44.2
43,9
44.7
QuRating-A
44.8
43.6
45.1
Doremi
44.5
43.7
44.9
Ours
45.2
44.2
45.5

---

Published as a conference paper at ICLR 2025
Model architecture. To further verify performance across different architectures, additionally adopt two large language models: Pythia (Biderman et al., 2023) with 1B parameters and OPT (Zhang et al., 2022) with 1.3B parameters. As shown in Table 4, compare the average commonsense reasoning performance of our DiSF with baselines, including DSIR, D4, QuRating-A, and Doremi, under the same training budget of 50B pre-trained tokens. The results demonstrate that DiSF consistently outperforms the baselines, highlighting its effectiveness and scalability across different model architectures.
Feature extractor. In our experiments, we uti- lize proxy model named Contriever (Izacard et al., 2021), with about 11M parameters, as feature extrac- tor to provide the embedding space for selected text samples. We also try two other pre-trained models: the text encoder of CLIP (Radford et al., 2021) with about 70M parameters and GPT-2 (Radford et al., 2019) with about 1 117M parameters. As shown in Ta- ble 5, CLIP fails to deliver satisfactory performance, because it struggles with longer text sequences. GPT- 2 does not outperform the Contriever model because it is optimized for autoregressive text generation rather than producing meaningful sentence represen- tations, whereas Contriever, specifically designed for measuring text similarity, excels at capturing text diversity.
Table 5: Ablation study of our DiSF with different feature extractor. We show results, respectively using Contriever with 1 110M pa- rameters, text encoder of CLIP with 70M pa- rameters and GPT-2 with 117M parameters.
Tiaya CLIP
Contriever
GPT-2
120 M
39.5
40,6
40.0
560 M
42.1
43.1
42.3
1.1
43.9
45.2
44.4
Selection scale. To manage the computational cost of file selection, we apply submodular optimization with classical greedy algorithm at the batch scale, as shown in Algorithm 1. introducing the hyper- parameter, selection scale b. The computational com- plexity of the selection process, defined in equation 7,
TinyLlama 120M, T= 50B
41.5
(3% 41.2)
Performance (%)
(1.5%.40.
(10%,40.5)
(25%.40 10.1)
is ( | ). where the larger increases the compu-
(1%,39.8)
tational burden. As ablation on b shown in Figure 7, a larger batch size significantly raises computational costs, while smaller batch size fails to select suffi- ciently diversified files. With fixed selection budget of 1.5%, training budget of 50B tokens, and Con- triever as the feature extractor, we identified an ideal point of 1024, marked by a red star, which balances computational time and near-optimal performance. This insight may provide valuable guidance for choos- ing an appropriate selection scale in other settings.
(0.5%.38.7)
1.5% Random Selection
38.5
|s|
Figure 8: Average commonsense reasoning performance of TinyL.lama 120M pre-trained under DiSF with different selection budget and 50B training budget. Blue and grey lines respectively denote Full Data pre-training and random selection with 1.5% selection budget.
CONCLUSION
In this work, we revisit recent innovations in file selection for pre-training large language models (LLMs) and identify a diversity dilemma: dimensional collapse, where performance improves on specific tasks but degrades overall across diverse tasks. To address this, we propose novel Diversified File selection method (DiSF) which selects decorrelated text files in the embedding space to enhance diversity. DiSF achieves more uniform eigenvalues of the feature covariance matrix by minimizing its Frobenius norm and solve it with greedy algorithm. We analyze its time complexity and approximation to optimal solution under 7-weakly submodular optimization, and establish a benchmark with TinyLlama architecture, evaluating performance across nine tasks from the Harness framework. Extensive experiments and ablation studies demonstrate the critical role of diversity in file selection for LLM pre- training and showcase DiSF's superior effectiveness and efficiency.
ACKNOWLEDGEMENT
The work is supported by the National Key R&D Program of China (No. 20222D016070),, STCSM (No. 22511 105700, No. 21DZ1 100100,,, 111 plan (No. BP0719010), National Natural Science
10

---

Published as a conference paper at ICLR 2025
Foundation of China (No. 62306178), Shenzhen Basic Research Project (Natural Science Foundation) Basic Research Key Project (NO, JC12204410012449941)) CCF-Baidu Open fund (NO. CCF- Baidu202413), and the National Research Foundation, Singapore, under its NRF Professorship Award No. NRF-P2024-001. Ziqing Fan is partially supported by Wu Wen Jun Honorary Doctoral Scholarship, Al Institute, Shanghai Jiao Tong University.
REFERENCES
Amro Kamal Mohamed Abbas, Kushal Tirumala, Daniel Simig, Surya Ganguli, and Ari Morcos. Semdedup: Data-efficient learning at web-scale through semantic deduplication. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024. 19
Ravikumar Balakrishnan, Tian Li, Tianyi Zhou, Nageen Himayat, Virginia Smith, and Jeff Bilmes. Diverse client selection for federated learning submodular maximization. In International Conference on Learning Representations, 2022. 5, 17
Adrien Bardes, Jean Ponce, and Yann Lecun. Vicreg: Variance-invarianc--ovvaii regularization for sele-supe learning. In ICLR 2022-Intenaaional Conference on Learning Representations, 2022. 2, 5, 17
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference Machine Learning, pp. 2397-2430. PMLR, 2023. 7, 10
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020. 2, 17, 18
Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005 2020a.
Tom B Brown. Language models are few-shot learners. arXiv preprint arXiv:2005.1416. 2020b. 1,
3, 17
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607. PMIR, 2020. 4
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023a.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung. Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023b. 1,3,8 8, 17
Christopher Clark, Kenton Lee, Ming-Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint arXiv:1905.10044 2019, 6, 18
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge.
8rXiv.1803.0045571,, 2018. 2, 6, 18
Together Computer. Redpajama: an open dataset for training large language models, 2023. URL https: http:/gggttt///t com/ toge there cemprrooppterr /eedP a ama-Data Data. 2, 3, 4, 6, 17
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with jo-awareness. Advances in Neural Information Processing Systems, 35: 35:16344-16359, 2022. 7
11

---

Published as a conference paper at ICLR 2025
A DAS. Submodular meets spectral: Greedy algorithms for subset selection, sparse approximation and dictionary selection. In Proc. 28th Int. Conf. on Machine Learning (ICML'II), pp. 1057-1064, 2011. 2, 6, 17
Ziqing Fan, Jiangehao Yao, Ruipeng Zhang, Lingjuan Lyu, Ya Zhang, and Yanfeng Wang. Fed- erated learning under partially class-disjoint data via manifold reshaping. arXiv preprint urXiv:2405.18988 2024. 2, 17
Satoru Fujishige. Submodular functions and optimization. Elsevier, 2005. 5, 17
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac' h. Haonam Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang. Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL httph:tpp///tm :ttp::::////eooooo// wttp://ewrrd////00/00000000 3. 6, 18
Xinyang Geng and Hao Liu. Openllama: An open reproduction of Ilama, May 2023. URL https: //github. com/ opwoomeemmmmmm www.ereeeeeeeeeeeeeeeeeeeaeaeere haee/ppee________ 11ama. 21
Andrew Guillory and Jeff A Bilmes. Online submodular set cover, ranking, and repeated active learning. Advances in neural information processing systems, 24, 2011. 17
Dan Hendrycks. Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations, 2021. 6, 18
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint wrXiv:2203.15533 2022.
Feng Hong, Yueming Lyu, Jiangehao Yao, Ya Zhang, Ivor Tsang, and Yanfeng Wang. Diversified batch selection for training acceleration. In Forty-first International Conference on Machine Learning, 2024. 5, 17
Gautier Izacard, Mathilde Caron, Lucas Hosseini, Sebastian Riedel, Piotr Bojanowski, Armand Joulin, and Edouard Grave. Unsupervised dense information retrieval with contrastive learning, 2021. URL https: / arxiv, orgg abs / 21 12. 09118. 2, 4, 7, 10
Eeshaan Jain, Tushar Nandy, Gaurav Aggarwal, Ashish Tendulkar, Rishabh Iyer, and Abir De. Efficient data subset selection to generalize training across models: transductive and inductive networks. Advances in Neural Information Processing Systems, 36, 2024. 17
Shuyang Jiang, Yusheng Liao, Yanfeng Wang, Ya Zhang, and Yu Wang. Fine-tuning with reserved ma- jority for noise reduction. In The Thirteenth International Conference on Learning Representations, 2025. URL https: opwnrevieweeeeeeeeee/eeee. net/ forum?
17
Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. In International Conference on Learning Representations, 2022. 2, 4, 17
Athresh Karanam, Krishnateja Killamsetty, Harsha Kokel, and Rishabh Iyer. Orient: Submodular mutual information measures for data subset selection under distribution shift, Advances in neural information processing systems, 11311796-188000 2022. 17
Vishal Kaushal, Rishabh Iyer, Suraj Kothawade, Rohan Mahadev, Khoshrav Doctor, and Ganesh Ramakrishnan. Learning from less data: A unified data subset selection and active learning framework for computer vision. In 2019 IEEE Winter Conference on Applications of Computer Vision (WACV), pp. 1289-1299. IEEE, 2019. 17
Suraj Kothawade, Vishal Kaushal, Ganesh Ramakrishnan, Jeff Bilmes, and Rishabh Iyer. Prism: A rich class of parameterized submodular information measures for guided data subset selection. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pp. 10238-10246, 2022. 17
12

---

Published as a conference paper at ICLR 2025
Andreas Krause and Daniel Golovin. Submodular function maximization. Tractability, 3(71-104):3, 2014. 17
Lilly Kumari, Shengjie Wang, Amav Das, Tianyi Zhou, and Jeff Bilmes. An end-to-end submodular framework for data-efficient in-context learning. In Findings of the Association for Computational Linguistics: NAACL 2024, pp.
2024. 17
Roy Ka-Wei Lee, Rui Cao, Ziqing Fan, Jing Jiang, and Wen-Haw Chong. Disentangling hate online memes, In Proceedings of the 29th ACM international conference multimedia, pp. 5138-5147, 2021. 17
Raymond Li, Loubna Ben Allal, Yangtian Zi, Niklas Muennighoff, Denis Kocetkov, Chenghao Mou, Mare Marone, Christopher Akiki, Jia Li, Jenny Chim, et al. Starcoder: may the source be with you! arXiv preprint arXiv:2305.06116/ 2023. 21
LightningAl. Litgpt. /ttppp:://ithhee htt.:tthub..oommoommmmm wtt..///tgtiinggg//////o 2023. 7 2023.
Hui Lin and Jeff A Bilmes. Optimal selection of limited vocabulary speech corpora. In INTER- SPEECH, pp. 1489-1492, 2011.
Hya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Reprexentations, 2019.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? a new dataset for open book question answering. In EMNLP, 2018. 2, 6, 18
George Nemhauser, Laurence A Wolsey, and Marshall L Fisher. An analysis of approximations for maximizing submodular set functions-i. Mathematical programming, 14 14:265-294 1978. 5, 17
Xianghe Pang, Shuo Tang, Rui Ye, Yuxin Xiong, Bolun Zhang, Yanfeng Wang, and Siheng Chen. Self-alignment of large language models via monopolylogue-bant -basdd social scene simulation. arXiv preprint 4rXiv:2402.05699 2024. 17
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
David Patterson, Joseph Gonzalez, Quoc Le, Chen Liang. Lluis-Miquel Munguia, Daniel Rothchild, David So, Maud Texier, and Jeff Dean. Carbon emissions and large neural network training. arXiv preprint arXiv:2 104. 10350, 2021. 1, 17
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog. 1(8):9, 2019, 8, 10
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. International conference on machine learning, pp. 8748-8763 PMLR, 2021. 8, 10
Jack Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, et al. Scaling language models: Methods, analysis & insights from training gopher. arXiv preprint arXiv:2112.11444 2021. 1, 17
HSVNS Kowndinya Renduchintala, Krishnateja Killamsetty, Sumit Bhatia, Milan Aggarwal, Ganesh Ramakrishnan, Rishabh K Iyer, and Balaji Krishnamurthy. Ingenious: Using informative data subsets for efficient pre-training of language models. In The 2023 Conference Empirical Methods in Natural Language Processing, 1991. 2, 4, 8,
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9)999-166, 2021. 6, 18
13

---

Published as a conference paper at ICLR 2025
Said Salhi. Discrete location theory. Journal of the Operational Research Sociery, 42(12):1 1991a. 19
Said Salhi. Discrete location theory. Journal of the Operational Research Sociery, 42(12): 122-1125,
1991b. 17
Richard Santiago and Yuichi Yoshida. Weakly submodular function maximization using local submodularity ratio. arXiv preprint 24te: 2020. 17
Yujun Shi, Kuangqi Zhou, Jian Liang, Zihang Jiang, Jiashi Feng, Philip HS Torr, Song Bai, and Vin- cent YF Tan. Mimicking the oracle: An initial phase decorrelation approach for class incremental learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. pp. 16722-16731, 2022. 2, 4, 5,
Ben Sorscher, Robert Geirhos, Shashank Shekhar, Surya Ganguli, and Ari Morcos. Beyond neural scaling laws: beating power law scaling via data pruning. Advances in Neural Information Processing Systems, 19523-1999,,,, 2022. 19
Aarohi Srivastava, Abhinav Rastogi, Abhishek Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alome et al. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. arXiv preprint arXiv:2206.0400 2022. 18
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc Le, Ed H Chi, Denny Zhou, et al. Challenging big-bench tasks and whether chain-of-thought can solve them. arXiv preprini arXiv:2210.09261, 2022. 6, 18
Shuo Tang, Xianghe Pang, Zexi Liu, Bohan Tang, Rui Ye, Xiaowen Dong, Yanfeng Wang, and Siheng Chen. Synthesizing post-training data for Ilms through multi-agent simulation. arXiv preprint arXiv:2410.14221,, 2024. 17
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng- Heng-Tze Cheng, Alicia Jin, Taylor Bos, Leslie Baker, Yu Du, et al. Lamda: Language models for dialog applications. arXiv preprint arXiv:2201.0822. 2022. 1, 17
Kushal Tirumala, Daniel Simig, Armen Aghajanyan, and Ari Morcos. D4: Improving Ilm pretraining via document de-duplication and diversification. Advances in Neural Information Processing Systems, 36, 2023. 2, 4, 8, 19
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Maric-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023a. 2, 3, 6, 17
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv: 2023,,
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008.
Pingjie Wang, Ziqing Fan, Shengehao Hu, Zhe Chen, Yanfeng Wang, and Yu Wang. Reconstruct the pruned model without any retraining. arXiv preprint arXiv:2407.1333/ 2024. 17
Kai Wei, Rishabh lyer, and Jeff Bilmes. Submodularity in data subset selection and active learning. In International conference on machine learning, pp. 1964-1963. PMLR, 2015. 17
Alexander Wettig, Aatmik Gupta, Saumya Malik, and Danqi Chen. Qurating: Selecting high-quality data for training language models. In Forty-first International Conference Machine Learning, 2024. 2, 4, 6, 8, 19 19
Sang Michael Xie, Hieu Pham, Xuanyi Dong. Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. Advances in Neural Information Processing Systems, 36, 2023a. 6, 8, 19
14

---

Published as a conference paper at ICLR 2025
Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. Advances Neural Information Processing Systems, 36: 34201-34227, 2023b. 2, 4, 8, 17, 19
Rui Ye, Xianghe Pang, Jingyi Chai, Jiaao Chen, Zhenfei Yin, Zhen Xiang, Xiaowen Dong, Jing Shao, and Siheng Chen. Are we there yet? revealing the risks of utilizing large language models in scholarly peer review. arXiv preprint arXiv:2412.01708, 2024. 17
Jingwei Yi, Rui Ye, Qisi Chen, Bin Zhu, Siheng Chen, Defu Lian, Guangzhong Sun, Xing Xie, and Fangzhao Wu. On the vulnerability of safety alignment in open-access Ilms. In Findings of the Association for Computational Linguistics ACL 2024, pp. 9236-9260, 2024.
Jure Zbontar, Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In International conference machine learning, pp. 12310- 12320. PMLR, 2021. 2, 4, 5, 17
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv: /0Xx::/90500083000 2019. 2, 6, 17, 18
Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model. arXiv preprint S24iv:22001022888 2024a. 2, 7,8
Qixin Zhang, Zengde Deng, Xiangru Jian, Zaiyi Chen, Haoyuan Hu, and Yu Yang. Communication- efficient decentralized online continuous dr-submodular maximization. In Proceedings of the 32nd ACM International Conference on Information and Knowledge Management, pp. 3330-3339, 2023.
Qixin Zhang, Zongqi Wan, Zengde Deng, Zaiyi Chen, Xiaoming Sun, Jialin Zhang, and Yu Yang. Boosting gradient ascent for continuous dr-submodular maximization, arXiv preprint arXiv:2401.08330, 2024b. 5
Qixin Zhang, Zongqi Wan, Yu Yang, Li Shen, and Dacheng Tao. Near-optimal online learning for multi-agent submodular coordination: Tight approximation and communication efficiency. arXiv preprint arXiv:2502..0000 2025. 5
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint arXiv:2205.01065 01068, 2022. 7, 10
Tianyi Zhou and Jeff Bilmes. Minimax curriculum learning: Machine teaching with desirable difficulties and scheduled diversity. In International conference on learning representations, 2018. 17
15

---

Published as a conference paper at ICLR 2025
A APPENDIX
In the Appendix, we provide additional information including dataset, baselines, model and training details, and the proof of Lemma I, As shown in the following, to help readers read, we provide table of contents of the full paper.
CONTENTS
1 Introduction
2 Rethinking File Selection for LLM Pre-training under Budget
2.1
Problem Statement
2.2 Recent Selection Methods
2.3
Diversity Dilemma: Dimensional Collapse
3
DiSF: Diversified File Selection
3.1
Method
3.2
Analysis
4
Experiment
4.1 Setup
4.2 Main Results
4.3
Efficiency Analysis
4.4
Ablation Study
Conclusion
10
6 Acknowledgement
10
A Appendix
16
A.1 Related Work
17
A.2
Dataset:SlimPajama
17
A.2.1 Evaluation
18
A.3 Baselines
19
A.3.1 DSIR.
19
A.3.2 QuRating.
19
A.3.3 Doremi.
19
A.3.4
INGENIOUS
19
A.3.5
D4.
19
A.4
Model and training details
20
A.5
More analysis
21
A.5.1
Performance on larger dataset and model scale
21
A.5.2 Training efficiency compared to all baselines
21
16

---

Published as a conference paper at ICLR 2025
A.5.3
Proof of lower bound on Submodular ratio
22
A.5,4
Complexity analysis
23
A.5.5
Comparison to submodular functions
23
A.5.6
Observation of dimensional collapse on Starcoder
23
A.6 Proof of Lemma
24
A.1 RELATED WORK
Efficient File Selection Pre-trained Large Language Models (LLMs) have exhibited exceptional performance across wide range of downstream tasks (Jiang al., 2025; Ye et al., 2024; Pang et al., 2024; Yi et al., 2024; Wang et al., 2024; Lee et al., 2021: Tang et al., 2024). However, their training process computationally intensive, with costs escalating significantly size and training data volume increase (Rae et al., 2021; Patterson et al., 2021; Thoppilan et al., 2022). To maximize the performance of L.MMs within constrained budgets, it is crucial to curate high-quality pre-training data from extensive text corpora, as this enhances both training efficiency and sample effectiveness. Recent advancements in pre-training file selection for LLMs primarily rely on leveraging existing or trained proxy models, coupled with proxy functions designed to assess the similarity of data to a target domain, which is typically considered high-quality. For instance, heuristic classification methods (Brown, 2020b; Chowdhery et al., 2023b) employ binary classifiers to filter content resembling text from domains such as Books and Wikipedia (Computer, 2023). Similarly. DSIR (Xie et al., 2023b) focuses on these domains by utilizing hashed n-gram extractor to quantify similarity. Another approach, QuRating (Wettig et al., 2024), utilizes GPT- 3.5-turbo to train judge model that evaluates the quality of domains like writing and education. Despite these innovations, methods anchored to specific domains face a diversity dilemma, akin to the phenomenon of dimensional collapse observed in representation learning (Jing al., 2022; Zbontar et al., 2021; Bardes et al., 2022; Shi et al., 2022; Fan et al., 2024). These approaches often result in feature vectors that span lower-dimensional subspace, reflecting reduced diversity. While this may enhance performance in Jomain-speciff tasks such as reading comprehension, it leads to significant performance degradation across broader, more diverse domains, particularly in tasks requiring physical world knowledge, such as PIQA (Bisk et al., 2020) and HellaSwag (Zellers et al., 2019). To tackle this issue and promote greater diversity, we introduce novel diversified file selection algorithm, termed DiSF, designed to curate text samples that span a more uniform distribution of feature dimensions within the embedding space.
both model
Submodular optimization. Submodularity is property of functions defined on set Ω that exhibit diminishing returns. Submodular functions like facility location, log determinant, and graph cut (Salhi, 1991b; Fujishige, 2005; Krause & Golovin, 2014; Kaushal et al., 2019; Karanam et al., 2022) are widely recognized for effectively modeling diversity (Nemhauser et al., 1978; Wei et al., 2015; Balakrishnan et al., 2022; Hong al., 2024). Moreover, submodular optimization has achieved significant success in various fields, such as summarization (Kothawade et al., 2022; Kumari et al., 2024), curriculum learning (Balakrishnan al., 2022; Zhou & Bilmes, 201 8), active learning (Wei et al., 2015; Guillory & Bilmes, 2011), I), and subset selection (Jain et al., 2024; Lin & Bilmes, 2011), where selecting diverse subsets is crucial. Given set Q
function f: 2 - R is y-weakly
submodular if and only if, for any subsets A B C Ω and a element x E Ω \ B. the following inequality holds:
(x)) (()) 2 (I(BU (x)) I(B)),
where 7 E (0. When y (Santiago Yoshida, 2020; DAS, 2011). the function is submodular function. A non-negative monotone y-weakly submodular maximization problem can be solved using the classical greedy algorithm, which guarantees a (1 appprooioxinn this work, our function aims to evaluate diversity, making well suited to be verified under this formulation.
A.2 DATASET:SLIMPAJAM
SlimPajama is a high-quality text corpus specifically created for pre-training large language mod- els. This corpus, derived from RedPajama (Computer, 2023), underwent additional cleaning and deduplication processes. The original RedPajama corpus, an open-source research project, was designed to replicate the pretraining data of Llama (Touvron et al., 2023a) and contains 1.2
17

---

Published as a conference paper at ICLR 2025
trilliin tokens. After extensive filtering to remove low-quality and duplicate content, SlimPajama retains only 50% the original tokens from RedPajama. Notably, compared to domains like ArXiv, GitHub, and StackExchange. Heuristic classification and DSIR tend to select higher proportion of files from Books and Wikipedia, and Qurating-V W also favors the Books domain. This bias toward specific domains may explain why these methods encounter dimensional collapse and performance degradation in overall task performance.
A.2.1 EVALUATION
Harness evaluation framework. To demon- strate the diversity dilemma and evaluate the general performance of the pre-trained LLMs, we utilize seven commonsense reasoning tasks from the widely recognized evaluation frame- work, Harness (Gao et al., 2024), including four reading comprehension tasks: OBQA (Mi- haylov et al., 2018): Inspired by open book ex- ams, this task tests the ability comprehend and apply knowledge similarly to human understanding; ARC-e and ARC-c (Clark et al., 2018):
Table 6: Commonsense performance of IN- GENIOUS when using Warmed-up model (INGENIOUS W) and Contriever (INGENIOUS- C) with 50B training budget.
Method
Tinyl 120M TinyL.ama-11118
INGENIOUS W
39.0
39.6
INGENIOUS-C
43.2
44.1
7,787 multiple-choice science questions at a grade-school level, divided into an easy set (ARC-e) and challenge set (ARC-c); BoolQ (Clark et al., 2019): A reading comprehension task that focuses on naturally occurring yes/no questions, and three physical world knowledge tasks: HellaSwag (Zellers et al., 2019): A collection to assess physically situated commonsense reasoning capabilities; WinoGrande (Sakaguchi et al., 2021): I): An expansion of the Winograd Schema Challenge (WSC) with increased scale and complexity; PIQA (Bisk et al., 2020): A task to measure the understanding and reasoning about physical interactions in the real world. To further verify our effectiveness, we also employ two problem-solving tasks from Harness: MMLU (Hendrycks et al., 2021): A task to measure the world knowledge and problem-solving capabilities across various subjects; BBH (Suzgun et al., 2022): A subset 23 challenging tasks from the BIG-Bench benchmark (Srivastava et al., 2022) to measure the ability of complex instruction following.
Table 7: Commonsense performance when pre-trained on SlimPajama with Open-Llama 3B, 1.5% selection ratio and 1OB training budget.
Method
Random
DSIR
QuRating-W
QuRating-A DiSF (Ours)
Commonsense ability
43.1
41.9
42.6
43.6
44.4
Table Commonsense performance when pre-trained on both StarcoderData, and SlimPajama with TinyLlama 1.1B, 1.5% selection ratio and 10B training budget.
Method
Random
DSIR
DiSF (Ours)
Commonsense ability
40.3
40.0
41.4
bbh
12.6
12.5
13.3
mmlu
23.0
22.1
23.7
code_x_glus
0.80
0.59
0.86
Table 9: Comparison on sample efficiency pre-trained after 50B tokens. We record selected tokens to achieve the performance of Full Data pre-trained with 50B tokens. The ablation interval is 0.5%.
Setting
Full-Data
DSIR
QuRating-W W QuRating-A
D4
DiSF (Ours)
TinyLlama- 1 120M
1.00
0.60
0.40
0.15
0.20
0.15
LlayLaama-5610
1.00
0.80
0.45
0.20
0.25
0.20
TinyL.ama-II.1
1.00
0.80
0.40
0.25
0.25
0.20
18

---

Published as a conference paper at ICLR 2025
Table 10: Comparison training efficiency pre-trained with 1.5% selection ratio. We record pre- trained tokens to achieve the performance of Full Data pre-trained with 50B tokens. denotes the method can not reach that performance. The ablation interval is IB.
Setting
Full-Data
DSIR
QuRating-W W QuRating-A A
D4
DiSF (Ours)
TinyLlama 120M
50B
32B
36B
27B
TinyLlama-5600M
50B
46B
47B
36B
TinyL.lmma-I.IB
50B
52B
56B
40B
A.3 BASELINES
A.3.1 DSIR.
DSIR (Xie al., 2023b) treats Books and Wikipedia as high-quality targets for file selection, employing hashed n-gram feature extractor to measure the similarity between the text features and the target distribution. In our experiments, following the selection procedures outlined DSIR, we calculate importance scores using data (SlimPajama) and target data (Wikipedia and Books) in an n-gram feature space. The importance weights are then applied to resample a subset of the raw dataset. As for files in Wikipedia and Books domains, proportionally integrate into the selected dataset.
A.3.2 QURATING.
QuRating (Wettig et al., 2024) queries GPT-3.5-turbo to train a judge model, that assess the specific quality of text samples, including four criteria: writing style, required expertise, facts & trivia, and educational value. this paper, for Qurating-WV W, we select samples with the highest scores for writing style, while for QuRating-A. A, we proportionally select top-scoring samples across all four criteria.
A.3.3 DOREMI.
Doremi (Xie et al., 2023a), recently proposed method that produces domain weights for pre-training on multiple text domains. our experiments, we follow the domain weights calculation process of Doremi, using the domain weights from the initial data distribution small reference model on the SlimPajama dataset. We then leverage this reference model to guide the training of small proxy model, designed to generate domain weights. Finally, these domain weights are employed as the random selection ratio for text domains to construct selected dataset.
initial reference to train
A.3.4 INGENIOUS
INGENIOUS (Renduchintala al., 1991) extracts features using a model after warm-up phase and employs Facility Location (Salhi, 1991a) to design proxy function for feature importance, which measures the similarity between samples the embedding space. Notably, as shown in Table 6, observed that INGENIOUS. when using warmed-up model for feature extraction (INGENIOUS- W), does not achieve satisfactory performance with our selected feature extractor, Contriever (INGENIOUS Although INGENIOUS-C performs competitively compared Ran- dom, DSIR, QuRating- W, and QuRating-A, our method consistently achieves the best performance across all settings.
A.3.5 D4.
The recent method D4 (Tirumala et al., 2023) notices the importance of diversified selection, involving SemDeDup (Abbas et al., 2024) and Prototypicality (Sorscher et al., 2022) to reduce file redundancy, but not achieve satisfactory uniform representations ours. In this paper, sequentially applied the SemDeDup and Prototypicality methods to filter the data, controlling the filtering ratios of these two steps be Ratedup 0.75 and Rymio 0.02, respectively.
19

---

Published as a conference paper at ICLR 2025
Table 11: Model structure and training details of pre-training TinyL 120M.
Parameter name
Value
Parameter number
121,129,728
Hidden size
768
Intermediate Hidden Size
2048
Context Len
2048
Heads
12
Layers
12
Vocab size
32000
Minimum learning rate
4e-5
Maximum learning rate
4e-4
Optimizer
AdamW
B1 of optimizer
0.9
B2 of optimizer
0.95
Warmup steps
2000
Batch size
2M tokens
Weight decay
0.1
Activation function
SwiGLU
Gradient clipping threshold
1.0
Platform
8 NVIDIA GeForce RTX 4090 GPUs
Training times 10B tokens
about 0.2 days
Training times on 50B tokens
about I days
Table 12: Model structure and training details of pre- vreitraining TinyLlama 560M.
Parameter name
Value
Parameter number Hidden size
561, 201,072,122
2048
Intermediate Hidden Size
2048
Context Len
2048
Heads
16
Layers
20
Vocab size
32000
Minimum learning rate
4e-5
Maximum learning rate
4e-4
Optimizer
AdamW
31 of optimizer
0.9
B2 of optimizer
0.95
Warmup steps
2000
Batch size
2M tokens
Weight decay
0.1
Activation function
SwiGLU
Gradient clipping threshold
1.0
Platform
8 NVIDIA GeForce RTX 4090 GPUs
Training times on 10B tokens
about 0.9 days
Training times on 50B tokens
about 4.5 days
A.4 MODEL AND TRAINING DETAILS
For better clarity and reproducibility, provide the model structures and training details for pretraining LinyLl with 120M, 560M. and 1.1B parameters, as shown Tables 11, 12, and 13, respectively. These tables detail the configurations used in our experiments, facilitating easier
20

---

Published as a conference paper at ICLR 2025
Table 13: Model structure and training details of pre-training TanyL.ama 1.1B.
Parameter name
Value
Parameter number
1,100,048,384
Hidden size
2048
Intermediate Hidden Size
5632
Context Len
2048
Heads
32
Layers
22
Vocab size
32000
Minimum learning rate
4e-5
Maximum learning rate
4e-4
Optimizer
AdamW
B1 of optimizer
0.9
B2 of optimizer
0.95
Warmup steps
2000
Batch size
2M tokens
Weight decay
0.1
Activation function
SwiGLU
Gradient clipping threshold
1.0
Platform
8 NVIDIA GeForce RTX 4090 GPUs
Training times 10B tokens
about 1.7 days
Training times 50B tokens
about 8.5 days
replication of our results. Please note that our goal in this paper is not to optimize all hyper-pale ppermpprrmettrs for the best LLM, but rather to compare selection methods under fair and reasonable conditions.
A.5 MORE ANALYSIS
A.5.1 PERFORMANCE ON LARGER DATASET AND MODEL SCALE
To evaluate the effectiveness of selection method on larger models and datasets, 1) pre-train Open-Llama 3B (Geng Liu, 2023) on SlimPajama, and 2) TinyL 1. on both SlimPajama and StarcoderData (Li al., 2023). StarcoderData is dataset created for code generation, containing approximately 210M files (500GB of data). We assess the pre-trained models on common sense ability (7 tasks) and an additional code task, code -oeex_gght ooduee using harness evaluation. As shown in Table 7, we pre-train open-llama 3B with 10B training budget and 1.5% selection budget with SlimPajama. We compare our method with Random, DSIR, QuRating-A, and QuRating- W. DSIR and QuRating-W W meet performance degradation due to dimensional collapse, while our method achieves the best performance on tasks. As shown in Table 8, we pre-train TinyL-Jamal1 1. on both SlimPajama and StarcoderData with 1.5% selection budget and 1OB training budget, compared to random selection and DSIR. Results of DSIR show performance degradation in code ability. We analyze that, DSIR tends ignore code files in StarcoderData due to the selection critetion is based on WikiPedia, which means dimensional collapse happens.
A.5.2 TRAINING EFFICIENCY COMPARED TO ALL BASELINES
To compare the computational and data efficiency of our approach with all baselines, we present additional results under the same settings as Figures 5 and 6, shown in Tables 9 and 10. The results demonstrate that, to achieve equivalent performance to Full Data with 50B training budget, our method requires the fewest samples and pre-training tokens compared to the baselines, showing promising efficiency. Additionally, We also compare the cost of our method with all baselines in the process of selecting data. As reported in QuRating, annotating the data using the GPT API costs 520 NVIDIA H100 hours with additional ranking procedures of 3 hours. For DSIR, it takes more than days using 48 CPUs in our platform. For DOREMI, training a 120M proxy model to provide weights for domains takes approximately week. In contrast, method utilizes public
21

---

Published as a conference paper at ICLR 2025
feature extractor and selects samples in about 26 hours using GPU and 48 CPUs. Combining these facts and our complexity analysis, we believe our method is practical among these methods for larger datasets.
A.5.3 PROOF OF LOWER BOUND ON SUBMODULAR RATIO
We first recall our function, which is based on F-norm of the covariance matrix. The gain on A, and with another sample e, can be formulated as following:
A भ|F Dis E (ErEd t0a+)));))
TAT V 2it 2n (Zize (
Disi Ei sxzytees)
Eisi E Eren( B.ts)
AB
where AC B C Q E Ω and e & B. We reformulate them as
A = VIV
EieA Fiz) -1),
A TVV DI RBt (e 1BT Σi E/ कहक Et Ef #ennenett 1).
We then define
(ceAA)
1
TE ee,)²²
1
Σ
IA+11 1|
|A|
1
ZEA
J
zEA
d Σ>
Aee BB + 11
1
Σ
J,z e,e,)²
1
Σ
Σ 2,,,)2
B
J
.
ZEB
Therefore, the original equation will be:
a = e VEEEE A
TAI VE
AeeAA)
e
(ee -
To provide theoretical guarantee, we aim to establish lower bound for the submodular ratio under any given set where the gain is positive, as Figure 4 generally suggest. Since the original formulation challenging to analyze directly, introduce two bounds on the function applied to the text files to derive a more concrete lower bound, as follows:
Assumption (Bounded Gain Difference). Assume the gain difference between any two sets are bounded:
[eeAA) - ((e|B)) s..
Assumption 2 (Bounded Average Utility). Assume average utility is μs, which means VU Ω have:
ΣΣ>> UEU
xır,)² S
With Assumption 1, we have:
A A
E. 15
A(e|BB) e- -
Ag
With Assumption 2, we have:
1
Σ 2
1
Σ
22 2 -2p,
B
[
BCB
uEA
22

---

Published as a conference paper at ICLR 2025
which means
AA Ag 2
2p
A(c|B) e-A(e)m)
We reformulate this as
A
e (e c-A(c]B) -1
AB
From Assumption 2, we have have:
EZEB x,0;)² S p. Therefore e A(c(B) 2 ed and we
[
V
A
- e2p e2u-1
24 (c
AB
Finally, we have a lower bound as e 2p -- which means given assumptions 1, 2, and positive gain, our function is at least e 22p -
weakly submodular function.
A.5.4 COMPLEXITY ANALYSIS
For detailed analysis of time complexity, we divide it into two parts: 1) computational complexity shown in Algorithm 1. In each batch, we initialize U, with a randomly selected sample and remove
| from the batch. Then, we iteratively apply (L 1) times the Argmax command on the batch of data with our proxy fuction. Denote the computational cost of our proxy function with k text samples as F (U) = OF&), the computation cost will be:
MSI
D b
DF 1
S
O(1 +
1)
O
DF K+1
O
a: 19
+1);
where b the batch scale, [D| is the total data scale, |S] is the selection budget. 2) The complexity of proxy function OFF.). Given text features and their feature dimension d. Frobenius norm and z zf are both )d Since our proxy function calculates times the z zz" O(Fs) kd", Finally, the complexity of our DiSF will be:
O(DiSF) & O(DISS T +1) O(IS)bba))
All terms in the time complexity are at most quadratic and independent of the overall dataset size, which we consider acceptable for applying to larger datasets. The space complexity largely depends on the stored features of all text files: ((D))
Table 14: Pre-trained performance compared DiSF-LD.
Setting
Random
DSIR
D4
QuRating-VV
QuRating-A DiSF-LD
DiSF (Ours)
TinyLlamal 120M
38.9
37.8
39.7
38.9
40,0
39.5
40.6
A.5.5 COMPARISON TO SUBMODULAR FUNCTIONS
In this section, we compare our proxy function with two strictly submodular functions: facility location and degeeeeermin For the facility location function, we compare with INGENIOUS. Additionally, we introduce variant defined as Fro==LoDe)((+C) (DiSF-LD), where I is the identity matrix and C covariance matrix. We evaluate this variant alongside our original proxy function on TanaL 120M, using 1.5% selection ratio and 50B training budget. As shown in Tables 2, and 14, the results demonstrate that both DiSF-LD and INGENIOUS help mitigate dimensional collapse and improve the performance of pre- traitreieed LLMs. However, neither method achieves the same level of performance as our original DiSF. This is due to their inability to directly optimize the uniformity of feature dimensions, leading trade-off between strict submodularity and the specific goal of optimizing dimensional uniformity.
A.5.6 OBSERVATION OF DIMENSIONAL. COLLAPSE ON STARCODER
As shown in Table 15, we present the selection results of DSIR on Stareoderdata, as well as on both Starcoderdata and SlimPajama, across two domains: Wikipedia and Python. The results demonstrate
23

---

Published as a conference paper at ICLR 2025
Table 15: File selection of DSIR Starcoder. We denote selected file ratios of DSIR based Wikipedia domain under both Starcoderdata and SlimPajama as ASIR-W-SS Wikipedia domain under Starcoderdata DSIR-W- and Python domain under Starcoderdata as DSIR-P-S.
Method
go
java
javascript
php
python
ruby
slimpajama
Original
0.59%
2.52%
2,45%
1.97%
1.61%
0.43%
74.07%
ASIR-W-SS
0.00%
0.77%
0,54%
0.00%
0.5%
0.00%
94.0%
DSIR-W-S W-$
0.03%
12.37%
17.15%
0.12%
7.58%
0.00%
DSIR-P-S
0.00%
25.04%
4.36%
2.20%
50.38%
0.00%
that DSIR, when applied to single domain, tends to select similar files, indicating dimensional collapse. Notably, during our analysis of the scores output by QuRater, we found that the writing style scores for most files in Starcoderdata negative. As result, these files unlikely to be selected when combined with SlimPajama, further highlighting the issue of dimensional collapse.
A.6 PROOF OF LEMMA 1
Lemma 2. Assuming covariance matrix M E R Rdxd computed from the feature of each sample with the standard normalization, and its eigenvalues A2,,22;., ddd. we will have the following equality that satisfied
1
21)
|M d.
Proof. Let M RRdxd be a covariance matrix computed from features that have been standard- ized-that is, the data has been centered (zero mean) and scaled to have unit variance along each dimension. This standard normalization implies that the trace of equals
d Σ
Tr(M)-
A
i=1
This means the average eigenvalue is:
d Σ x; =1
Left-Hand Side (LHS) Calculation:
The left-hand side involves the sum of squared deviations of the eigenvalues from their mean:
d
d
A
= Σ (2s x) = Σ (A: - 1)2.
1=1
Expanding each term on the right, we get the following equation:
(A: 1)2 = Σ (X) - 2A, +1
i=1
Σ x
Xj Σ 1.
i=1
iml
Simplifying with Edn X dnd Σ
d, we have:
24

---

Published as a conference paper at ICLR 2025
Σ
x
2d
Σ
X d.
Right-Hand Side (RHS) Calculation:
First, the Frobenius norm of a matrix M is defined as
|M
Mf = TrM M).
i=1j=1
Since M is symmetric (M M,,, this simplifies to:
| Tr(M²).
Using the spectral theorem, we have
M-UAU
where U is an orthogonal matrix whose columns are the eigenvectors of M, and A = diag (A1, Ag...... contains the eigenvalues. Since U I, we have the following equation:
M (UAUU)(UAUT) UAU = UATU,
Therefore, the trace of M is
Tr(M3) = THUAPTT = -) = IT(A3) = Σ x?.
Thus, the right-hand side
|
Conclusion:
Comparing both sides:
2
d
A
= -
x2 d=
M1 d.
==1
1=1
This completes the proof.
25


Reviews:
[{"review_id": "Pc1QuJWNzy", "reviewer": "Reviewer_2LFP", "summary": "This paper introduces DiSF, a method aimed at selecting diverse training data to enhance the effectiveness of pre-training large language models (LLMs). The authors begin by analyzing recent data selection methods and identify a critical issue: the selected data often suffer from dimensional collapse in the feature space. This collapse leads to improved performance on domain-specific tasks but causes a significant degradation in general performance due to a lack of diversity.\n\nTo address this problem, the authors propose improving the diversity of the data selection process by formulating it as a submodular optimization problem. Specifically, their goal is to prevent dimensional collapse by ensuring more uniform eigenvalues in the covariance matrix of the selected samples. They employ a classical greedy algorithm for submodular optimization, combined with mini-batch sampling, to make the approach tractable.\n\nDue to hardware limitations, the authors tested their approach on models with sizes of 120M, 560M, and 1.1B parameters. Despite the smaller scale, the DiSF method achieves superior performance over six baselines on nine standard benchmarks consistently, across different training budgets (10B and 50B tokens). The authors also provide extensive ablation studies to analyze the effects of data selection budget, different architectural backbones, feature extractors, and more.", "questions": "The paper appears to lack a detailed analysis of the time and computational costs associated with the data selection process. How long does it take to complete the data selection using DiSF, and is it efficient enough to be practical for larger datasets?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ Well-Motivated Research: The paper addresses a significant issue in data selection for LLMs—dimensional collapse due to lack of diversity—which has substantial implications for model performance.\n+ Clear Presentation: The methodology is clearly explained, making the paper easy to read and understand.\n+ Strong Experimental Results: The results are convincing, demonstrating that DiSF outperforms existing methods consistently across multiple benchmarks and model sizes.\n+ Thoughtful Comparisons: Comparisons with six baseline methods are thorough, providing a solid context for the effectiveness of DiSF.\n+ Comprehensive Ablation Studies: The extensive ablation studies offer valuable insights into the impact of various factors such as data selection budget, model architecture, and feature extractors.", "weaknesses": "+ Limited Scale of Experiments: The experiments are limited to models up to 1.1B parameters. While the results are promising, it would be beneficial to see how the method scales to larger models commonly used in current LLM research.\n+ Scalability Concerns: There is uncertainty about whether the proposed method can maintain its efficiency and effectiveness when applied to larger datasets.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "SCSIuGRPmV", "reviewer": "Reviewer_H11o", "summary": "The paper introduces the DiSF (Diversified Submodular File Selection) algorithm to combat dimensional collapse in Large-Language Model (LLM) pretraining. The DiSF algorithm is designed as a subset selection task with a goal of selecting a diverse subset from an unlabeled training corpus and is achieved by maximizing a submodular information function over the complete unlabeled set with sample interactions modeled through a normalized covariance matrix. The paper demonstrates upto 1.5x training efficiency and upto 5x data efficiency on several benchmark datasets.", "questions": "Please refer to the weaknesses for detailed questions.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The strengths of this paper can be summarized as below - \n\n1. The paper address the task of data-efficient LLM pretraining while boosting performance on several benchmark tasks by minimizing the impact of dimensionality collapse which is a critical challenge in this domain. \n\n2. The authors clearly demonstrate the impact of dimensionality collapse in existing approaches targeting data-efficient LLM training and establish the requirement for decorrelated text samples during the training process to minimize its effect.\n\n3. The DiSF algorithm models this requirement as a diversity driven selection problem and achieves this my training the underlying LLM on a subset of data selected by maximizing a submodular function $C(U, M)$ over the complete unlabeled set $\\mathbb{D}$. This proposed methodology is strongly supported by experimental evidence on multiple benchmark datasets.", "weaknesses": "Although the results demonstrated in the paper are impressive the following weaknesses/ questions need to be addressed.\n\n1. Although the subset selection mechanism in DiSF depends on submodular maximization of $F_{M}^{DiSF}(U)$ but the proof of submodularity of $F_{M}^{DiSF}(U)$ has not been provided. The paper presents only an empirical proof of monotonicity through Figure 4 but that is not deemed sufficient (Fujishige, 2005) to prove a function to be submodular. The authors must provide proof of $F_{M}^{DiSF}(U)$ demonstrating the diminishing marginal returns property described in lines 232- 234.\n\n2. Diversity driven subset selection for data-efficient LLM pretraining is not novel to the current paper. Prior work like INGENIOUS (Renduchintala et al., 2023) have demonstrated significant improvements in model performance while being data-efficient. The authors should compare against such methods to study the impact of dimensionality collapse and the diversity selection algorithm.\n\n3. Submodular Functions like Log-Determinant, Facility-Location etc. have been shown to model diversity in several tasks like summarization (Kumari et al., 2024, Kothawade et al., 2022) and subset selection (Jain et al., 2023, Lin and Bilmes, 2011) to name a few. What is the justification of introducing only $F_{M}^{DiSF}(U)$ (Equation 7 in the paper) rather than any known submodular functions ? The The authors should compare the performance of DiSF when the underlying function is chosen to be some/all of the popular submodular functions proven to be effective for diversity selection.\n\n4. Although the experimental results in Table 2 show a strong average performance across datasets, when the budget increases from 10B to 50B (5x increase in number of available training examples) the performance gain (overall) is quite small. Did the authors perform any additional experiments or theoretical analysis to justify this result ?\n\nSome minor questions/ suggestions -\n1. Section 2.3 can be moved higher up in the paper or a brief explanation on the cause and effects of dimensionality collapse should be highlighted in Section 1 for better readability.\n\n2. Phrases like \"Diversified Submodular File Selection\" (line 19, 88, 183-184 ...), \"decorrelated text files\" (line 20, 89, 184 ...) has been repeated several times in the paper which should be reduced to improve readability.\n\n3. Do the authors plan on releasing the code for reproducing the results from the paper in the near future ?\n\n4.  In Equation 8 the $\\text{argmax}$ is applied over $\\mathbb{D} \\setminus U$ which seems to be redundant. To the best of my knowledge, during the course of the greedy optimization process in Algorithm 1 of the paper, DiSF selects the sample with the highest information gain which inherently excludes duplicate samples which are already selected in U.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "KZioZNk0TY", "reviewer": "Reviewer_7MY9", "summary": "The authors propose a diversified submodular file selection algorithm (DiSF) to select sample set with more diverse data. It can alleviate dimensional collapse problem in feature space and enable LLM model to have generic performance.", "questions": "1. It is recommended to provide the sample size and performance analysis experiments under the full data case and data selection case in different datasets.\n2. It is suggested to provide time complexity and space complexity metrics on Table 2. The experiment in Section 'Selection scale' is not enough to illustrate the computational efficiency of the proposed method.\n3. The performance improvement in Table 2 is limited, most of which are within ±0.5 (within the error range), and reasonable analysis is needed to illustrate the effectiveness of the method.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The paper addresses an important problem in LLM pretraining and proposes a simple yet effective method. The writing is clear and the narrative is easy to follow, facilitating understanding of complex concepts.", "weaknesses": "No major weakness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "NnNK0kqdnu", "reviewer": "Reviewer_8JRU", "summary": "This paper proposes a novel file selection method for LLM pretraining to avoid the dimensional collapse issue. The method can achieve greater uniformity across feature dimensions, while the other methods usually get a long narrow feature space, limiting their performance on all the tasks.", "questions": "What does (1-e^-1) approximation mean in the abstract.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper observes the drawbacks of the other file selection methods for LLM pretraining, which usually obtain quite narrow feature space, and are hard to obtain good results for all the tasks.\n\n2. The proposed method is interesting, which converts the uniformity of the eigenvalues of the covariance matrix into minimizing the F norm of the covariance matrix, then uses a submodular optimization to solve it.\n\n3. Experimental results on 7 tasks show the proposed method achieves a higher average perfomance than the other baselines, and is quite data efficient compared to full data pretraining.", "weaknesses": "1. The observation in figure 2 is evaluated on the SlimPajama. Do these methods have the same phenomenan on the other pretraining corpus.\n\n2. In the efficiency analysis, comparing with full data pretraining is not enough, the comparison with the other baselines is not given.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Oral)