[论文解读] Staged Training for Transformer Language Models

Sheng Shen, Pete Walsh+4 more|arXiv (Cornell University)|Mar 11, 2022

Topic Modeling被引用 4

一句话总结

本文提出了一种针对Transformer语言模型的分阶段训练方法，即先训练一个小模型，然后通过学习得到的算子逐步增加其深度或宽度，同时保持损失和训练动态不变。通过重用先前的计算资源，并利用缩放定律优化阶段转换，该方法在训练从头开始时实现了高达22%的计算量节省，同时性能保持不变或有所提升。

ABSTRACT

The current standard approach to scaling transformer language models trains each model size from a different random initialization. As an alternative, we consider a staged training setup that begins with a small model and incrementally increases the amount of compute used for training by applying a "growth operator" to increase the model depth and width. By initializing each stage with the output of the previous one, the training process effectively re-uses the compute from prior stages and becomes more efficient. Our growth operators each take as input the entire training state (including model parameters, optimizer state, learning rate schedule, etc.) and output a new training state from which training continues. We identify two important properties of these growth operators, namely that they preserve both the loss and the "training dynamics" after applying the operator. While the loss-preserving property has been discussed previously, to the best of our knowledge this work is the first to identify the importance of preserving the training dynamics (the rate of decrease of the loss during training). To find the optimal schedule for stages, we use the scaling laws from (Kaplan et al., 2020) to find a precise schedule that gives the most compute saving by starting a new stage when training efficiency starts decreasing. We empirically validate our growth operators and staged training for autoregressive language models, showing up to 22% compute savings compared to a strong baseline trained from scratch. Our code is available at https://github.com/allenai/staged-training.

研究动机与目标

• 降低大规模Transformer语言模型训练的高计算成本。
• 解决标准训练中的低效问题，即大模型在初期计算效率较低，仅在后期才变得更为高效。
• 设计能够同时保持损失和训练动态的模型增长算子。
• 利用缩放定律制定合理的阶段转换调度，以最大化计算效率。
• 在GPT-2风格自回归模型上通过标准基准（如Wikitext-103和LAMBADA）对方法进行实证验证，实现可测量的计算量节省。

提出的方法

• 该方法使用深度和宽度增长算子，以完整训练状态（权重、优化器状态、学习率调度）作为输入，输出一个更大规模的新模型状态。
• 增长算子被设计为保持损失和训练动态，确保增长后损失及其下降速率与从头训练的模型保持一致。
• 阶段调度基于Kaplan等人（2020）提出的缩放定律，用于预测训练效率开始下降的时刻，并决定新阶段的启动时机。
• 在实际场景中，当精确缩放定律未知时，推导出最优调度的实用近似方法。
• 该方法在GPT-2风格自回归语言模型上进行了验证，使用Wikitext-103和LAMBADA等标准基准。
• 该方法通过重用先前的训练阶段，避免冗余计算，显著降低了大模型的训练计算量。

实验结果

研究问题

• RQ1 使用模型增长的分阶段训练能否降低大规模Transformer语言模型的总计算成本？
• RQ2 增长算子必须具备哪些特性，才能确保在各阶段之间实现稳定且高效的训练？
• RQ3 如何确定最优的阶段转换调度以实现最大计算量节省？
• RQ4 在模型增长过程中，完整训练状态（包括优化器和学习率）在多大程度上可以被保留？
• RQ5 保持训练动态是否能带来比简单增长方法更好的收敛性和性能表现？

主要发现

• 所提出的增长算子能够同时保持损失和训练动态，确保增长后的损失下降速率与从头训练的模型一致。
• 利用缩放定律确定阶段转换可实现相比从头训练高达22%的计算量节省。
• 在GPT-2风格模型上的实证结果表明，增长后的模型在显著减少计算量的前提下，实现了相当或更优的验证损失和零样本困惑度。
• 该方法仅使用全量GPT-2基础模型15%的计算量，即可在Wikitext-103上达到37.5的PPL；在LAMBADA上实现等效准确率时，仅需33%的计算量。
• 过早或过晚增长会导致计算量节省的损失，证实了精确阶段调度的重要性。
• 该方法具有鲁棒性和泛化能力，在不同模型规模和训练目标下均能保持性能稳定。