Combatting Dimensional Collapse in LLM Pre-Training Data via Submodular File Selection

Thanks for the constructive feedback provided by the Reviewer 8JRU. We sincerely appreciate the time and effort you dedicated to evaluating our work. Below, we provide detailed responses to the weaknesses and questions, which we denote as W and Q. Besides, we also added the following discussions into revised submission.

W1: The observation in figure 2 is evaluated on the SlimPajama. Do these methods have the same phenomenan on the other pretraining corpus.

Reply to W1: Thanks for the valuable comments! We observe the similar phenomenon (dimensional collapse) with additional text corpora, starcoderdata (210M files, approximately 500GB) [1], which is specialized for code generation. In the following, we verify the phenomenan from three views:

1. Selecting similar files on other pretraining corpora. As shown in following Table T1, we present the selection results of DSIR on Starcoderdata and both Starcoderdata and Slimpajama involving two domains: Wikipedia and Python. All results demonstrate that DSIR, when operating on a single domain, tends to select similar files, which indicates dimensional collapse. Experiments with QuRating are still ongoing, as it requires running its judger model (QuRater-1.3B) to evaluate all files, and we expect to complete it in the coming days. Notably, during our analysis of the scores output by QuRater, we found that the writing style scores for most files in Starcoderdata are negative. This makes these files unlikely to be selected when combined with SlimPajama, further underscoring the issue of dimensional collapse.
2. Larger dominance score on other pretraining corpora. As shown in following Table T2, we present the dominance score$\frac{\sum_{i=1}^k\lambda_i}{\sum_{j=1}^d\lambda_j}$ of feature covariance matrix for files selected by DSIR respectively based on the Wikipedia domain and the Python domain, compared to random selection and our DiSF. Results show that regardless of the domain DSIR is based on, the dominance score of its selected files is higher than that of random selection, let alone our method. This indicates a severe issue of dimensional collapse in DSIR's selected files.
3. Performance degradation on other pretraining corpora. As shown in following Table T3, we compare our method with DSIR (based on the Python domain), random selection, and Full Data pre-training. The model is evaluated on the 'code_x_glue' task [2] using the Harness framework on TinyLlama-120M. This task consists of six sub-tasks: Go, Java, JavaScript, PHP, Python, and Ruby. The results show that while DSIR performs better on Python and Java, it achieves poor average performance across the sub-tasks, highlighting dimensional collapse. In contrast, our method achieves the best average performance among all compared baselines.

Table T1: Selected file ratios of DSIR based on Wikipedia domain under both Starcoderdata and SlimPajama (DSIR-W-SS), on Wikipedia domain under Starcoderdata (DSIR-W-S) and Python domain under Starcoderdata (DSIR-P-S).

Table T2 Dominance score compared on starcoderdata.

Table T3: Code ability on six sub-tasks with 1.5% selection and 50B training budget on TinyLlama120M.

[1] Starcoder: may the source be with you!

[2] CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation

Thanks for the valuable comments provided by the Reviewer 7MY9. We sincerely appreciate the time and effort you dedicated to evaluating our work. Below, we provide detailed responses to the questions, which we denote as Q. Besides, we also added the following discussions into revised submission. The corresponding reference is listed at last response.

Q1: It is recommended to provide the sample size and performance analysis experiments under the full data case and data selection case in different datasets.

Reply to Q1: We conduct sample size and performance analysis on another dataset, Starcoderdata [1] (210M files, approximately 500GB), which is specialized for code generation. The model is evaluated on the 'code_x_glue' [2] task using the Harness framework on TinyLlama-120M. This task comprises six sub-tasks, including Go, Java, JavaScript, PHP, Python, and Ruby. We present results compared to Full Data, Random Selection, and DSIR. Comparison with QuRating is ongoing, as it requires running its judger model (QuRater-1.3B) to evaluate all files, and we anticipate completing it in the coming days. Below, we provide detailed analysis on: 1) experiments related to dimensional collapse, 2) performance comparisons on sub-tasks, and 3) performance comparisons across different sample sizes.

1. Dimensional collapse is observed in another dataset with DSIR and QuRating-W. As shown in following Table T1, we present the selected results for SlimPajama and Starcoderdata using DSIR across two domains: Wikipedia and Python. All results demonstrate that DSIR, when operating on a single domain, tends to select similar files, indicating dimensional collapse. Additionally, while processing the scores output by QuRater, we found that the writing style scores for most files in Starcoderdata are negative. This makes them unlikely to be selected when combined with SlimPajama, further highlighting the issue of dimensional collapse.
2. Best averaged performance on another dataset. As shown in following Table T2, we compare our method with DSIR (based on the Python domain), random selection, and Full Data pre-training on Starcoderdata. The results show that DSIR performs better on Python and Java but achieves poor averaged performance across subtasks. In contrast, our method achieves the best averaged performance among all selection baselines.
3. Scalability with Sample Size on another dataset. As shown in following Table T3, we compare our method with Full Data, Random selection, and DSIR (based on the Python domain) across different sample sizes. The results demonstrate our method consistently achieves the best performance under different sample sizes, compared to selection methods. Besides, under a larger sample scale (exceed 3%), our method outperforms Full-Data.

Table T1: Selected file ratios of DSIR based on Wikipedia domain under both Starcoderdata and SlimPajama (DSIR-W-SS), on Wikipedia domain under Starcoderdata (DSIR-W-S) and Python domain under Starcoderdata (DSIR-P-S).

Table T2: Programming ability on six sub-tasks with 1.5% selection and 50B training budget on TinyLlama120M.

Table T3: Problem solving and programming ability (averaged performance on six sub-tasks) on TinyLlama-120M pre-trained on starcoderdata with different sample size and 50B training budget.

[1] Starcoder: may the source be with you!

[2] CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation

Thanks for the constructive feedback provided by the Reviewer H11o. We sincerely appreciate the time and effort you dedicated to evaluating our work. Below, we provide detailed responses to the weaknesses and questions, which we denote as W and Q. Besides, we also added the following discussions into revised submission.

W1: Although the subset selection mechanism in DiSF depends on submodular maximization of F but the proof of submodularity of F has not been provided. ..empirical proof of monotonicity is not deemed sufficient to prove a function to be submodular.

Reply to W1: We thank for the constructive comments for improving our submission. In the original manuscript, we intend to provide approximation analysis when our function showing diminishing gain on set of feature (original description: 'our proxy function aims to capture the diversity in a given set, making it well-suited for submodular optimization'), instead of claiming it is a strictly submodular function on all sets. To enhance clarity and practicality, we have improved the descriptions in the revised submission, emphasizing that our function may not exhibit the submodular property across all file sets. Additionally, we employ more general $\gamma$-weakly submodular optimization properties to empirically verify the approximation to the optimal solution during the execution of our algorithm, specifically on SlimPajama. We provide more details in the following.

1. Revisions are:

• We revised section 3 for better clarity. Specifically, we divide original section 3 into method and analysis, and move all contents related to submodular optimization to analysis. Besides, for a more concrete and practical verification to estimate its approximation to optimal solution during execution, we step further from submodular to $\gamma$-weakly submodular optimization (See details in the next item).
• To avoid misunderstanding, we changed some section names from 'submodular selection' to 'diversified selection'.
• To avoid misunderstanding, we changed the descriptions from 'model it as a submodular optimization problem' in title, abstract, and introduction to 'solve it with classical greedy algorithm and analyze the approximation from the view of $\gamma$-weakly submodular optimization'.

2. Non-negative monotone $\gamma$-weakly submodular optimization is more general for set functions, which provides a guarantee of $(1 − e^{-\gamma})$ approximation ($\gamma \in$(0,1]) [1] to the optimal solution with classical greedy algorithm. Submodular ratio $\gamma$ is defined as $\gamma = \min_{A \subseteq B \subseteq \Omega, e \notin B} \frac{f(A \cup \{e\}) - f(A)}{f(B \cup \{e\}) - f(B)},$ where $A\subset B\subset \Omega$.When $\gamma=1$, the function is submodular over the given set. When $0<\gamma< 1$, the function exhibits weak submodularity over the given set. However, when $\gamma\leq 0$, the function is no longer suitable for analysis using this property. In the revision, we estimate $\gamma$ during selection on SlimPajama in batch files to empirically verify the level of approximation. We estimate the submodular ratio $\gamma$ in 1000 selection batches. In each batch $\Omega=b_i$, we randomly sample 1000 cases of $A$, $B$ and $e$ from $b_i$, recording the smallest $\gamma$. Although submodular ratio often shows very pessimistic bound in applications [2], as shown in the following Table T1, most cases are within (0,1] and show a quite large value, with an average submodular ratio $\hat{\gamma}=0.478$. As shown in the following Table T2, we also calculate the gains $\delta_i$ after adding $i$ times of 32 selected samples across 5 different cases. All cases demonstrate diminishing returns in the gains. As shown in the following Table T2, we also calculate the gains $\delta_i$ after adding $i$ times of 32 selected samples for 5 cases. All cases demonstrate the diminishing gains.

We hope this provide more concrete and practical analysis of our algorithm and can address the concerns of the reviewer.

Table T1 Estimation of submodular ratio $\gamma$ in 1000 batch files.

Table T2 Diminishing return property during execution of algorithm.

[1] Submodular meets Spectral: Greedy Algorithms for Subset Selection, Sparse Approximation and Dictionary Selection

[2] Weakly Submodular Function Maximization Using Local Submodularity Ratio

We are grateful for the endorsement of the reviewer and the increased score. We here to provide a more concrete theoretical guarantee on the submodular ratio.

We first recall our function, which is based on F-norm of the covariance matrix. The gain on A, and B with another sample e, can be formulated as following:

$\Delta_A = e^{-\frac{1}{|A| + 1} \sqrt{\sum_{i=1}^d \sum_{j=1}^d \left( \sum_{x \in A} x_i x_j + e_i e_j \right)^2}} - e^{-\frac{1}{|A|} \sqrt{\sum_{i=1}^d \sum_{j=1}^d \left( \sum_{x \in A} x_i x_j \right)^2}}$, $\Delta_B = e^{-\frac{1}{|B| + 1} \sqrt{\sum_{i=1}^d \sum_{j=1}^d \left( \sum_{x \in B} x_i x_j + e_i e_j \right)^2}} - e^{-\frac{1}{|B|} \sqrt{\sum_{i=1}^d \sum_{j=1}^d \left( \sum_{x \in B} x_i x_j \right)^2}}$,

where $A\subset B \subset \Omega$, $e\in \Omega$, and $e \notin B$. We reformulate them as

$\Delta_A = e^{-\frac{1}{|A|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j)^2}} (e^{\frac{1}{|A|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j)^2}-\frac{1}{|A+1|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j+e_ie_j)^2}} - 1)$, $\Delta_B = e^{-\frac{1}{|B|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j)^2}} (e^{\frac{1}{|B|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j)^2}-\frac{1}{|B+1|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j+e_ie_j)^2}} - 1)$

Therefore, we have:

$\frac{\Delta_A}{\Delta_B}=e^{\frac{1}{|B|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j)^2}-\frac{1}{|A|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j)^2}}*\frac{e^{\frac{1}{|A|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j)^2}-\frac{1}{|A+1|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j+e_ie_j)^2}} - 1}{e^{\frac{1}{|B|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j)^2}-\frac{1}{|B+1|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j+e_ie_j)^2}} - 1}$

We then define:

$\Delta(e|A)=\frac{1}{|A+1|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j+e_ie_j)^2}-\frac{1}{|A|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j)^2}$

$\Delta(e|B)=\frac{1}{|B+1|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j+e_ie_j)^2-\frac{1}{|B|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j)^2}}$

Therefore, the original equation will be: $\frac{\Delta_A}{\Delta_B}=e^{\frac{1}{|B|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j)^2}-\frac{1}{|A|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j)^2}} *\frac{e^{-\Delta(e|A)}-1}{e^{-\Delta(e|B)}-1}$.

To provide a theoretical guarantee, we aim to establish a lower bound for the submodular ratio under any given set where the gain is positive, as most empirical verifications suggest. Since the original formulation is challenging to analyze directly, we introduce two bounds on the function applied to the text files to derive a more concrete lower bound, as follows:

Assumption 1) Assume $|\Delta(e|A)-\Delta(e|B)|\leq \epsilon$, which means the gain difference between any two sets ($A\subset B \subset \Omega$) are bounded.

Assumption 2) Assume average utility is $\mu$, which means $\forall U \in \Omega$, we have: $\frac{1}{|U|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in U}x_i x_j)^2}\leq \mu.$

With Assumption 1, we have: $\frac{\Delta_A}{\Delta_B}\geq e^{\frac{1}{|B|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j)^2}-\frac{1}{|A|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j)^2}} *\frac{e^{-\Delta(e|B)-\epsilon}-1}{e^{-\Delta(e|B)}-1}$.

With Assumption 2, we have:$\frac{1}{|B|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j)^2}-\frac{1}{|A|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in A}x_i x_j)^2}\geq -2\mu$, which means $\frac{\Delta_A}{\Delta_B}\geq e^{-2\mu} *\frac{e^{-\Delta(e|B)-\epsilon}-1}{e^{-\Delta(e|B)}-1}.$We reformulate this as

$\frac{\Delta_A}{\Delta_B}\geq e^{-2\mu} *(e^{-\epsilon}+\frac{e^{-\epsilon}-1}{e^{-\Delta(e|B)}-1}).$

From Assumption 2, we have $\frac{1}{|B|} \sqrt{\sum_{i}^d \sum_{j}^d (\sum_{x\in B}x_i x_j)^2}\leq \mu$. Therefore $e^{-\Delta(e|B)}\leq e^{2\mu}$, and we have: $\frac{\Delta_A}{\Delta_B}\geq e^{-2\mu} *(e^{-\epsilon}+\frac{e^{-\epsilon}-1}{e^{2\mu}-1}).$

Finally, we have a lower bound as $e^{-2\mu}\frac{e^{2\mu-\epsilon}-1}{e^{2\mu}-1}$, which means under defined assumptions, our function is at least $e^{-2\mu}\frac{e^{2\mu-\epsilon}-1}{e^{2\mu}-1}$-weakly submodular function. It can be seen that, the theoretical results are influenced by the gain difference and the average utility. Larger difference or utility indicate a smaller lower bound.

Thanks for the constructive feedback provided by the Reviewer 2LFP. We sincerely appreciate the time and effort you dedicated to evaluating our work. Below, we provide detailed responses to the weaknesses and questions, which we denote as W and Q. Besides, we also added the following discussions into revised submission.

W1 and W2: 1) While the results are promising, it would be beneficial to see how the method scales to larger models commonly used in current LLM research. 2) Whether the proposed method can maintain its efficiency and effectiveness when applied to larger datasets.

Reply to W1 and W2:

To evaluate the effectiveness of our selection method on larger models and datasets, we try our best to 1) pre-train Open-Llama 3B [1] on SlimPajama, and 2) TinyLlama-1B on both SlimPajama and StarcoderData [2]. StarcoderData is a dataset created for code generation, containing approximately 210M files (500GB of data). We assess the pre-trained models on common sense ability (7 tasks), mmlu, bbh and an additional programming task, code_x_glue [3], using harness evaluation. Due to the limited rebuttal time, the cost of pre-training, and the effort involved in selecting data under other baselines, we present the most results we can at this stage. We hope this can address the reviewer’s concerns. Comparison to QuRating in some settings are still ongoing, as it requires running its judger model (QuRater-1.3B) to evaluate all files, and we expect to complete it in the coming days.

1. As shown in the following Table T1, we pre-train open-llama 3B with 10B training budget and 1.5% selection budget with SlimPajama (on 4xA100 devices with about 4 days) . We compare our method with Random, DSIR, QuRating-A, and QuRating-W. DSIR and QuRating-W meet performance degradation due to dimensional collapse, while our method achieves the best performance on all tasks.
2. As shown in the following Table T2, we pre-train TinyLlama-1B on both SlimPajama and StarcoderData with 1.5% selection budget and 10B training budget, compared to random and DSIR. Performance of DSIR and QuRating-W also meet performance degradation in code generation ability. We analyze that DSIR tends to ignore code files in StarcoderData because its selection criterion is based on Wikipedia, leading to dimensional collapse. Besides, our DiSF can mitigate this collapse and achives the best performance on all tasks.

Notably, in Table 4 of our original submission (Table 5 in the revised submission), we also evaluate our method on Pythia-1B and OPT-1.3B in terms of average commonsense reasoning performance, which might also helps illustrate the scalability of our method on current models in the LLM community.

Table T1: Comparison on openllama 3B with 1.5% selecting ratio and 10 B pre-training budget on slimpajama.

Table T2: Comparison on TinyLlama-1B with 1.5% selecting ratio and 10 B pre-training budget on both slimpajama and starcoderdata.

[1] Openllama: An open reproduction of llama

[2] Starcoder: may the source be with you!

[3] CodeXGLUE: A Machine Learning Benchmark Dataset for Code Understanding and Generation