Power Scheduler: A Batch Size and Token Number Agnostic Learning Rate Scheduler
作者: Yikang Shen, Matthew Stallone, Mayank Mishra, Gaoyuan Zhang, Shawn Tan, Aditya Prasad, Adriana Meza Soria, David D. Cox, Rameswar Panda
分类: cs.CL, cs.AI, cs.LG
发布日期: 2024-08-23 (更新: 2024-09-11)
💡 一句话要点
提出Power Scheduler，一种与Batch Size和Token数量无关的学习率调度器
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 学习率调度器 语言模型预训练 超参数优化 幂律关系 零样本迁移
📋 核心要点
- 大型语言模型预训练中，学习率等超参数搜索成本巨大，现有方法难以有效迁移小规模实验结果。
- 论文提出Power Scheduler，基于变量间的幂律关系，实现学习率对训练Token数量和Batch Size的无关性。
- 实验表明，Power Scheduler结合muP，能用同一组超参数在不同模型大小和架构上取得良好效果。
📝 摘要（中文）
为语言模型预训练寻找最优学习率极具挑战，这不仅因为学习率、Batch Size、训练Token数量、模型大小和其他超参数之间存在复杂的关联，还因为对数十亿或数万亿参数的大型语言模型进行超参数搜索的成本过高。最近的研究提出使用小型代理模型和小语料库进行超参数搜索，并将最优参数迁移到大型模型和大型语料库。虽然模型大小相关超参数（如深度和宽度）的零样本迁移在理论和经验上都得到了证明，但从小语料库到大语料库的零样本迁移仍未得到充分探索。本文研究了最近提出的WSD调度器的最优学习率、Batch Size和训练Token数量之间的相关性。经过数千次小型实验，我们发现了变量之间的幂律关系，并证明了其在不同模型大小之间的可迁移性。基于此，我们提出了一种新的学习率调度器Power Scheduler，它与训练Token数量和Batch Size无关。实验表明，将Power Scheduler与最大更新参数化(muP)相结合，无论训练Token数量、Batch Size、模型大小甚至模型架构如何，都可以通过一组超参数始终如一地实现令人印象深刻的性能。我们使用Power Scheduler训练的3B密集和MoE模型实现了与最先进的小型语言模型相当的性能。我们在https://ibm.biz/BdKhLa开源了这些预训练模型。
🔬 方法详解
问题定义：现有大型语言模型预训练中，最优学习率的确定面临挑战。超参数（如学习率、Batch Size、训练Token数量）之间存在复杂关联，且直接在大模型上进行超参数搜索成本过高。利用小模型和小数据集进行搜索，并将结果迁移到大模型的方法，在训练数据量上的迁移效果不佳。
核心思路：论文的核心思路是发现并利用最优学习率、Batch Size和训练Token数量之间的幂律关系。通过大量实验，确定这些变量之间的关系，并基于此设计一种新的学习率调度器，使其对训练Token数量和Batch Size的变化不敏感。这样，就可以在小规模实验中确定合适的超参数，并将其直接应用于大规模训练，从而降低成本。
技术框架：Power Scheduler的核心在于其学习率调整策略，该策略基于对最优学习率、Batch Size和训练Token数量之间关系的建模。整体流程包括：1) 通过大量小规模实验，分析最优学习率与Batch Size和训练Token数量之间的关系；2) 建立基于幂律关系的数学模型；3) 设计Power Scheduler，使其学习率调整策略符合该模型；4) 将Power Scheduler与muP相结合，进行大规模语言模型预训练。
关键创新：Power Scheduler的关键创新在于其学习率调整策略与训练数据量和Batch Size的解耦。传统学习率调度器通常需要针对不同的训练数据量和Batch Size进行调整，而Power Scheduler通过建模变量之间的幂律关系，实现了对这些因素的自适应。这使得超参数的迁移更加容易，降低了大规模预训练的成本。
关键设计：Power Scheduler的具体实现细节未知，但核心在于其学习率更新公式。该公式需要体现最优学习率与Batch Size和训练Token数量之间的幂律关系。此外，论文强调了与Maximum Update Parameterization (muP)的结合，muP是一种参数化方法，可以提高模型训练的稳定性和效率。具体的损失函数和网络结构取决于具体的模型架构，但Power Scheduler的设计目标是使其与不同的模型架构兼容。
🖼️ 关键图片
📊 实验亮点
实验结果表明，结合Power Scheduler与muP，可以使用同一组超参数在不同模型大小、Batch Size和训练Token数量下训练模型，且性能表现优异。使用Power Scheduler训练的3B密集和MoE模型达到了与最先进的小型语言模型相当的性能，证明了该方法的有效性。
🎯 应用场景
Power Scheduler可应用于大规模语言模型的预训练，尤其是在资源受限的情况下。它降低了超参数搜索的成本，加速了模型开发周期。该方法还可推广到其他深度学习任务，例如图像识别和语音识别，只要这些任务也存在类似的超参数依赖关系。
📄 摘要（原文）
Finding the optimal learning rate for language model pretraining is a challenging task. This is not only because there is a complicated correlation between learning rate, batch size, number of training tokens, model size, and other hyperparameters but also because it is prohibitively expensive to perform a hyperparameter search for large language models with Billions or Trillions of parameters. Recent studies propose using small proxy models and small corpus to perform hyperparameter searches and transposing the optimal parameters to large models and large corpus. While the zero-shot transferability is theoretically and empirically proven for model size related hyperparameters, like depth and width, the zero-shot transfer from small corpus to large corpus is underexplored. In this paper, we study the correlation between optimal learning rate, batch size, and number of training tokens for the recently proposed WSD scheduler. After thousands of small experiments, we found a power-law relationship between variables and demonstrated its transferability across model sizes. Based on the observation, we propose a new learning rate scheduler, Power scheduler, that is agnostic about the number of training tokens and batch size. The experiment shows that combining the Power scheduler with Maximum Update Parameterization (muP) can consistently achieve impressive performance with one set of hyperparameters regardless of the number of training tokens, batch size, model size, and even model architecture. Our 3B dense and MoE models trained with the Power scheduler achieve comparable performance as state-of-the-art small language models. We open-source these pretrained models at https://ibm.biz/BdKhLa.