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

📄 arXiv: 2504.20644v1 📥 PDF

作者: Ziqing Fan, Siyuan Du, Shengchao Hu, Pingjie Wang, Li Shen, Ya Zhang, Dacheng Tao, Yanfeng Wang

分类: cs.LG

发布日期: 2025-04-29

💡 一句话要点

提出DiSF算法，通过多样化文件选择解决LLM预训练数据中的维度坍塌问题。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 预训练数据选择 维度坍塌 数据多样性 贪婪算法 子模优化 训练效率

📋 核心要点

  1. 现有LLM预训练数据选择方法存在维度坍塌问题，导致通用性能下降。
  2. DiSF算法通过选择特征空间中最不相关的文本文件，增强数据多样性，防止维度坍塌。
  3. 实验表明，DiSF算法在TinyLlama模型上显著提高了整体性能，并提升了训练和数据效率。

📝 摘要（中文）

为提升大语言模型(LLM)在有限计算资源下的性能，本文研究了高质量预训练数据的选择问题。现有方法主要依赖代理模型评估样本与目标领域（如BookCorpus和Wikipedia）的相似度。然而，这些基于领域相似性的选择标准存在维度坍塌问题，即提升了领域相关任务的性能，却严重降低了通用性能。为防止坍塌并增强多样性，本文提出了一种多样化文件选择算法(DiSF)，该算法选择特征空间中最不相关的文本文件。通过经典的贪婪算法，DiSF旨在使所选文本的特征协方差矩阵具有更均匀的特征值，并分析了其在$γ$-弱子模优化问题下的近似最优解。在TinyLlama架构上，通过120M到1.1B参数的模型，并在Harness框架的九个任务上进行评估，DiSF显著提高了整体性能。具体而言，DiSF在SlimPajama中节省了98.5%的590M训练文件，在50B的训练预算内优于全数据预训练，实现了约1.5倍的训练效率和5倍的数据效率。

🔬 方法详解

问题定义：现有的大语言模型预训练数据选择方法，例如基于与BookCorpus和Wikipedia等高质量数据源的相似度进行选择，虽然能在特定领域任务上取得较好效果，但会导致特征空间的维度坍塌，从而损害模型的通用性能。现有方法的痛点在于缺乏对数据多样性的有效考虑，导致模型过拟合于特定领域的数据分布。

核心思路：本文的核心思路是通过选择特征空间中尽可能不相关的文本文件来增加预训练数据的多样性。这种方法旨在避免模型过度依赖某些特定的特征维度，从而提高模型的泛化能力。作者认为，通过最大化所选文本特征协方差矩阵的特征值的均匀性，可以有效地防止维度坍塌。

技术框架：DiSF算法采用贪婪算法来选择文本文件。该算法迭代地选择与已选文件相关性最低的文件，直到达到预定的文件数量。算法的核心是计算文本文件在特征空间中的相关性，并选择能够最大程度降低整体相关性的文件。作者将该问题形式化为一个$γ$-弱子模优化问题，并分析了贪婪算法对最优解的近似程度。

关键创新：DiSF算法的关键创新在于其多样性驱动的文件选择策略。与传统的基于相似性的选择方法不同，DiSF算法直接优化所选数据的多样性，从而更有效地防止维度坍塌。此外，作者将该问题形式化为$γ$-弱子模优化问题，为算法的理论分析提供了基础。

关键设计：DiSF算法的关键设计包括：1) 使用预训练的语言模型（例如，BERT）提取文本文件的特征向量；2) 使用余弦相似度或相关系数来衡量文本文件之间的相关性；3) 使用贪婪算法迭代地选择文件，每次选择与已选文件集合相关性最低的文件；4) 可以设置选择的文件数量作为超参数，以控制预训练数据集的大小。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，DiSF算法在TinyLlama模型上取得了显著的性能提升。在50B的训练预算内，DiSF算法仅使用SlimPajama数据集中1.5%的文件（节省了98.5%的训练文件），就超过了使用整个数据集进行预训练的性能。具体而言，DiSF算法实现了约1.5倍的训练效率和5倍的数据效率。在Harness框架的九个任务上，DiSF算法也表现出优于其他基线方法的整体性能。

🎯 应用场景

DiSF算法可应用于各种大语言模型的预训练数据选择，尤其是在计算资源有限的情况下，可以显著提高模型的训练效率和性能。该方法可以帮助研究人员和开发者构建更通用、更强大的语言模型，并应用于自然语言处理的各个领域，如文本生成、机器翻译、问答系统等。此外，该方法也适用于其他机器学习任务的数据选择，例如图像分类和语音识别。

📄 摘要（原文）

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.