实践中,通过减小模型大小来提高模型训练效率的这种常用做法,实际上与计算效率最佳的训练策略背道而驰。
加州伯克利分校研究者提出一种“先训练更大的模型,然后再压缩模型”的思路,来加速 Transformer 模型训练和推理,这种方法能够兼顾训练效率和训练成本。
在深度学习中,使用更多的计算(例如增加模型大小、数据集大小或训练步骤)通常会提高准确度。随着像 BERT 这样的无监督预训练方法大获成功,能够将训练扩展到非常大的模型和数据集上,更体现了这一思路的有效性。
然而遗憾的是,大规模训练的计算成本非常昂贵,尤其是在没有大型工业界研究实验室提供硬件资源支持的条件下。因此,实践中进行大规模训练的目标往往是,在不超过硬件预算和训练时长的情况下,实现高准确度。
一般的训练预算几乎都无法满足大型模型的训练。因而最大化训练效率的“入门”策略,往往是使用隐藏层尺寸小或网络层少的模型,因为这些模型运行速度更快,并且占用的内存更少。
然而,伯克利博士生 Eric Wallace 等人在最新论文《Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers》中,发现实践中通过减小模型大小来提高模型训练效率的这种常用做法,实际上与计算效率最佳的训练策略背道而驰。
论文链接:https://arxiv.org/pdf/2002.11794.pdf
在论文中,他们还提出“训练更大的模型,然后再压缩模型”的思路来加速Transformer 模型训练和推理,能够兼顾训练效率和训练成本。
在论文《Train Large, Then Compress: Rethinking Model Size for Efficient Training and Inference of Transformers》中,作者发现实践中通过减小模型大小来提高模型训练效率的这种常用做法,实际上与计算效率最佳的训练策略背道而驰。
研究者在计算预算内训练Transformer模型,无非是希望能够大幅增加模型大小,与此同时还能减少训练时长。换言之,即通过论证牺牲收敛性的同时增加模型大小的几率,重新考虑模型在收敛前必须训练的隐含假设。
发生这种现象的原因,是在较少的梯度更新中,较大的模型比较小的模型更能够收敛到较低的测试误差。此外,模型收敛性上的提高胜过使用更大模型的额外计算成本。因此,在考虑时钟训练时长时,较大的模型可以更快地实现较高的准确度。
下图中的两个训练曲线展示了这一趋势。左图绘制了BERT的变体——预训练 RoBERTa的验证错误。对于给定的时钟时间,较深的RoBERTa模型比较浅的模型实现了更低的困惑度(这篇论文表明,针对较宽的模型同样如此)。这种趋势也适用于机器翻译。
右图绘制了一个英-法 Transformer机器翻译模型的验证BLEU分数(越高越好)。在相同的训练时间下,较深和较宽的模型比较小的模型获得了更高的BLEU分数。
有趣的是,对于预训练RoBERTa,增加模型的宽度或深度都会加快训练速度。其中对于机器翻译,较宽的模型要优于较深的模型。因此,作者建议研究者在加深模型之前先尝试增加模型的宽度。
作者还建议增加模型大小而不是批量大小。具体而言,他们证实了,一旦批量大小接近临界范围,增加批量大小只能实现时钟训练时间上的边际改进。因此,在资源限制下,他们建议在临界范围内使用批量大小,之后再考虑尝试增加模型大小。
虽然较大的模型训练效率更高,但它们也增加了模型推理的计算和内存需求。这是有问题的,因为推理的总成本远大于大多数实际应用程序的训练成本。
然而,这篇论文表明了,RoBERTa 可以使用模型压缩来实现这一权衡。特别是,对于模型压缩技术而言,更大的模型比更小的模型鲁棒性更强。因此,研究者可以通过先训练非常大的模型,再对模型进行大量压缩,从而达到两全其美的效果。
其中,作者使用了量化和剪枝的压缩方法。量化方法能够以低精度格式存储模型权重;而剪枝方法则将某些神经网络权重设置为零。这两种方法都可以降低存储模型权重的推理延迟和内存需求。
他们首先在相同的时钟训练时长下,对不同尺寸的RoBERTa 模型进行了预训练。之后在下游文本分类任务(MNLI)上对这些模型进行微调,并进行了剪枝或量化。他们最终发现,在给定的测试时间预算内,训练模型最佳方法是先训练非常大的模型,训练后再对模型进行大量压缩。
例如,下图中的橙色曲线展示了最深的模型的剪枝结果。在不对模型剪枝的情况下,它达到了很高的准确度,不过使用了大约2亿个参数,因此需要大量的内存和计算。
然而,这个模型可以进行大量剪枝(左图,剪枝程度沿着曲线向左变化),而不会严重影响准确度。这与较小的模型形成了鲜明的对比,例如粉线表示的6层模型,剪枝后其准确度会严重下降。右图展示的量化,也有类似的趋势。
总的来说,在测试时间预算内训练模型的最佳方法,也是首先训练非常大的模型,再对模型进行大量压缩。
研究表明,增大Transformer模型的大小可以提高训练和推理的效率,即“先训练大的模型,再压缩模型”。
这一发现还引出了很多其他有趣的问题,比如为什么较大的模型收敛得更快,压缩得更好。
本文对这一现象进行了初步探讨,但仍需进一步研究。另外值得一提的是,这篇论文的发现目前针对NLP,后续可以探索这些结论如何推广到计算机视觉等其他领域。