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

TL;DR

The paper addresses dimensional collapse in LLM pretraining data caused by domain-focused selection under budget constraints. It proposes DiSF, a greedy, batch-based method that decorrelates text features by minimizing the covariance Frobenius norm, analyzed within a $\gamma$-weakly submodular optimization framework to guarantee a $(1-e^{-\gamma})$-approximation. Empirically, DiSF on TinyLlama models across $\sim$120M to $1.1$B parameters achieves higher Harness task performance while delivering about $1.5\times$ training efficiency and $5\times$ data efficiency with a $1.5\%$ file budget, outperforming full-data pretraining under a $50$B-token budget. The work highlights the importance of data diversity in budget-constrained pretraining and offers a scalable, practical approach for diversifying LLM training data.

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

Figures (9)

• Figure 1: The t-SNE tsne visualization of text features (normalized to the unit sphere) selected by different methods on SlimPajama. We use Contriever contriever 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 a long narrow band, indicating dimensional collapse. (d) and (e) represent D4 and our DiSF. For D4, we display 500 random samples after reducing redundancy, while for DiSF, we select samples with the highest values based on equation \ref{['eq:proxy']}. Both methods, especially DiSF, show more uniformly scattered features, indicating improved diversity.
• Figure 2: Commonsense reasoning ability of pre-trained TinyLlama 1B using various selection methods evaluated 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.
• Figure 3: The dominance score for recent methods calculated as $\frac{\sum_{i=1}^k\lambda_i}{\sum_{j=1}^d\lambda_j}$, where $\lambda_i$ represents the $i$-th largest eigenvalue of the feature covariance matrix, and $d$ is the dimension of the feature space. We use Contriever model to extract features. We select the top 500 text samples based on their respective selection criteria. For D4, we select 500 random samples after reducing redundancy.
• Figure 4: Proxy value, as defined in equation \ref{['eq:proxy']}, calculated on different methods. 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 selection method achieves significantly larger proxy values, indicating much better uniformity of feature dimensions.
• 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-training on SlimPajama.

Theorems & Definitions (3)

• Lemma 1
• Lemma 2
• proof