Efficient Stagewise Pretraining via Progressive Subnetworks

📄 arXiv: 2402.05913v2 📥 PDF

作者: Abhishek Panigrahi, Nikunj Saunshi, Kaifeng Lyu, Sobhan Miryoosefi, Sashank Reddi, Satyen Kale, Sanjiv Kumar

分类: cs.CL, cs.LG

发布日期: 2024-02-08 (更新: 2024-10-13)

💡 一句话要点

提出渐进子网络训练以提升预训练效率

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 阶段性训练 子网络训练 高效预训练 随机部分训练 大型语言模型 自然语言处理 模型性能提升

📋 核心要点

  1. 现有的阶段性训练方法如层丢弃被认为效果不佳，限制了预训练的效率和性能。
  2. 本文提出渐进子网络训练框架，通过逐步训练子网络并增加其规模，挑战传统观点，展示了层丢弃的潜力。
  3. 实验结果表明，RAPTR方法在训练速度上提升了33%，并在UL2的下游任务中提高了1.5%的性能。

📝 摘要（中文）

近年来，大型语言模型的发展引发了对高效预训练方法的关注。阶段性训练方法，如逐步堆叠和层丢弃，虽然受到关注，但普遍认为层丢弃策略效果不佳。本文挑战这一观点，提出渐进子网络训练框架，通过逐步增加子网络的规模，展示了层丢弃策略在适当设计下的竞争力。我们提出的随机部分训练（RAPTR）方法在每个训练步骤中选择并训练随机子网络，显著加快了BERT和UL2等标准基准的训练速度，且在UL2的下游任务中表现优于标准训练，提供了更好的归纳偏置。

🔬 方法详解

问题定义：本文旨在解决现有阶段性训练方法（如层丢弃）在效率和性能上的不足，尤其是其在大型语言模型预训练中的应用局限性。

核心思路：提出渐进子网络训练框架，通过逐步训练子网络并在训练过程中逐渐增加其规模，旨在提高训练效率并改善模型性能。

技术框架：整体架构包括随机部分训练（RAPTR）方法，该方法在每个训练步骤中选择并训练随机子网络，逐步增加网络的深度或宽度，直到训练完整个网络。

关键创新：最重要的创新在于提出了渐进子网络训练的理论基础，证明了逐步增加子网络复杂性和在阶段转换中保持损失稳定的有效性，与传统的层丢弃方法本质上不同。

关键设计：在RAPTR中，关键设计包括随机选择子网络的策略、逐步增加子网络规模的机制，以及在训练过程中保持残差连接和层归一化等现代架构组件的稳定性。

📊 实验亮点

实验结果显示，RAPTR方法在标准基准上训练速度提升了33%，并在UL2的问答任务和SuperGLUE上提高了1.5%的性能，表明其在效率和效果上的显著优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的大型语言模型预训练，尤其是在需要高效训练和快速迭代的场景中。通过提升训练效率和下游任务性能，RAPTR方法能够为实际应用提供更具竞争力的解决方案，推动相关技术的发展。

📄 摘要（原文）

Recent developments in large language models have sparked interest in efficient pretraining methods. Stagewise training approaches to improve efficiency, like gradual stacking and layer dropping (Reddi et al, 2023; Zhang & He, 2020), have recently garnered attention. The prevailing view suggests that stagewise dropping strategies, such as layer dropping, are ineffective, especially when compared to stacking-based approaches. This paper challenges this notion by demonstrating that, with proper design, dropping strategies can be competitive, if not better, than stacking methods. Specifically, we develop a principled stagewise training framework, progressive subnetwork training, which only trains subnetworks within the model and progressively increases the size of subnetworks during training, until it trains the full network. We propose an instantiation of this framework - Random Part Training (RAPTR) - that selects and trains only a random subnetwork (e.g. depth-wise, width-wise) of the network at each step, progressively increasing the size in stages. We show that this approach not only generalizes prior works like layer dropping but also fixes their key issues. Furthermore, we establish a theoretical basis for such approaches and provide justification for (a) increasing complexity of subnetworks in stages, conceptually diverging from prior works on layer dropping, and (b) stability in loss across stage transitions in presence of key modern architecture components like residual connections and layer norms. Through comprehensive experiments, we demonstrate that RAPTR can significantly speed up training of standard benchmarks like BERT and UL2, up to 33% compared to standard training and, surprisingly, also shows better downstream performance on UL2, improving QA tasks and SuperGLUE by 1.5%; thereby, providing evidence of better inductive bias.