深度学习的大批量优化：76分钟训练 BERT
机器学习
2020-01-06 v5 人工智能
计算与语言
机器学习
摘要
在大规模数据集上训练大型深度神经网络在计算上极具挑战性。近期，利用大批量随机优化方法解决该问题的研究兴趣激增。该研究方向中最突出的算法是 LARS，其通过采用分层自适应学习率，在几分钟内于 ImageNet 上训练 ResNet。然而，LARS 在 BERT 等注意力模型中表现不佳，表明其性能提升在不同任务上并不一致。本文中，我们首先研究了一种原则性的分层自适应策略，以利用大 mini-batch 加速深度神经网络的训练。基于该策略，我们开发了一种称为 LAMB 的新型分层自适应大批量优化技术；随后我们给出了 LAMB 与 LARS 的收敛性分析，表明在一般非凸设定下收敛到驻点。我们的实证结果表明，LAMB 在 BERT 和 ResNet-50 训练等多种任务中具有优越性能，且只需极少的超参数调优。特别地，对于 BERT 训练，我们的优化器支持使用高达 32868 的极大批量大小而不产生任何性能下降。通过将批量大小增加到 TPUv3 Pod 的内存极限，BERT 的训练时间可从 3 天缩短至仅 76 分钟（表 1）。LAMB 的实现见 https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/lamb.py
引用
@article{arxiv.1904.00962,
title = {Large Batch Optimization for Deep Learning: Training BERT in 76 minutes},
author = {Yang You and Jing Li and Sashank Reddi and Jonathan Hseu and Sanjiv Kumar and Srinadh Bhojanapalli and Xiaodan Song and James Demmel and Kurt Keutzer and Cho-Jui Hsieh},
journal= {arXiv preprint arXiv:1904.00962},
year = {2020}
}
备注
Published as a conference paper at ICLR 2020