QuaDMix：为高效大语言模型预训练而生的质量-多样性平衡数据选择方法

这项研究出自字节跳动的刘丰泽、周卫东、刘斌斌、于知淼、张一帆、林浩彬、余一峰、周晓欢、王泰峰和曹勇等研究人员之手，发表于2025年4月23日的arXiv预印本平台（arXiv:2504.16511v1）。这篇论文探讨了如何优化大语言模型（LLM）预训练数据的选择策略，特别注重质量与多样性之间的平衡。

一、研究背景：大模型训练数据的双重挑战

想象一下你在制作一道美食。你需要什么？当然是优质的食材，但同时也需要各种不同类型的食材来确保菜肴丰富多样。大语言模型的训练也面临着类似的挑战。

研究人员发现，现有的大语言模型训练方法往往存在一个根本性问题：它们要么专注于提高数据质量，要么致力于增加数据多样性，却很少同时考虑这两个因素之间的复杂互动关系。就像做菜时如果只关注食材的新鲜度而忽略了种类的丰富性，最终的菜肴可能会单调乏味；反之，如果只追求食材种类的多样，却不在意质量，那么即使种类再多，做出的菜肴也不会可口。

研究团队指出，在实际训练大语言模型时，高质量数据资源是有限的，这就导致了一个不可避免的权衡问题：是优先选择更高质量的数据，还是优先保证数据的多样性？正是这个问题促使研究团队开发了一种新的方法，能够在固定的训练配额下，自动权衡数据质量和多样性，找到最优的数据分布。

二、传统方法的局限性

在深入了解QuaDMix之前，让我们先看看现有方法存在的问题。传统的数据选择策略主要有两种做法：一种是先进行质量过滤，然后再调整不同领域数据的比例；另一种是直接优化数据混合配比，追求模型性能最大化。

这些方法就像是在烹饪过程中，要么先挑选最新鲜的食材，再决定每种食材放多少；要么直接根据经验配比各种食材，不太关心每种食材本身的新鲜程度。这两种方法都有其局限性。

首先，质量评判标准的定义本身就很模糊。有些研究者用正则表达式匹配来筛选数据，有些用教育价值来评判，还有些研究者将与指令微调数据的相似度作为质量指标。这就像不同的厨师对"新鲜"的定义各不相同：有人看保质期，有人看外观，有人则看气味。

其次，质量标准的选择会直接影响最终选出的数据分布。正如论文图1所示，当使用Fineweb-edu分类器选择排名前5%的数据时，与原始数据相比，健康、就业和教育领域的数据比例大幅增加。这就像是我们用"外观"来筛选食材的话，可能会偏爱某些特定种类的蔬菜，而忽略了其他虽然外观不佳但营养丰富的食材。

最关键的是，不同领域的数据质量水平差异很大，而高质量数据又是有限的。这就使得在质量和多样性之间进行权衡变得尤为重要，而现有方法很难同时优化这两个维度。

三、QuaDMix：兼顾质量与多样性的统一框架

研究团队提出的QuaDMix框架就像是一位懂得平衡的大厨，能够根据食材的新鲜度和种类多样性，科学地调配出最佳的烹饪配方。这个框架主要包含三个步骤：特征提取、质量排序和数据采样。

首先是特征提取。QuaDMix使用多种质量评分器和领域分类器为训练语料库中的每个文档打分并分类。这就像是大厨对仓库里的各种食材进行全面检查，不仅看它们的新鲜度，还看它们属于哪种食材类别。

接着是质量排序。QuaDMix采用一个加权平均的方式合并不同质量评分器的分数，这些权重是可调整的参数。然后，系统会根据合并后的质量分数，为每个领域内的数据进行排序。想象一下，厨师不仅根据食材的外观、气味、保质期等多重标准评估它们的新鲜度，还会将同一类食材（如肉类、蔬菜、水果等）分开排序，因为不同类别的食材之间很难直接比较新鲜度。

最后是数据采样。QuaDMix使用一个参数化的采样函数，基于数据的质量和领域标签来决定每个文档的采样频率。研究团队假设，质量较高的数据应该被更频繁地采样，而采样函数的参数决定了随着质量下降，采样频率如何变化。不同领域有独立的参数控制，这就使得系统能够灵活地调整不同领域数据的比例，从而控制整体的多样性。

为了探索这个庞大的参数空间，研究团队采用了一个两步法：首先，他们在小模型上进行了大量实验，尝试不同的参数配置；然后，他们训练了一个回归模型，用来预测不同参数配置下大模型的表现。这种方法大大提高了寻找最优参数的效率，避免了直接在大模型上进行大量实验的高昂成本。

四、实验设计：小模型模拟与参数优化

如何找到最优的QuaDMix参数？如果直接在大模型上尝试各种参数组合，计算成本将高得惊人。研究团队的解决方案充满智慧：他们使用小规模模型作为"试验田"，通过观察小模型在不同参数下的表现，来预测大模型的性能。

具体来说，研究团队使用了RefinedWeb数据集，选择了三种质量评估工具：AskLLM、Fineweb-Edu和DCLM，以及一个将数据分为26个不同领域的分类器。他们训练了3000个只有100万参数的小模型，每个模型使用不同的QuaDMix参数从数据集中采样10亿个标记进行训练。

这就像是厨师在正式烹饪大餐之前，先用小份量的食材进行多次试验，尝试各种不同的配方和烹饪方法，从中找出最佳组合。这些小规模实验只需要每个1个NVIDIA H100 GPU小时，大大降低了探索成本。

训练完这些小模型后，研究团队使用LightGBM回归器来拟合参数与模型性能之间的关系。这个回归模型就像是一本详细的烹饪笔记，记录了各种配方组合与最终菜肴口感之间的关系，厨师可以根据这本笔记快速找出最佳配方，而不需要反复试验。

通过这种方法，研究团队能够在庞大的参数空间中高效地搜索最优参数配置，为大规模模型训练提供指导。

五、实验结果：多方位的性能提升

QuaDMix的表现如何？研究结果令人印象深刻。研究团队使用找到的最优参数在530M参数的模型上进行了训练，并与多种现有方法进行了对比。

如表1所示，与仅关注数据质量或仅关注数据混合的方法相比，QuaDMix在所有基准测试上都表现出色。具体来说，当使用OpenHermes作为验证集时（QuaDMix-OH），QuaDMix在阅读理解、常识推理、知识密集型任务和数学问题上的平均性能提升了7.2%。更令人兴奋的是，当使用下游任务的训练集作为验证集时（QuaDMix-BMK），性能进一步提升，显示出QuaDMix针对特定任务优化的能力。

研究团队还分析了最优参数配置的特点。如图4所示，在QuaDMix-BMK的最优配置中，健康和科学领域的数据被大幅提升采样，而体育和计算机领域的数据则被降低采样。这表明，不同领域的数据对模型性能的贡献各不相同，QuaDMix能够自动找到最佳的数据分布。

此外，研究团队发现，在三种质量过滤器中，DCLM对合并质量得分的贡献最大，而AskLLM的贡献相对较小。这种自动权衡不同质量标准的能力，是QuaDMix区别于其他方法的重要特点。

研究团队还进行了深入的消融实验，以验证质量得分合并和数据量选择的重要性。结果表明，合并所有三种质量过滤器的策略表现最佳，尽管单独使用某一种过滤器可能在特定任务上表现较好（如DCLM在数学任务上）。同时，他们发现选择300亿标记（即排名前5%的高质量数据）比选择更多标记但质量较低的数据更有效，这表明在数据量和质量之间，质量往往更为重要。

六、小模型的预测能力和未来探索方向

QuaDMix的一个关键假设是：小模型的表现可以预测大模型的表现。这一假设得到了实验验证。研究团队训练了5个单独的回归模型，每个模型使用一个基准测试的训练集损失作为目标。结果表明，除了HellaSwag任务外，QuaDMix-BMK在所有任务上的预测性能都优于QuaDMix-OH，这与大模型的实际表现一致。

HellaSwag任务上的不一致性可能是因为这个任务的预测损失方差较大，使得代理能力较低。这提示我们，如何进一步提高小模型的代理能力是未来值得探索的方向。

研究团队坦诚地指出了QuaDMix框架的几点局限性。首先，QuaDMix参数空间的设计仍有改进空间，目前的采样函数参数可能在不同参数下生成相似的函数，导致冗余并影响回归模型的确定性。其次，在高维参数空间中寻找最优参数的效率仍然较低，目前采用的随机猜测方法可能只能找到局部最优解，如何更有效地搜索参数空间仍是一个开放问题。

尽管如此，QuaDMix为大语言模型的预训练数据选择提供了一个实用的解决方案，能够有效平衡数据质量和多样性，从而提高模型性能。

七、结论与启示

回顾整个研究，QuaDMix的创新之处在于将数据质量和多样性纳入一个统一的框架中考虑，自动优化数据分布，解决了传统方法各自为营的局限性。这项研究不仅改进了模型性能，还揭示了一些有价值的见解：

首先，不同的质量标准在不同下游任务上有各自的优势和局限，但适当合并这些标准可以在各个任务上都取得一致的改进，这说明不同质量标准提供了互补的信息。

其次，最优的数据混合配比随质量标准的变化而变化，这表明同时优化质量和多样性的重要性。如果我们只固定其中一个方面，很难达到最佳性能。

第三，回归模型的目标可以指导对特定下游任务的偏好，这使得针对特定应用场景优化数据选择成为可能。

归根结底，QuaDMix提供了一种实用的方法来解决大语言模型预训练中的数据选择问题，它不仅提高了模型性能，还为我们理解数据质量和多样性之间的平衡提供了新的视角。对于普通人来说，这意味着未来的AI助手可能会变得更加全面和高效，能够更好地理解和回应我们的需求。

感兴趣的读者可以通过arXiv:2504.16511v1访问完整论文，深入了解QuaDMix的技术细节和实验结果。随着大语言模型技术的不断发展，如何更有效地利用训练数据将继续是一个重要的研究方向，QuaDMix在这个方向上迈出了重要的一步。