使用 LLM 估算的效用优化预训练数据混合

图 1:我们提出的数据混合方法与流行基线的扩展比较。
预训练数据混合为何重要?
训练大型语言模型(LLM)需要海量数据集。这通常意味着需要结合不同的来源——例如网络数据、学术论文和编程代码。然而,并非所有数据集对模型性能的贡献都相同。挑战在于:在计算限制下,**我们如何决定用于训练的最佳数据混合?**
以前的工作通过**手动筛选、基于启发式的采样或学习型数据混合模型**来探索数据混合。在我们广泛的基线评估中,我们发现最简单的基于启发式的采样方法,尽管假设所有数据源都同等有用,但其性能出人意料地优于以往的工作。在这项工作中,我们引入了 **UtiliMax** 和 **模型估算数据效用(MEDU)**——这些方法能够自动估算数据效用,并利用这些信息计算出改进的数据混合。
问题:如何选择最佳数据混合?
给定不同的数据集,一个关键挑战是**如何在每个数据集之间分配训练资源**,以最大化模型性能。例如,在这项工作中,我们使用 Dolma V1.7,这是用于训练 OLMo 的数据集,它由来自不同领域的 15 个不同来源组成。
在解决数据混合问题时,之前的工作探索了:
- 人工筛选——专家决定每个来源要包含多少数据。
- 基于启发式的采样——例如按比例分配 token 的方法。
- 学习型数据混合模型——在训练过程中动态调整数据分布。
然而,以前的方法尚未在受控设置中进行比较,因此它们之间的比较尚不清楚!为了解决这个问题,我们首先进行了广泛的基线比较。我们发现,最简单的方法,即仅根据数据集大小推导出的 UniMax,出人意料地优于所有其他现有方法!
图 2:基线数据混合方法的比较。UniMax 在计算受限和数据受限设置中都优于其他基线。
值得注意的是,UniMax 不区分不同数据集的质量或领域。在我们工作的其余部分,我们探讨了将这些因素(概括为数据集效用)纳入考虑是否会带来更强的结果。
我们的方法:UtiliMax 和 MEDU
UtiliMax:平衡效用、多样性和规模
UtiliMax 通过将**从小规模实验中得到的单个数据集效用估算**作为数据效用估算值,扩展了基于启发式的数据混合方法。然后,我们将寻找最佳数据混合的问题框定为投资组合优化问题。在金融投资组合中,资产的效用是其预期回报,而投资组合的风险是其多样化程度的函数!利用为投资组合优化设计的凸优化工具,我们推导出了一个方法,该方法在**采样所有数据集**、**偏好高效用数据集**以及**避免过度重复小型数据集**之间取得平衡。
我们发现,这种公式化方法始终优于其他优化程序,例如贪婪采样高效用数据或 UniMax。
图 3:UtiliMax 与替代优化程序的比较。UtiliMax 始终优于其他方法。
MEDU:基于 LLM 的效用估计
虽然 UtiliMax 提高了效率,但对每个数据集进行消融研究的计算成本很高。MEDU 利用现有的 LLM 来消除此成本,通过估算训练数据的有用性而无需额外训练运行,**从而使计算成本比基于消融的方法降低约 200 倍**。
图 4:MEDU 与直接从消融实验中得出的成本更高的数据混合进行比较。
MEDU 首先使用 LLM 根据基准问题描述领域所需的高级技能和知识。然后,它使用此描述将来自单个数据集的文档分类到效用类别(优秀、良好、一般、差、无用)中。通过少量样本,这使我们能够估算单个数据集的效用,而无需训练新模型!
主要发现
1. 简单启发式方法通常优于复杂方法
UniMax 仅平衡数据多样性和重复约束,其性能优于许多手动和学习型数据混合模型。这表明许多更复杂的数据混合方法未能捕捉到真实的训练动态!
2. UtiliMax 显著节省计算成本
通过使用小规模的效用估算,UtiliMax 能够进行数据分配,从而以**更少的 FLOPs 获得更好的模型**。
3. LLM 可以有效估算数据效用
MEDU 取代了成本高昂的消融研究,同时实现了可比的性能,使得数据选择**更快、更经济**。
4. 多样性和规模对泛化能力很重要
优先考虑数据集多样性和大小的混合方法比仅关注效用分数的方法能带来更好的结果。
启示与未来工作
这些发现为**自动化、计算高效的数据混合**奠定了基础,该方法能够适应计算和数据受限的训练设置,并可在训练前进行计算。
我们预计未来的研究将揭示**高质量数据的新信号**——例如 Thrush 等人提出的开源模型之间的损失相关性。UtiliMax 优化程序是一种原则性的方法,可以将这些信号纳入其中,以获得更好的结果。
欲了解更多详情,请查看我们的完整研究论文此处。