Power Scheduler：一种与批量大小和标记数无关的学习率调度器
计算与语言
2024-09-13 v2 人工智能
机器学习
摘要
为语言模型预训练寻找最佳学习率是一个具有挑战性的任务。这不仅是因为学习率、批量大小、训练标记数、模型规模以及其他超参数之间存在复杂关联，而且因为对拥有数十亿或数万亿参数的大型语言模型进行超参数搜索在成本上是不可行的。最近的研究提出使用小型代理模型和小型语料库进行超参数搜索，并将最优参数迁移到大型模型和大型语料库上。虽然对于与模型规模相关的超参数（如深度和宽度）已在理论和实证上证明了零样本可转移性，但从小型语料库到大型语料库的零样本迁移仍未得到充分探索。本文研究了最佳学习率、批量大小和训练标记数之间的相关性，针对最近提出的 WSD 调度器进行了研究。在数千次小规模实验后，我们发现变量之间存在幂律关系，并展示了其在不同模型规模之间的可转移性。基于此，我们提出一种新的学习率调度器 Power 调度器，能够忽略训练标记数和批量大小的影响。实验表明， 将 Power 调度器与 Maximum Update Parameterization (muP) 结合使用，可以在一种超参数设置下，无论训练标记数、批量大小、模型规模甚至模型架构如何，都能始终实现出色的性能。我们在 3B 参数的密集模型和混合专家 (MoE) 模型上使用 Power 调度器训练所得模型，性能可与最先进的小型语言模型相媲美。我们已在 https://ibm.biz/BdKhLa 开源这些预训练模型。
引用
@article{arxiv.2408.13359,
title = {Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler},
author = {Yikang Shen and Matthew Stallone and Mayank Mishra and Gaoyuan Zhang and Shawn Tan and Aditya Prasad and Adriana Meza Soria and David D. Cox and Rameswar Panda},
journal= {arXiv preprint arXiv:2408.13359},
year = {2024}
}