谷歌核心算法新消息(谷歌实习生新算法提速惊人)
郭一璞 安妮 发自 凹非寺
量子位 报道 | 公众号 QbitAI
从头开始训练一遍当今最强的语言AI需要多久?
现在,这个时间从三天三夜,一下缩短到一个多小时!
带来这个进步的,是Google Brain实习生,尤洋。这位小哥来自中国河南,曾经的清华计算机硕士第一,目前在加州大学伯克利分校读博。
他最近完成的一项研究,把预训练BERT的速度足足提高了64倍。训练时间从4860分钟,变成了76分钟11秒。
而且训练完成后,在机器问答数据集SQuAD-v1上测试一下,F1得分比原来的三天三夜版还要高一点点。
别人家的实习生,究竟又使出了什么神技?
费时到省时
要缩短神经网络的训练时间,本来有成熟的方法可用:
双管齐下,一是堆上大量CPU、GPU或者TPU增加算力,二是增加批大小减少迭代次数。
在计算机视觉研究中,这种方法就很常用。前两天,富士通的研究人员还用这种方法,74.7秒在ImageNet上训练完ResNet-50。
但是,视觉领域的这些训练方法,搬到BERT上是行不通的。BERT是目前工业界训练起来最耗时的应用,计算量远高于ImageNet。
此外,大批量训练还有个“通病”,就是会产生泛化误差(Generalization Gap),导致网络泛化能力下降,如此直接优化往往会导致测试集上准确度下降。
怎么办?
为了能用大批量训练BERT,尤洋和他的同事们提出了LAMB优化器。这是一个通用的神经网络优化器,无论是大批量还是小批量的网络都可以使用,也无需在学习率之外调试超参数。
靠超大批量也适用的LAMB,他们将批大小由512扩展到了65536。
65536是什么概念呢?这已经达到了TPU内存的极限,也是第一次有研究用2000以上的超大批量来训练BERT。
于是,迭代次数大大降低。此前,BERT-Large模型需要1000000次迭代才能完成预训练过程,耗时81.4小时。有了LAMB加持用上大批量,只需要进行8599次迭代,预训练时间直接缩短到76分钟。
这样看来,加速64倍!
因吹斯听,这LAMB到底是何方神器?
LAMB优化器
它的全称是Layer-wise Adaptive Moments optimizer for Batch training,和大家熟悉的SGD、Adam属于同类,都是机器学习模型的优化器(optimizer)。
原本,三天三夜的BERT训练,用的是权重衰减的Adam优化器。
而这一次的新型优化器LAMB,是在论文一作尤洋2017年的一项研究启发下产生的。当时,他提出了一种用于大批量卷积神经网络的优化器LARS。
LARS使用系数eeta控制信任率(trust ratio),但是这种做法可能会导致一些问题,造成一些差异。
因此,在LAMB里,研究团队删除了eeta,对于0|w|或者0|g|的层,直接把信任率设置成1.0,消除了BERT训练的差异。
另外,LARS里用权重衰减来计算信任率:
考虑到权重衰减,LAMB里的信任率公式改成了这样:
另外,LARS虽然在ImageNet上运行的很好,但是在ImageNet上用的模型参数比BERT少得多,因此,这一次,研究团队把LARS中的
改成了:
通过一系列改动,LAMB优化器的最大批量被提升到了32K。
真·优化
是骡子是马,也该拉出来溜溜了。
研究人员用常规训练与混合批训练两种方法,测试LAMB优化器的优化成果,实测效果不错。
在测试中,他们大大增加了算力,选择1024核的TPUv3 Pod进行训练,其中的1024个TPU核心,每秒可以提供超过10亿次浮点(100 petaflops)的混合精度运算。
最后研究人员敲定,和原版BERT模型一样,用Wikipedia和ooksCorpus数据集预训练,然后利用斯坦福的SQuAD-v1数据集进行测试,测试得出的F1 Score的值则用来衡量准确度。
结果显示,随着批量的增大,迭代次数逐减少,F1 Score的波动并不显著,F1值维持在90以上,但训练时间明显缩短。
△ 测试结果
当批大小大于65536、序列长度达到128时,训练时间就没有明显的减少了。
当用了1024块TPU,批大小为32768或65536、迭代8599次时,训练时间缩减到最小,预训练过程只需要76.19分钟。
并且,最终达到了101.8%的弱缩放效率(weak scaling efficiency)。
学霸一作:本硕都是第一名
这项研究的作者是UC伯克利计算机科学部的在读博士尤洋,同时他也是Google Brain的实习生。
小哥哥是一位大学霸,他本科就读于中国农业大学计算机专业,是班里的第一名,硕士保送了清华计算机系,在134名入读清华的佼佼者中,他依然是第一名。
作为第一名的学霸,在申请博士的时候,尤洋一下子喜提了UC伯克利、CMU、芝加哥大学、UIUC、佐治亚理工、西北大学六所名校的全奖offer,简直是名校任挑。
于是,他从六所名校里选择了UC伯克利,UC伯克利刚好位于湾区,尤洋也因此有机会游走于Google Brain、英特尔实验室、微软研究院、英伟达、IBM沃森研究中心等知名企业、研究院实习,趁实习的机会为TensorFlow、英伟达GPU上部署caffe、英特尔CPU部署caffe等大型知名开源项目做出了贡献。
甚至,还有机会趁实习去皮衣哥黄仁勋家里开爬梯~真是让人羡慕啊。
△ 今天没穿皮衣
另外,尤洋也是一位论文高产选手,光是一作的顶会论文就有十几篇,其中还包含去年ICPP的最佳论文,还会有IPDPS 2015的最佳论文。而且,他还拿到了2014年的西贝尔学者奖(Siebel Scholar)。
传送门
Reducing BERT Pre-Training Time from 3 Days to 76 Minutes
Yang You, Jing Li, Jonathan Hseu, Xiaodan Song, James Demmel, Cho-Jui Hsieh
https://arxiv.org/abs/1904.00962
— 完 —
诚挚招聘
量子位正在招募编辑/记者,工作地点在北京中关村。期待有才气、有热情的同学加入我们!相关细节,请在量子位公众号(QbitAI)对话界面,回复“招聘”两个字。
量子位 QbitAI · 头条号签约作者
վ'ᴗ' ի 追踪AI技术和产品新动态
,免责声明:本文仅代表文章作者的个人观点,与本站无关。其原创性、真实性以及文中陈述文字和内容未经本站证实,对本文以及其中全部或者部分内容文字的真实性、完整性和原创性本站不作任何保证或承诺,请读者仅作参考,并自行核实相关内容。文章投诉邮箱:anhduc.ph@yahoo.com