大语言模型持续预训练研究：如何优化预热策略？
大语言模型持续预训练研究：如何优化预热策略？
大语言模型（LLM）的持续预训练是一个重要的研究方向，特别是在新数据不断涌现的今天。来自加拿大蒙特利尔大学的研究团队在论文《Continual Pre-Training of Large Language Models: How to (re)warm your model？》中，深入探讨了不同预热策略对模型性能的影响。本文将详细介绍该研究的实验设置、结果和结论。
研究背景与动机
大语言模型通常需要在数十亿个token上进行预训练。当有新数据可用时，从头开始重新训练整个模型不仅成本高昂，而且效率低下。因此，研究者们提出了持续预训练的概念，即用新数据更新预训练模型，而不是从头开始重新训练。然而，新数据引起的分布变化可能会导致模型在旧数据上的性能下降。为了克服这一挑战，该研究重点研究了不同的预热策略。
实验设置
数据集
- 上游数据集：Pile（Gao，2020），包含300B个token。
- 下游数据集：SlimPajama（Soboleva，2023），包含297B个token。SlimPajama是RedPajama（Together.xyz，2023）的去重版本，基于LLama数据集（Touvron，2023）构建。
模型架构
使用Pythia 410M语言模型架构，该模型在Pile数据集上进行了预训练（Biderman，2023）。未使用flash attention（Dao，2022）。
超参数设置
- 优化器：AdamW，β1 = 0.9，β2 = 0.95，ε = 10−8，权重衰减为0.1。
- 最大学习率：{1.5 · 10−4, 3 · 10−4, 6 · 10−4}。
- 学习率衰减：余弦衰减到最大学习率的10%，在240B个token时达到最小学习率。
- 其他设置：梯度裁剪为1.0，训练以半精度（FP16）进行，无dropout操作。
持续预热实验
1. 预热时长的影响
研究者尝试了不同比例的预热数据量（0%、0.5%、1%和2%），并测量了前50B个token之后的性能。结果显示，预热阶段的长度对Pile和SlimPajama的验证损失没有显著影响。这表明，即使在大规模数据集上，预热阶段的长度对最终性能的影响有限。
要点1：
- 预热阶段的长度似乎对Pile和SlimPajama验证损失没有显著影响。
2. 最大学习率的影响
研究者固定预热阶段的长度为1%，并改变最大学习率（1.5 · 10−4、3 · 10−4和6 · 10−4）。结果显示，较大的最大学习率会提高下游数据的性能，但会损害上游数据的性能；相反，较小的最大学习率会提高上游数据的性能，但限制了对下游数据的适应性。这表明，选择合适的最大学习率是一种平衡上下游性能的有效方法。
要点2：
- 重新预热然后降低学习率似乎是在下游任务上学习良好的必要条件。
- 虽然保持恒定的学习率最初在Pile上有利，但当在SlimPajama上训练足够长时间时，这种优势就会消失。
- 仅在SlimPajama上学习的模型在SlimPajama上的表现比在Pile上预训练的模型更差，突出了两个数据集之间的正向迁移。
3. 与从头开始训练的模型比较
研究者将微调模型与从头开始训练的模型进行了比较。结果显示，所有经过预热的微调模型都比从头开始训练的模型表现更好。这表明，即使下游数据集与上游数据集规模相同且存在重叠，微调也可能提高性能。
4. 在相同数据上重新预热
为了评估分布变化对性能的影响，研究者在Pile数据集上进行了预热实验。结果显示，重新预热学习率会导致性能下降，这表明优化动态性是损失增加的一个重要原因。
要点3：
- 重新调整学习率似乎是导致性能下降的一个重要原因。
- 在同一数据集上进行训练时，模型似乎无法从重新预热学习率导致的性能损失中恢复过来。
5. 评估早期检查点
研究者测试了从不同预训练检查点开始微调的效果。结果显示，选择较早的检查点并不会改善下游性能，最新的检查点仍然是最佳选择。
要点4：
- 在Pile上进行预训练时使用较早的检查点不会导致在SlimPajama上学习得更快。
局部结论
即使下游数据与上游数据的来源相似，在新数据上重新预热预训练模型也是一项艰巨的任务。研究结果表明：
- 用于预热的token数不会显著改变性能。
- 增加最大学习率会提高最终模型的下游性能，而降低最大学习率会提高上游性能。
- 选择较早的检查点会降低上游和下游数据的性能。
这项研究为大语言模型的持续预训练提供了重要的理论和实践指导，特别是在如何平衡新旧数据的性能方面。