理解温度-稳定-衰减学习率：河谷损失景观视角

当前训练语言模型需要预先确定固定的计算预算，因为典型的余弦学习率计划依赖于总步数。在 contrast，Warmup-Stable-Decay（WSD）计划使用恒定学习率生成一系列主迭代，理论上可以无限延续，而无需预先指定计算预算。然后，对于任何计算预算，都可以在恰当时间从主分支分支出来，并使用快速衰减学习率，以获得强大的模型。经验上，WSD 生成非传统的损失曲线：在稳定阶段损失保持在较高水平，但在衰减阶段损失急剧下降。为了解释这一现象，我们推测预训练损失呈河谷景观，这类似于深谷底部有河流的景观。在此假设下，我们展示在稳定阶段，迭代因高学习率而产生大幅振荡，却沿着河流迅速前进。在衰减阶段，学习率的快速下降最小化迭代的振荡，使其靠近河流，显示出真实的优化进展。因此，持续的高学习率阶段和快速衰减阶段分别负责河流和山间方向的进展，二者都至关重要。我们的分析预测了与实证观察一致的现象，并显示这种景观可从对简单二元语法数据集的预训练中产生。鼓励理论应用，我们引入 WSD-S，即 WSD 的一种变体，复用先前检查点的衰减阶段，仅保留一个主分支，其中我们从衰减后的检查点恢复。在参数从 0.1B 到 1.2B 的 various compute budgets 下，WSD-S 在获取多个语言模型检查点方面经验上超过 WSD 和 Cyclic-Cosine。