中文

一种用于加速 Transformer 模型训练的多层级框架

机器学习 2024-04-15 v1 计算与语言

摘要

大规模深度学习模型(如 Bert、GPT 和 ViT)能力的快速增长正在彻底改变 NLP、CV 及许多其他领域的格局。然而,训练此类模型对计算能力提出了前所未有的需求,导致能源成本和二氧化碳排放量呈指数级增长。因此,开发高效的训练解决方案以降低训练成本至关重要。受从典型训练过程中可识别的特征图与注意力之间的层间及层内相似性的一组关键观察所启发,我们提出了一种用于加速训练的多层级框架。具体而言,该框架基于三个基本算子:合并、解合并与插值,通过编排这些算子可以构建一个多层级训练框架。该框架包含一个 V 循环训练过程,该过程逐步对模型尺寸进行下采样与上采样,并通过合并与解合并将参数在相邻层级的模型之间进行投影。其核心思想是,一个能够被快速训练至收敛的较小模型,其训练后的参数可为下一级更大的网络提供高质量的中间解。插值算子旨在打破由解合并引起的神经元对称性,以获得更好的收敛性能。我们在基于 Transformer 的语言模型(如 Bert、GPT)以及视觉模型(如 DeiT)上的实验证明,所提出的框架在训练 BERT/GPT-Base 模型时减少了约 20% 的计算成本,在训练 BERT-Large 模型时最高可减少 51.6% 的计算成本,同时保持了性能。

关键词

引用

@article{arxiv.2404.07999,
  title  = {A Multi-Level Framework for Accelerating Training Transformer Models},
  author = {Longwei Zou and Han Zhang and Yangdong Deng},
  journal= {arXiv preprint arXiv:2404.07999},
  year   = {2024}
}

备注

ICLR 2024