Domain reweighting is an emerging research area aimed at adjusting the relative weights of different data sources to improve the effectiveness and efficiency of LLM pre-training. We show that data mixtures that perform well at smaller scales may not retain their advantage at larger scales, challenging the existing practice of determining competitive mixtures in small-scale experiments and directly applying them at much larger scales. To address this, we propose AutoScale, a two-stage, scale-aware data composition framework. First, AutoScale fits a parametric model that predicts the model's loss under different data compositions, then uses it to find an approximate best allocation at smaller, more manageable budgets. Next, leveraging a novel theoretical analysis of how optimal compositions evolve with scale, AutoScale extrapolates that composition to larger budgets without further retraining. Empirically, AutoScale accelerates convergence and improves downstream performance. For instance, when pre-training GPT-2 Large, it achieves a 28% faster perplexity reduction than baselines and up to a 38% speed-up over unweighted training, while yielding best-average results on various downstream tasks. Overall, our findings illustrate how domain importance shifts with training scale, underscoring the need for scale-dependent data curation in LLM training. Our code is open-sourced.

翻译：领域重加权是一个新兴研究方向，旨在调整不同数据源的相对权重以提升大语言模型预训练的效果与效率。我们发现，在小规模训练中表现优异的数据混合方案，在更大规模下可能无法保持其优势，这对当前在小规模实验中确定优势混合方案并直接应用于更大规模训练的既有实践提出了挑战。为解决此问题，我们提出了AutoScale——一个两阶段的规模感知数据组合框架。首先，AutoScale拟合一个参数化模型，用于预测不同数据组合下模型的损失，并利用该模型在较小且更易管理的计算预算下寻找近似最优的数据分配方案。随后，基于对最优组合如何随规模演变的新颖理论分析，AutoScale将该组合方案外推至更大计算预算，而无需重新训练。实证结果表明，AutoScale能加速收敛并提升下游任务性能。例如，在预训练GPT-2 Large时，其困惑度下降速度比基线方法快28%，相较于未加权训练最高可提速38%，同时在多种下游任务上取得最佳平均性能。总体而言，我们的研究揭示了领域重要性如何随训练规模变化，强调了大语言模型训练中需要依赖规模的数据策展。代码已开源。