POWER SCHEDULER：一种与批次大小和token数量无关的学习率调度器

摘要

为语言模型预训练找到最优的学习率是一项具有挑战性的任务。这不仅因为学习率、批次大小、训练token数量、模型大小以及其他超参数之间存在复杂的关联，还因为对具有数十亿或数万亿参数的大型语言模型进行超参数搜索的成本极高。最近的研究提出使用小型代理模型和小规模语料库进行超参数搜索，并将最优参数迁移到大型模型和大规模语料库上。虽然与模型大小相关的超参数（如深度和宽度）的零样本迁移性已在理论和实验上得到验证，但从小规模语料库到大规模语料库的零样本迁移性尚未得到充分探索。在本文中，我们研究了最近提出的WSD调度器中学习率、批次大小和训练token数量之间的相关性。经过数千次小型实验，我们发现变量之间存在幂律关系，并证明了其在模型大小间的可迁移性。基于这一观察，我们提出了一种新的学习率调度器——Power调度器，它对训练token数量和批次大小不敏感。实验表明，将Power调度器与最大更新参数化（µP）结合，可以仅用一组超参数在不同训练token数量、批次大小、模型大小甚至模型架构下持续取得优异性能。我们使用Power调度器训练的3B密集模型和MoE模型在性能上可与最先进的小型语言模型相媲美。我们在https://ibm.biz/BdKhLa上开源了这些预训练模型。

1 引言

学习率是深度神经网络训练中的一个关键超参数。在大型语言模型（LLMs）的背景下，余弦学习率调度器是最常用的策略。它已被证明在多个最先进的模型中有效，包括Llama 3（Dubey等，2024）、Gopher（Rae等，2021）等。然而，余弦调度器需要预先定义训练步数以达到最优损失。这导致两个主要缺点：1）中间训练检查点次优；2）现有语言模型的持续训练变得复杂。
MiniCPM（Hu等，2024）提出了Warmup-Stable-Decay（WSD）学习率调度器（如图1所示）以解决这些问题。WSD学习率调度分为三个阶段：1）预热阶段，将学习率从0线性增加到峰值；2）稳定阶段，将学习率保持在峰值并训练模型大部分时间；3）衰减阶段，在相对较短的时间内将学习率退火至0。该调度器的主要优势在于无需提前指定训练步数。这对于大规模训练尤为方便，因为可以随时应用衰减以观察模型性能并决定是否停止。它还默认支持持续学习，因为可以从稳定阶段的检查点恢复训练。此外，可以在衰减阶段更改数据混合比例以增加高质量数据的比例。这种数据课程在最近的几种语言模型中被证明是有效的（Team等，2023；Hu等，2024；Dubey等，2024；Shen等，2024）。

然而，WSD是否真的对token数量不敏感？在我们的实验中，我们发现尽管WSD调度器在理论上可以无限期地保持在稳定阶段，但不同训练token数量的最优学习率是不同的。换句话说，最优学习率调度器与训练token数量密切相关。因此，WSD调度器仍然面临与余弦调度器相同的问题：如果token数量与原始计划相差太大，中间和持续训练的检查点将是次优的。

此外，确定最优学习率对于大规模预训练仍然具有挑战性。正如我们的实验将展示的那样，超参数（包括学习率、模型大小、批次大小和训练步数）之间存在复杂的相关性。现代LLM的规模和训练成本使得在目标模型大小和训练计算量上进行超参数搜索变得不可能。研究人员提出使用小型代理模型进行超参数搜索，并从搜索结果中预测大型模型的最优超参数（Dey等，2023；Hu等，2024；Yang等，2022）。在这些方法中，µTransfer（Yang等，2022；2023）被提出以促进不同模型大小之间的零样本超参数迁移。µTransfer已成功应用于多种语言模型，包括Cerebras-GPT（Dey等，2023）、miniCPM（Hu等，2024）和AFM（Gunter等，2024）。

在本文中，我们首先结合WSD调度器和µTransfer来研究代理模型与大型模型之间的学习率迁移。我们的大量实验表明，µTransfer并未为WSD优化器提供跨token数量和批次大小的直接零样本学习率迁移性。相反，最优学习率 ${\eta }_{\mathrm{opt}}$与批次大小$\beta$ 和token数量T之间满足幂律关系：
KaTeX parse error: Undefined control sequence: \label at position 61: …ta \cdot a T^b \̲l̲a̲b̲e̲l̲{eq:wsd-opt-lr}…
其中，a和 b是幂律系数。此外，我们的实验证实了幂律关系在不同模型大小之间的零样本可迁移性，不同模型大小的系数非常相似。受这一观察启发，我们提出了一种新的学习率调度器——PowerLR，它对批次大小和token数量不敏感。它允许直接跨批次大小、token数量和模型大小迁移最优学习率调度。因此，昂贵的预训练可以在无需指定训练token数量的情况下进行，从而支持提前停止和持续预训练，而不会牺牲收敛性。

2 背景

2.1 最大更新参数化（$\mu \mathrm{P}$）

最大更新参数化（Maximal Update Parameterization, µP）（Yang & Hu, 2020; Yang et al., 2022; 2023）通过控制初始化、逐层学习率和激活幅度，确保训练在分析上保持稳定，且与模型的宽度和深度无关。除了提高训练稳定性外，$\mu \mathrm{P}$还改善了从小型代理模型到大型模型的训练超参数的可迁移性，这种技术被称为µTransfer。µP的超参数可迁移性在理论上已被证明适用于宽度（Yang et al., 2022）和深度（Yang et al., 2023）。先前的研究还提供了跨批次大小、训练步数和训练序列长度的可迁移性的实验证据。

在本文中，我们遵循CerebrasGPT（Dey等，2023）中使用的$\mu \mathrm{P}$配置，研究了批次大小和学习率在不同训练token数量和模型大小之间的可迁移性。表1列出了我们对模型初始化、学习率和乘数所做的$\mu \mathrm{P}$调整。

2.2 预热-稳定-衰减（WSD）调度器

MiniCPM（Hu等，2024）提出了WSD学习率调度器，将预训练分为三个阶段：预热阶段、稳定训练阶段和剩余的衰减阶段。WSD调度器定义如下：
$$\text{ WSD }(n)=\{\begin{array}{ll}\frac{n}{N_{\text{warmup}}}\cdot \eta & \text{ if }n<N_{\text{warmup}}\\ \eta & \text{ if }{N}_{\text{warmup}}<n\le N-N_{\text{decay}}\\ f(n,N,N_{\text{decay}})\cdot \eta & \text{if }n>N-N_{\text{decay}}\end{array},\text{\hspace{1em}(2)}$$

其中， n是当前步数，$\eta$是稳定学习率，N是总步数，$N_{\text{wamup}}$ 是衰减步数，$N_{\mathrm{decay}}$是衰减步数,$f(n,N,N_{\text{decay}})$是学习率衰减函数。WSD调度器相较于其他调度器（如余弦和线性调度器）的主要优势在于无需提前指定训练步数。由于稳定训练阶段的学习率不依赖于训练步数，WSD调度器无需提前确定总训练步数。这对于大规模训练尤为方便，因为可以随时启动冷却阶段以观察模型行为并决定是否停止（Hagele等，2024）。它还便于扩展预训练，可以应用于稳定训练阶段的最后一个检查点。此外，在冷却阶段可以更改数据混合比例，以在训练结束时加入更多高质量数据。这种数据混合策略在最近的许多语言模型中被证明是有效的（Hu等，2024；Dubey等，2024；Team等，2024）。

3 使用$\mu \mathrm{P}$为WSD调度器搜索最优学习率

在本节中，我们专注于为小型代理模型寻找最优学习率$\eta$。由于我们使用了$\mu \mathrm{P}$，我们期望最优学习率在不同模型大小、批次大小和训练步数之间保持不变。为了验证这一点，我们进行了大量实验，以寻找不同模型大小和训练token数量的最优学习率。表2展示了我们的整体超参数搜索配置。我们使用WSD调度器，并按照先前的工作（Hu等，2024；Hagele等，2024）在衰减阶段使用10%的token。本节及下一节中的所有模型均在RedPajama（Computer，2023）语料库上训练，并在保留测试集上进行测试。

3.1 最优学习率是否具有可迁移性？

我们进行了两项对照实验来验证最优学习率的可迁移性。首先，我们固定批次大小$\beta =128$，并遍历所有学习率$\eta$和训练token数量$T$的组合。图2(a)显示，最优学习率${\eta }_{\text{opt}}$在广泛的$T$范围内并不具有可迁移性。

${\eta }_{\mathrm{opt}}$倾向于随着训练token数量的增加而减小。例如，对于20亿训练token，${\eta }_{\mathrm{opt}}$为0.0128，而对于2560亿训练token，${\eta }_{\text{opt}}$为0.0008。其次，我们将训练token数量$T$固定为1280亿token，并遍历所有学习率$\eta$和批次大小的组合。图2(b)显示，最优学习率在广泛的$\beta$范围内并不具有可迁移性，但倾向于随着批次大小的增加而增加。对于批次大小16，最优学习率${\eta }_{\mathrm{opt}}$为0.0002，而对于批次大小512，${\eta }_{\text{opt}}$为0.0064。

3.2 $et{a}_{\text{opt}}$、β 和 $T$之间的关系是什么？

为了更好地理解${\eta }_{\text{opt}}$、$\beta$和$T$之间的关系，我们遍历了这三个变量的所有可能组合。图3展示了每种批次大小和训练token组合下的最优学习率。与我们之前的观察一致，最优学习率在不同批次大小下始终随着训练token数量的增加而减小。此外，

我们还注意到，对于每种训练token数量，最优学习率${\eta }_{\text{opt}}$与批次大小$\beta$之间的比例相对稳定。基于这一观察，我们提出了第一个假设：

假设1：对于WSD调度器和给定的$(T,\beta )$组合，最优学习率${\eta }_{\text{opt}}$与训练批次大小$\beta$成正比。

因此，我们定义$\gamma$为${\eta }_{\text{opt}}$与$\beta$的比值：
$$\gamma =\frac{{\eta }_{\text{opt}}}{\beta }(3)$$
为了验证假设1，我们对三种模型大小（12M、36M和121M）进行了广泛的超参数搜索。在为每种$(T,\beta ,\text{模型大小})$组合找到最优学习率${\eta }_{\text{opt}}$后，我们仅保留三种最佳批次大小以聚焦于最优情况。每种$T$和模型大小下的三种最佳批次大小的$\gamma$值分别绘制在图4(a)、图4(b)和图4©中。这些结果表明，在给定固定训练token数量$T$的情况下，$\gamma$值落在一个相对较小的范围内。

此外，我们注意到$\gamma$与训练token数量$T$之间近似遵循幂律关系。因此，我们提出了第二个假设：

假设2：学习率与批次大小的比值$\gamma$与$T$之间存在幂律相关性：
$$\gamma =a{T}^b(4)$$
当使用µP时，这种相关性可以在不同模型大小之间迁移。
为了验证假设2，我们计算了最佳批次大小的平均$\gamma$值，以估计每种token数量$\beta$下的实际$\gamma$。图4(d)显示，所有三种模型大小都共享相似的幂律相关性。在汇总三种模型大小的结果后，我们得到了以下相关性：
$$\gamma =4.6T^{-0.51}(5)$$
其中，$x$ 是训练token的数量。
换句话说，WSD调度器的最优学习率${\eta }_{\text{opt}}$、批次大小$\beta$和训练token数量$T$之间的关系可以用以下方程近似建模：
$${\eta }_{\text{opt}}=\beta \cdot a{T}^b(6)$$
公式6提供了一种简单的方法，可以根据WSD调度器的训练token数量和训练批次大小预测最优学习率。与先前的工作类似，(a)和(b)可以通过在小型代理模型上进行超参数搜索轻松估计。

然而，训练token数量与最优学习率之间的这种关系限制了在训练前需要指定训练步数。我们还注意到，该公式在大规模训练语料库上往往会提供较小的学习率。例如，给定一个10万亿token的语料库和批次大小1024，预测的最优学习率为0.0011。如果我们考虑到µP会将模型内部矩阵的学习率除以一个因子$m_{\text{width}}=d_{\text{model}}/d_{\text{base}}$，那么对于大型模型，实际学习率将非常小。虽然较小的学习率有利于最终收敛，但也可能在训练初期导致探索不足。为了解决这些问题，我们将在下一节提出一种新的幂律学习率调度器。

4 幂律调度器
基于前一节的观察，我们提出了一种新的幂律学习率调度器：
$${\eta }_{\text{power}}(n)=\min \left({\eta }_{\text{max}},\beta \cdot a{n}^b\right)(7)$$
其中，$\beta$ 是批次大小，$n$ 是已训练的token数量，$a$ 是学习率的振幅，$b$ 是学习率随训练token数量衰减的幂律指数，${\eta }_{\text{max}}$ 是学习率的上限，用于避免过大的学习率。与恒定学习率类似，幂律学习率也不需要提前指定训练步数或总训练token数量，因为学习率仅依赖于当前已训练的token数量。

与WSD调度器类似，我们可以将幂律学习率与预热和衰减结合，得到最终的Power调度器：
$$\text{Power}(n)=\{\begin{array}{ll}\frac{n}{N_{\text{warmup}}}\cdot {\eta }_{\text{power}}(N_{\text{warmup}})& \text{如果 }n<N_{\text{warmup}}\\ {\eta }_{\text{power}}(n)& \text{如果 }{N}_{\text{warmup}}<n\le N-N_{\text{decay}}\\ f(n,N,N_{\text{decay}})\cdot {\eta }_{\text{power}}(N-N_{\text{decay}})& \text{如果 }n>N-N_{\text{decay}}\end{array}(8)$$

Power调度器需要三个独立的超参数$({\eta }_{\text{max}}$、$a$、$b$），而余弦调度器和WSD调度器仅需要一个超参数。然而，我们期望这些超参数能够在不同模型大小、训练token数量和批次大小之间迁移。此外，由于图1显示${\eta }_{\text{max}}$对整体学习率曲线的影响非常有限，我们可以简单地将其设置为一个足够大的值。在本文中，我们将${\eta }_{\text{max}}$设置为0.02。然后，我们对$a$和$b$进行了另一次超参数搜索，以找到最优的学习率超参数。表3展示了Power调度器超参数搜索中使用的配置。

图5展示了320亿训练token的超参数搜索结果。我们观察到，在搜索范围$[3,5]$内，测试困惑度对$a$的不同选择不敏感。相比之下，测试困惑度对$b$的选择更为敏感。最优的$b$值在$-0.52$到$-0.51$之间。在不同训练token数量的实验中，我们也观察到了类似的现象。图6显示，对于不同的$\beta$和$T$，最优的$b$值始终落在$[-0.49,-0.53]$范围内，仅有少数例外。结合这些观察结果，我们在本文的其余部分选择了$a=4$和$b=-0.51$。

5 预训练实验

本节将Power调度器与WSD和余弦调度器在不同场景下进行比较。在第一部分中，我们进行了对照实验，使用不同的学习率调度器训练1B参数的Transformer模型和1B的混合专家（MoE）模型，以展示Power调度器与其他调度器相当或更优的性能。在第二部分中，我们采用更实际的设置，训练了一个3B的Transformer语言模型和一个3B的MoE语言模型，并在衰减阶段使用高质量数据，与强大的开源语言模型进行比较。在本节的所有Power调度器实验中，我们将使用第4节中的超参数：$a=4$、$b=-0.51$以及${\eta }_{\text{max}}=0.02$。

5.1 1B对照实验

我们使用WSD、余弦和Power调度器训练了一系列1B参数的密集模型和混合专家（MoE）Transformer模型。所有模型均使用1万亿token进行训练。我们采用了MiniCPM（Hu等，2024）中提出的最优超参数，并通过µTransfer将其适配到我们的1B设置中。所有模型的批次大小均为1024。对于WSD和余弦调度器，我们设置$\eta =0.01$，这是MiniCPM中的最优学习率。对于Power和WSD调度器，我们在最后1000亿token中将学习率指数衰减至0。MoE模型使用ScatterMoE（Tan等，2024）实现。更多模型细节见表4。

我们在语言模型任务和LM评估工具（Gao等，2024）中的多项选择任务上评估了所有1B模型。多项选择任务包括小学科学问题（ARC，Clark等，2018）、是非问题（BoolQ，Clark等，2019）、常识推理（Hellaswag，Zellers等，2019）、开放书籍问答（OpenBookQA，Mihaylov等，2018）、物理问题（PIQA，Bisk等，2020）以及Winograd模式任务（Winogrande，Sakaguchi等，2021）。表5展示了性能结果。Power调度器在语言建模和下游任务中均表现出一致更好或相当的性能。令人惊讶的是，

尽管我们仅在一个小型密集模型上进行了超参数搜索，但这些超参数在MoE模型上仍然表现良好。

5.2 3B实际实验

为了与强大的开源语言模型进行性能比较，我们预训练了两个3B语言模型：1）PowerLM-3B，一个3B的密集语言模型；2）PowerMoE-3B，一个3B的MoE语言模型。我们使用Hu等（2024）中提出的两阶段训练模式和Power调度器对这两个模型进行预训练。第一阶段线性预热学习率，然后应用幂律衰减。训练语料库是混合了大规模、中等质量的开源数据集（具有宽松许可证）。PowerLM-3B训练了1万亿token，PowerMoE-3B训练了2.5万亿token。第二阶段将学习率指数衰减至零。训练语料库是第一阶段数据与少量高质量开源/合成语料库（具有宽松许可证）的混合。PowerLM-3B训练了2500亿token，PowerMoE-3B训练了5000亿token。训练批次大小为1024。

表6展示了我们的模型与最先进模型在多项选择任务中的性能。表7展示了在数学和代码领域的MMLU和生成性能。结果表明，尽管训练token数量相对较少，我们的PowerLM-3B仍然在性能上与最先进的2B到4B语言模型相当。此外，我们的PowerMoE-3B仅使用了8亿活跃参数，但性能与最先进的1B到2B密集模型相当。

6 结论

在本文中，我们系统地研究了最优学习率、批次大小和训练token数量之间的关系。在使用WSD学习率调度器时，我们观察到这些变量之间存在幂律关系。受此启发，我们提出了一种新的学习率调度器——Power调度器，它对token数量和批次大小不敏感。实验表明，将Power调度器与最大更新参数化（$\mu \mathrm{P}$）结合，可以仅用一组超参数在不同训练token数量、批次大小、模型大小甚至模型架构下持续取得优异性能。