Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler

寻找语言模型预训练的最佳学习率是一项具有挑战性的任务。这不仅因为学习率、批量大小、训练令牌数量、模型大小和其他超参数之间存在复杂的相关性，而且因为对于具有数十亿或数万亿参数的大型语言模型进行超参数搜索是代价高昂的。最近的研究提出使用小型代理模型和小型语料库来执行超参数搜索，并将最佳参数转换到大型模型和大型语料库中。虽然零-shot可迁移性在理论上和经验证明适用于与模型大小相关的超参数，例如深度和宽度，但从小型语料库到大型语料库的零-shot迁移还未被深入探讨。本文研究了最近提出的WSD调度器的最佳学习率、批量大小和训练令牌数量之间的相关性。经过成千上万次的小型实验，我们发现了变量之间的幂律关系，并证明了其在模型大小方面的可迁移性。基于这一观察，我们提出了一种新的学习率调度器，Power调度器，它对训练令牌数量和批量大小是不可知的。实验证明，将Power调度器与最大更新参数化（muP）相结合，可以在不考虑训练令牌数量、批量大小、模型大小甚至模型架构的情况下，始终实现令人印象深刻的性能。我们使用Power调度器训练的3B密集型和MoE模型达到了与小型语言模型的最新技术水平相当的性能。我们在https://ibm.biz/BdKhLa上开源了这些预训练模型。

本文旨在解决语言模型预训练中学习率的优化问题，特别是针对大规模语言模型的优化问题。同时，本文还探讨了小规模数据集与大规模数据集之间的学习率转移问题。

本文提出了一种基于幂律关系的学习率调度器——Power scheduler，它可以在不考虑训练令牌数量和批量大小的情况下，实现在各种规模的模型上的良好表现。同时，本文还提出了一种最大更新参数化（muP）的方法，将Power scheduler与muP相结合可以在不同规模、不同架构的模型上实现一组超参数的一致性表现。

本文通过大量小规模实验，发现了学习率、批量大小和训练令牌数量之间的幂律关系，并证明了该关系可以在不同规模的模型上进行转移。同时，本文还开源了使用Power scheduler和muP训练的3B dense和MoE模型，并取得了与小型语言模型相当的表现。

在相关研究方面，最近的一些研究集中在语言模型的训练技巧上，如探索更好的优化器、更好的初始化方法等。例如，GPT-3使用的AdaScale优化器就是一种针对大规模语言模型的优化器。