Transformer 语言模型的分阶段训练
计算与语言
2022-03-15 v1
摘要
当前扩展 transformer 语言模型的标准方法是从不同的随机初始化分别训练每个模型规模。作为替代，我们考虑一种分阶段训练设置，从一个小模型开始，通过应用“增长算子”增加模型深度与宽度来逐步增加训练所用计算量。通过将每个阶段初始化为上一阶段的输出，训练过程有效地重用了先前阶段的计算并变得更高效。我们的增长算子各自将整个训练状态（包括模型参数、优化器状态、学习率调度等）作为输入，并输出一个新的训练状态以继续训练。我们确定了这些增长算子的两个重要性质，即它们在应用算子后同时保持损失与“训练动态”。虽然保损失性质此前已被讨论，但据我们所知，本工作是首个确定保持训练动态（训练期间损失下降的速率）重要性的。为寻找各阶段的最优调度，我们使用 (Kaplan et al., 2020) 的缩放定律来寻找精确调度，通过在训练效率开始下降时启动新阶段以节省最多计算。我们在自回归语言模型上经验验证了我们的增长算子与分阶段训练，显示出相比从头训练的强基线最高 22% 的计算节省。我们的代码可在 https://github.com/allenai/staged-training 获取。
引用
@article{arxiv.2203.06211,
title = {Staged Training for Transformer Language Models},
author = {Sheng Shen and Pete Walsh and Kurt Keutzer and Jesse Dodge and Matthew Peters and Iz Beltagy},
journal= {arXiv preprint arXiv:2203.06211},
year = {2022}
}