Optimizing Pretraining Data Mixtures with LLM-Estimated Utility
作者: William Held, Bhargavi Paranjape, Punit Singh Koura, Mike Lewis, Frank Zhang, Todor Mihaylov
分类: cs.CL, cs.AI
发布日期: 2025-01-20 (更新: 2025-01-23)
备注: 10 pages, 8 figures
💡 一句话要点
提出UtiliMax和MEDU，高效优化LLM预训练数据混合比例，加速训练并降低计算成本。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 大型语言模型 预训练数据 数据混合 效用估计 自动化机器学习
📋 核心要点
- 现有方法在平衡预训练数据的质量、数量和多样性方面存在不足，尤其是在计算资源有限的情况下。
- 论文提出UtiliMax和MEDU两种方法，分别通过缩减规模消融和LLM评估来估计数据效用，优化数据混合比例。
- 实验表明，UtiliMax相比手动基线加速高达10.6倍，MEDU在计算成本降低200倍的情况下，性能与消融实验相当。
📝 摘要（中文）
大型语言模型（LLM）的性能随着高质量训练数据的增加而提升。然而，利用更大的数据集需要在来源之间平衡质量、数量和多样性。在计算和数据受限的情况下评估了九种基线方法后，我们发现token计数启发式方法优于手动和学习混合方法，表明考虑数据集大小和多样性的简单方法非常有效。基于这一洞察，我们提出了两种互补的方法：UtiliMax，它通过结合缩减规模消融实验的效用估计来扩展基于token的启发式方法，与手动基线相比，实现了高达10.6倍的加速；以及模型估计数据效用（MEDU），它利用LLM从小型样本中估计数据效用，在将计算需求降低约200倍的同时，匹配了基于消融实验的性能。总之，这些方法建立了一个新的框架，用于自动、计算高效的数据混合，该框架在各种训练方案中都具有鲁棒性。
🔬 方法详解
问题定义：现有的大型语言模型预训练方法在数据混合方面面临挑战。简单地增加数据量并不一定能提升模型性能，因为不同来源的数据质量和多样性存在差异。手动调整数据混合比例耗时且低效，而现有的自动化方法，如基于token计数的启发式方法，虽然简单有效，但忽略了数据的内在价值（utility）。因此，如何高效地确定最优的数据混合比例，以在有限的计算资源下最大化模型性能，是本文要解决的核心问题。
核心思路：本文的核心思路是利用模型本身来估计数据的效用（utility），并以此为依据来优化数据混合比例。具体来说，论文提出了两种互补的方法：UtiliMax和MEDU。UtiliMax通过缩减规模的消融实验来估计数据效用，而MEDU则直接利用LLM来评估数据样本的质量和价值。这两种方法都旨在克服传统方法中数据效用评估的难题，从而实现更高效的数据混合。
技术框架：整体框架包含两个主要方法：UtiliMax和MEDU。UtiliMax首先使用缩减规模的数据集进行消融实验，评估每个数据源对模型性能的影响，从而估计其效用。然后，它使用这些效用估计来调整数据混合比例，以最大化整体训练效果。MEDU则直接利用一个预训练的LLM来评估数据样本的质量和价值，并基于这些评估结果来确定数据混合比例。两种方法都旨在自动化数据混合过程，并减少人工干预。
关键创新：本文的关键创新在于提出了利用LLM来估计数据效用的思想。与传统的基于token计数的启发式方法相比，这种方法能够更准确地反映数据的内在价值，从而实现更优的数据混合比例。此外，MEDU方法通过直接利用LLM进行数据评估，显著降低了计算成本，使其在资源有限的情况下也能够有效应用。
关键设计：UtiliMax的关键设计在于如何有效地进行缩减规模的消融实验，以获得可靠的效用估计。MEDU的关键设计在于如何设计合适的prompt，使得LLM能够准确地评估数据样本的质量和价值。此外，两种方法都需要设计合适的优化算法，以基于效用估计来确定最优的数据混合比例。具体的参数设置和损失函数选择可能取决于具体的应用场景和数据集。
🖼️ 关键图片
📊 实验亮点
实验结果表明，UtiliMax方法相比手动基线方法，在模型训练速度上实现了高达10.6倍的加速。MEDU方法在将计算需求降低约200倍的同时，性能与基于消融实验的方法相当。这些结果表明，本文提出的方法在计算效率和模型性能方面都具有显著优势。
🎯 应用场景
该研究成果可广泛应用于大型语言模型的预训练阶段，尤其是在数据资源丰富但计算资源有限的情况下。通过自动优化数据混合比例，可以显著提升模型性能，降低训练成本，加速模型开发周期。此外，该方法还可以应用于其他机器学习任务，例如图像分类和自然语言处理等。
📄 摘要（原文）
Large Language Models improve with increasing amounts of high-quality training data. However, leveraging larger datasets requires balancing quality, quantity, and diversity across sources. After evaluating nine baseline methods under both compute- and data-constrained scenarios, we find token-count heuristics outperform manual and learned mixes, indicating that simple approaches accounting for dataset size and diversity are surprisingly effective. Building on this insight, we propose two complementary approaches: UtiliMax, which extends token-based heuristics by incorporating utility estimates from reduced-scale ablations, achieving up to a 10.6x speedup over manual baselines; and Model Estimated Data Utility (MEDU), which leverages LLMs to estimate data utility from small samples, matching ablation-based performance while reducing computational requirements by $\sim$200x. Together, these approaches establish a new framework for automated, compute-efficient data mixing that is robust across training regimes.