Recent developments in large language models have sparked interest in efficient pretraining methods. A recent effective paradigm is to perform stage-wise training, where the size of the model is gradually increased over the course of training (e.g. gradual stacking (Reddi et al., 2023)). While the resource and wall-time savings are appealing, it has limitations, particularly the inability to evaluate the full model during earlier stages, and degradation in model quality due to smaller model capacity in the initial stages. In this work, we propose an alternative framework, progressive subnetwork training, that maintains the full model throughout training, but only trains subnetworks within the model in each step. We focus on a simple instantiation of this framework, Random Path Training (RaPTr) that only trains a sub-path of layers in each step, progressively increasing the path lengths in stages. RaPTr achieves better pre-training loss for BERT and UL2 language models while requiring 20-33% fewer FLOPs compared to standard training, and is competitive or better than other efficient training methods. Furthermore, RaPTr shows better downstream performance on UL2, improving QA tasks and SuperGLUE by 1-5% compared to standard training and stacking. Finally, we provide a theoretical basis for RaPTr to justify (a) the increasing complexity of subnetworks in stages, and (b) the stability in loss across stage transitions due to residual connections and layer norm.
翻译:近期大语言模型的发展引发了人们对高效预训练方法的兴趣。一种有效的范式是进行分阶段训练,即在训练过程中逐步增加模型规模(例如渐进堆叠法(Reddi等,2023))。尽管这种方法在资源节省和实际运行时间方面颇具吸引力,但仍存在局限性,特别是无法在早期阶段评估完整模型,以及因初始阶段模型容量较小而导致模型质量下降。本文提出一种替代框架——渐进子网络训练,该框架在整个训练过程中保持完整模型,但每一步仅训练模型内的子网络。我们聚焦于该框架的简单实例——随机路径训练(RaPTr),该方法每一步仅训练一层子路径,并分阶段逐步增加路径长度。实验表明,对于BERT和UL2语言模型,RaPTr在预训练损失上表现更优,同时相比标准训练减少20-33%的浮点运算量,且与其他高效训练方法相比具有竞争力或更优表现。此外,RaPTr在UL2下游任务中表现更佳,相比标准训练和堆叠法,问答任务和SuperGLUE基准提升1-5%。最后,我们从理论上为RaPTr提供依据,论证了:(a) 分阶段子网络复杂度的递增合理性,以及(b) 因残差连接和层归一化带来的阶段转换期间损失的稳定性。