论文

WSqD:大型 模型训练 的无地平线学习率计划

WSqD: A Horizon-Free Learning Rate Schedule for Large Model Training

模型训练预训练

摘要

标准学习率计划(例如余弦退火)与固定的训练范围相关,限制了它们适应事后范围扩展的能力。预热稳定衰减(WSD)通过在短暂的线性冷却之前维持较长的恒定速率阶段来部分解决此问题,从而允许训练从衰减前检查点恢复。然而,其峰值学习率仍然根据原始训练范围进行调整,并且当训练延长时可能会变得次优。受随机凸优化的启发,我们提出了 WSqD(带有平方根基数和线性衰减的预热),这是一种学习率计划,它用移位的逆平方根基数取代了 WSD 的恒定稳定阶段,同时保留了最终的线性冷却时间。在随机凸设置中,WSqD 可证明达到极小极大最优 $O(1/\sqrt{T})$ 最后迭代收敛率。重要的是,其基本学习率计划与时间范围无关,并且仅需要训练时间范围来确定何时开始最终冷却。根据经验,在使用 SlimPajama 语料库进行语言模型预训练时,WSqD 在多个训练范围内匹配或优于精心调整的 WSD 和其他基线,同时重复使用单个峰值学习率。