[论文解读] Progressively Stacking 2.0: A Multi-stage Layerwise Training Method for BERT Training Speedup

Cheng Yang, Shengnan Wang+4 more|arXiv (Cornell University)|Nov 27, 2020

Topic Modeling参考文献 22被引用 10

一句话总结

本文提出多阶段逐层训练（MSLT），一种通过逐步堆叠浅层模型并在每个阶段仅微调顶层层来加速BERT预训练的方法，同时冻结先前训练过的层。该方法减少了反向传播计算和通信开销，在性能损失可忽略的情况下实现了超过110%的训练加速。

ABSTRACT

Pre-trained language models, such as BERT, have achieved significant accuracy gain in many natural language processing tasks. Despite its effectiveness, the huge number of parameters makes training a BERT model computationally very challenging. In this paper, we propose an efficient multi-stage layerwise training (MSLT) approach to reduce the training time of BERT. We decompose the whole training process into several stages. The training is started from a small model with only a few encoder layers and we gradually increase the depth of the model by adding new encoder layers. At each stage, we only train the top (near the output layer) few encoder layers which are newly added. The parameters of the other layers which have been trained in the previous stages will not be updated in the current stage. In BERT training, the backward computation is much more time-consuming than the forward computation, especially in the distributed training setting in which the backward computation time further includes the communication time for gradient synchronization. In the proposed training strategy, only top few layers participate in backward computation, while most layers only participate in forward computation. Hence both the computation and communication efficiencies are greatly improved. Experimental results show that the proposed method can achieve more than 110% training speedup without significant performance degradation.

研究动机与目标

• 解决大型BERT模型训练计算成本高和训练时间长的问题。
• 通过减少分布式设置下反向传播和通信开销，提升训练效率。
• 克服现有渐进堆叠方法在后期因全模型更新而变慢的低效问题。
• 利用底层层在训练早期即趋于稳定的现象，避免对已稳定层的重复训练。
• 开发一种在显著缩短训练时间的同时保持模型性能的训练策略。

提出的方法

• 该方法将BERT训练分解为多个阶段，从仅包含少数编码器层的浅层模型开始。
• 在每个阶段，仅更新最顶层新添加的编码器层，而所有先前训练过的层均被冻结，仅参与前向传播。
• 通过在每个阶段逐步增加新层来提升模型深度，同时保持相同的预训练目标。
• 当所有层添加完毕后，进行一次最终的微调阶段，联合更新所有参数，持续几步以优化模型。
• 该方法利用交替优化原理，每次仅更新一小部分参数，从而提高收敛效率。
• 通过先训练未训练的顶层层，再以小学习率联合微调所有层，避免学习率冲突。

实验结果

研究问题

• RQ1 在深层BERT模型中，仅更新最顶层层是否能显著提升训练效率？
• RQ2 在后期阶段冻结已训练过的层是否能在不损害模型性能的前提下减少训练时间？
• RQ3 与标准训练相比，渐进堆叠策略在速度和准确率方面表现如何？
• RQ4 能否利用底层注意力分布在早期即趋于稳定的现象来加速训练？
• RQ5 多阶段逐层训练中，加速与性能下降之间的最优权衡是什么？

主要发现

• 所提出的MSLT方法在相同训练步数下，相比标准BERT训练实现了超过110%的训练加速。
• 在相同训练时间内，MSLT的性能优于标准训练，表现出更优的样本效率。
• 该方法保持了与标准BERT相当的模型性能，仅在GLUE基准测试中出现轻微下降（例如，BERT-Base：79.5 vs. 79.8的平均得分）。
• 最终微调阶段（更新所有层）仅占总训练步数的约20%，表明该阶段前模型已基本优化完成。
• 由于每阶段仅限参数更新，该方法显著减少了反向计算和梯度同步开销，尤其在分布式训练中效果明显。
• 实验结果表明，底层注意力分布在此后阶段基本保持不变，验证了冻结这些层的核心假设。