LLM训练数据配比策略的优劣，竟由学习率调度决定

（1/9）假设你在有限的数学数据上用网络数据训练大语言模型。你应该把所有数学数据放在训练末尾，还是更早地与网络数据混合？答案会随学习率（LR）调度而改变！在固定计算量和中期训练（MT）token 预算的情况下，当学习率在整个训练过程中衰减时，更早混合能带来更低的中期训练损失。在恒定学习率下，把所有中期训练数据留到最后效果更好。

我们随后改变了学习率衰减开始的时间。测试出的最佳策略是：在中期训练数据加入后保持高学习率，然后在接近结束时衰减。详见我们的新论文：《数据调度与学习率调度应协同设计》！

（2/9）中期训练（MT）是在预训练（PT）之后，在规模更小、价值更高的数据集（如数学或代码）上继续训练。我们应该把所有中期训练数据留到最后，还是更早引入并与预训练数据混合？我们在固定总计算量和中期训练 token 预算的情况下比较了这些选择。通过对密集模型和混合专家模型（MoE）进行 250 多次运行，我们研究了答案如何依赖于学习率调度。

（3/9）这里有一个具体例子。当学习率在整个持续训练过程中衰减时，更早引入中期训练数据并与预训练数据混合，比把所有中期训练数据留到最后的纯中期训练阶段能带来更低的最终中期训练损失。在恒定学习率下，结果相反：最后的纯中期训练阶段效果更好。数据调度没有改变。仅仅改变学习率就改变了哪个方案更优。

（4/9）因此我们希望这两种调度能协同工作：在相同计算量和中期训练 token 数下，实现更低的最终中期训练损失。我们保持每种数据调度不变，测试了在首次引入中期训练数据之前、之时或之后开始学习率衰减的情况。在所有 19 次扫描中，测试出的最佳数据与学习率调度组合都遵循相同的模式：开始使用中期训练数据，在一段时间内保持高学习率，然后在接近结束时衰减。衰减本身仅覆盖从首次加入中期训练数据到训练结束这段时间的最后 10%–40%。

（5/9）我们的理论模型有助于解释这两个发现，首先是为什么偏好的数据调度依赖于学习率。在固定中期训练预算下，我们的二次随机梯度下降（SGD）模型在恒定学习率下有一个简单的最优解：把所有中期训练数据放在最后。在足够强的线性学习率衰减和额外假设下，最优解则分为三个阶段：仅预训练 → 中期训练数据占比逐渐增加的混合阶段 → 仅中期训练。

（6/9）同一个模型也解释了在中期训练最初几次更新中保持高学习率的好处。对于从仅预训练到仅中期训练的固定切换，在切换后保持高学习率有助于中期训练更新去降低预训练更新无法降低的误差。在接近结束时降低学习率可以限制新梯度噪声的影响，因为后期更新几乎没有时间去降低这些噪声。在我们的假设下，当中期训练预算较小且训练时间足够长时，最优的线性衰减在切换之后、训练结束之前开始。

（7/9）附带说明：渐进式数据混合也能减少混合专家模型（MoE）负载不平衡的峰值。混合专家模型的路由器将 token 分配给不同的专家。在我们的实验中，中期训练数据占比的突然跳升会导致负载不平衡出现急剧峰值。而渐进式增加则不会出现类似的峰值。我们还看到，在渐进式混合运行中，专家负载排名的重新洗牌更少：在数据转变前接收更多中期训练 token 的专家，在转变后往往仍保持较高的使用率。

（8/9）这是我们在训练调度系列工作中的一次新探索！我们最初提出了一个将学习率调度映射到损失曲线的缩放定律，随后研究了数据复用如何在多个 epoch 中改变缩放规律，以及学习率衰减如何削弱数据课程的优势。现在，我们发现仅仅改变学习率调度就能翻转哪种数据混合调度能带来更低的损失！

（9/9）要点：在比较数据调度时，请检查结论在改变学习率调度后是否依然成立。在调整学习率时，请考虑新数据何时进入。与 Kairong Luo（@openhonor）和 Tingkai Yan 的合作非常愉快！非常感谢我的导师 Kaifeng Lyu（@vfleaking）和 Wenguang Chen 教授的指导。📄 Arxiv 链接：随时欢迎讨论训练调度及其背后的理论！

好问题！我们即使在 630M 的密集模型上也观察到了这种反转，所以它似乎并不需要非常大的模型。然而，我们尚未测试窄任务微调。我们的设置涉及通用数据与专业数据的混合，看看同样的现象是否在微调中成立将会很有趣。