Efficient Online Data Mixing For Language Model Pre-Training

📄 arXiv: 2312.02406v2 📥 PDF

作者: Alon Albalak, Liangming Pan, Colin Raffel, William Yang Wang

分类: cs.CL, cs.LG

发布日期: 2023-12-05 (更新: 2023-12-09)

💡 一句话要点

提出在线数据混合(ODM)算法，提升语言模型预训练效率并优化数据配比。

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

关键词: 语言模型预训练 在线学习 数据混合 多臂老虎机 数据选择

📋 核心要点

  1. 现有数据选择方法计算成本高昂，难以适应大型模型和数据集的预训练需求。
  2. 论文提出在线数据混合(ODM)算法，利用多臂老虎机优化训练过程中的数据混合比例。
  3. 实验表明，ODM算法能显著减少训练迭代次数，提高模型在MMLU基准测试上的准确率。

📝 摘要（中文）

大型语言模型的预训练数据对模型下游性能有决定性影响，因此涌现了大量数据选择方法，旨在自动确定最适合预训练的数据。然而，现有数据选择方法速度慢且计算成本高昂，随着模型和预训练数据集规模的增大，这个问题更加严重。数据混合通过将数据点分组并确定整个组的采样概率来降低数据选择的复杂性。但是，数据混合比例通常在训练前固定，因此无法适应不断变化的训练动态。为了解决这些限制，我们开发了一种高效的在线数据混合（ODM）算法，该算法结合了数据选择和数据混合的元素。基于多臂老虎机算法，我们的在线方法在训练期间优化数据混合比例。值得注意的是，我们的方法训练的模型以比次优方法少19%的训练迭代次数达到最终困惑度，并在5-shot MMLU基准测试中将相对准确率提高了1.9%，同时在预训练期间几乎不增加挂钟时间。

🔬 方法详解

问题定义：现有数据选择方法，例如数据选择和数据混合，在应用于大型语言模型预训练时存在效率问题。数据选择方法计算成本高，而传统的数据混合方法无法动态调整数据比例以适应训练过程中的变化。这导致预训练过程耗时且可能无法达到最优性能。

核心思路：论文的核心思路是利用在线学习的方法，在预训练过程中动态地调整不同数据源的混合比例。通过将数据混合问题建模成一个多臂老虎机问题，算法可以根据每个数据源在训练过程中的表现（例如，对损失函数的贡献）来调整其采样概率，从而优化整体训练效率。

技术框架：ODM算法的核心是一个多臂老虎机模型，其中每个“臂”代表一个数据源（或数据组）。算法在每个训练迭代中，根据当前的老虎机策略选择一个数据源进行训练。训练后，根据该数据源的性能（例如，训练损失）更新老虎机策略，从而调整各个数据源的采样概率。整个过程是一个在线学习的过程，数据混合比例随着训练的进行而不断优化。

关键创新：关键创新在于将多臂老虎机算法应用于语言模型预训练的数据混合问题。与传统的数据混合方法相比，ODM算法能够动态地调整数据混合比例，从而更好地适应训练过程中的变化。与数据选择方法相比，ODM算法的计算效率更高，因为它不需要对每个数据点进行评估，而是对整个数据组进行评估。

关键设计：ODM算法的关键设计包括：1) 如何定义每个数据源的“奖励”（reward），通常使用训练损失的负值；2) 如何选择多臂老虎机算法，例如，可以使用UCB（Upper Confidence Bound）或Thompson Sampling等算法；3) 如何设置多臂老虎机算法的参数，例如，探索率（exploration rate）等。这些设计直接影响算法的性能和收敛速度。

📊 实验亮点

实验结果表明，ODM算法能够以更少的训练迭代次数达到与现有最佳方法相同的困惑度（perplexity），减少了19%的训练迭代次数。此外，在5-shot MMLU基准测试中，ODM算法将相对准确率提高了1.9%。这些结果表明，ODM算法能够显著提高语言模型预训练的效率和性能。

🎯 应用场景

该研究成果可广泛应用于大型语言模型的预训练，尤其是在数据资源丰富但质量参差不齐的情况下。通过动态优化数据混合比例，可以提高预训练效率，降低计算成本，并提升模型在下游任务中的性能。该方法还可应用于其他机器学习任务，例如，迁移学习和领域自适应。

📄 摘要（原文）

The data used to pretrain large language models has a decisive impact on a model's downstream performance, which has led to a large body of work on data selection methods that aim to automatically determine the most suitable data to use for pretraining. Existing data selection methods suffer from slow and computationally expensive processes, a problem amplified by the increasing size of models and of pretraining datasets. Data mixing, on the other hand, reduces the complexity of data selection by grouping data points together and determining sampling probabilities across entire groups. However, data mixing proportions are typically fixed before training and therefore cannot adapt to changing training dynamics. To address these limitations, we develop an efficient algorithm for Online Data Mixing (ODM) that combines elements from both data selection and data mixing. Based on multi-armed bandit algorithms, our online approach optimizes the data mixing proportions during training. Remarkably, our method trains a model that reaches the final perplexity of the next best method with 19\% fewer training iterations, and improves performance on the 5-shot MMLU benchmark by 1.9% relative accuracy, while adding negligible wall-clock time during pretraining.