[论文解读] Reuse, Don't Retrain: A Recipe for Continued Pretraining of Language Models

Jupinder Parmar, Sanjev Satheesh+3 more|arXiv (Cornell University)|Jul 9, 2024

Natural Language Processing Techniques被引用 4

一句话总结

本文提出了一种实用的大语言模型持续预训练方法，强调复用已有模型而非从头开始训练。通过优化数据分布变化和学习率调度——具体而言，采用两阶段数据混合策略，优先侧重知识密集型任务，并结合早期预热的余弦衰减学习率调度——该方法在MMLU上的平均准确率相比基线持续预训练提升了9%，显著提升了模型优化效率，且无需完全重新训练。

ABSTRACT

As language models have scaled both their number of parameters and pretraining dataset sizes, the computational cost for pretraining has become intractable except for the most well-resourced teams. This increasing cost makes it ever more important to be able to reuse a model after it has completed pretraining; allowing for a model's abilities to further improve without needing to train from scratch. In this work, we detail a set of guidelines that cover how to design efficacious data distributions and learning rate schedules for continued pretraining of language models. When applying these findings within a continued pretraining run on top of a well-trained 15B parameter model, we show an improvement of 9\% in average model accuracy compared to the baseline of continued training on the pretraining set. The resulting recipe provides a practical starting point with which to begin developing language models through reuse rather than retraining.

研究动机与目标

• 为应对大语言模型预训练日益增长的计算成本，通过有效复用已训练模型来实现改进。
• 识别能提升模型能力而不引发灾难性遗忘的最优数据分布策略。
• 确定在持续预训练过程中能平衡学习速度与模型稳定性的有效学习率调度方案。
• 提供一种可迁移、实用的配方，通过持续预训练提升通用语言模型的能力。

提出的方法

• 设计两阶段数据分布：第一阶段在原始预训练数据分布上进行持续预训练；第二阶段切换至更强调知识密集型、任务导向型数据（如问答数据集）的分布。
• 实施带有早期预热的余弦学习率衰减调度，以平衡学习幅度与衰减陡度，防止性能下降。
• 在3000亿个持续预训练token处设置数据分布切换点，模型由此从预训练数据过渡到任务导向数据。
• 使用E5-large-v2嵌入和FAISS进行文档挖掘，识别与问答样本最相似的50篇非问答文档，以提升长周期训练中的数据效用。
• 采用在8TB token上预训练的150亿参数模型作为基础模型，以确保研究结果在不同模型规模上的泛化能力。
• 在MMLU和各类任务基准上评估不同持续预训练token数量（1000亿至10000亿token）下的性能表现。

Figure 1: Breakdown of the various distributions considered for the General Blend (GB). We use Upweight Non Web w/ High Quality Web as the GB moving forward given its strong performance across all evaluation areas.

实验结果

研究问题

• RQ1 何种数据分布策略能在不损害现有能力的前提下最大化持续预训练的性能增益？
• RQ2 哪种学习率调度方案能在持续预训练过程中实现最高效且稳定的训练？
• RQ3 在训练过程的哪个时间点切换数据分布可实现最优模型性能？
• RQ4 从现有预训练数据中挖掘文档是否能提升长周期持续预训练中的模型准确率？

主要发现

• 与在原始预训练分布上持续训练的基线相比，所提出的配方在MMLU上的平均准确率提升了9%（从56.1%提升至54.6%），最佳结果达到56.1%。
• 在3000亿个持续预训练token处切换数据分布可获得最高性能，优于更早或更晚的切换时机。
• 采用最大学习率十分之一的预热余弦衰减学习率调度方案，能产生最有效的学习动态。
• 该配方在不同训练周期内均能持续提升性能，在10000亿个持续预训练token时，相比基线模型准确率提升达16%。
• 通过挖掘预训练数据中与问答样本相似的文档，MMLU上的准确率从56.8%提升至57.9%，证明了数据效用的显著增强。
• 该配方在从1000亿到10000亿个持续预训练token的广泛训练规模范围内均表现出强鲁棒性与有效性，显示出优异的可迁移性。

Figure 2: Various distributions of QA data. We use Blend 3.