CLIMB：基于聚类的迭代数据混合引导语言模型预训练

预训练数据集通常是从网络内容中收集的，缺乏内在的领域划分。例如，广泛使用的数据集如Common Crawl并不包含明确的领域标签，而手动整理带有标签的数据集如The Pile则需要大量的人力。因此，尽管在预训练性能方面有显著的好处，但确定最优的预训练数据混合仍然是一个具有挑战性的问题。为了解决这些问题，我们提出了基于聚类的迭代数据混合引导（CLIMB），这是一个自动化的框架，用于在预训练环境中发现、评估和优化数据混合。具体来说，CLIMB将大规模数据集嵌入到语义空间中并进行聚类，然后使用较小的代理模型和预测器迭代地搜索最优混合。当我们的10亿参数模型在持续4000亿个标记的数据混合上进行训练时，其性能超过了当前最先进的Llama-3.2-1B模型2.0%。此外，我们发现针对特定领域（例如社会科学）进行优化可以比随机采样提高5%的性能。最后，我们介绍了ClimbLab，这是一个经过过滤的1.2万亿标记语料库，并包含20个聚类作为研究平台；以及ClimbMix，一个紧凑而强大的4000亿标记数据集，专门设计用于高效预训练，在同等标记预算下提供更优性能。我们对最终的数据混合进行了分析，并阐明了最优数据混合的特点。我们的数据可在以下链接获取：https://research.nvidia.com/labs/lpr/climb/