BERT训练猛提速！谷歌新研究将BERT预训练时间从3天缩短到76分钟

YnHzGpYq

新智元报道
编辑：肖琴
" G4 m6 m# y' v
【新智元导读】BERT是目前最强大的NLP预训练模型，也是工业界目前最耗时的应用，计算量远高于ImageNet。谷歌的研究人员提出新的优化器，使用1024块TPU，将BERT的训练时间从3天成功缩短到76分钟，提速 65.2 倍！
! k# [& `+ b' u: k4 f
去年，谷歌发布了最强预训练模型 BERT，宣告了NLP领域的一项重大突破。
! W5 A, c3 m' e1 HBERT 在 33 亿文本的语料上训练语言模型，再分别在不同的下游任务上微调，在11个不同的 NLP 任务均得到了目前为止最好的结果。
不过，在 33 亿文本的语料上预训练一个 BERT 模型的成本是非常大的，谷歌用了 16 个自己的 TPU 集群（一共 64 块 TPU）来训练大号版本的 BERT，一共花了约4天的时间。
如此巨大的训练成本，让普通研究者难以尝试自己去训练一个BERT。
7 U3 a+ b7 o4 ?- C9 Y. [有没有办法加快BERT的训练呢？近日，来自Google、UC Berkeley、UCLA的几位研究人员提出新的优化器——LAMB 优化器，将训练的batch size推到硬件的极限，使用 TPU Pod ( 1024块 TPUv3 芯片)，将BERT的训练时间从3天缩短到了76分钟！
! X& j1 a! C6 X) [) J
0 \0 V: `9 i( \- `. W$ W论文地址：
3 O5 q2 d2 s/ P1 s4 @' V* N其中一作尤洋(Yang You)来自UC Berkeley，这项工作于他在Google Brain实习期间完成。
接下来，新智元带来对这篇论文的译介：
; i/ J8 R2 m, E5 b) n加快深度神经网络最有效的方法 大批量训练 (large-batch training) 是加快大型分布式系统中深度神经网络训练的关键。然而， large-batch 训练是很困难的，因为它会产生一种泛化差距 (generalization gap)。直接优化通常会导致测试集的准确性下降。
BERT 是一种最先进的深度学习模型，建立在用于语言理解的深度双向 transformers 之上。对 BERT 来说，当扩大批大小 (例如超过 8192) 时，以前的 large-batch 训练技术效果并不好。BERT 的预训练也需要很长时间才能完成 (使用 16 个 TPUv3 芯片大约需要 3 天)。
5 w# R5 r) \* x! e4 I; z* M, h, b2 P2 y为了解决这个问题，我们提出了 LAMB 优化器，它帮助我们将批大小扩大到 65536，而不会丢失准确性。
LAMB 是一个通用的优化器，适用于小批量和大批量，并且除了学习率外不需要超参数调优。基线 BERT-Large 模型需要 100 万次迭代才能完成预训练，而 batch size 为65536/32768 的 LAMB 只需要 8599 次迭代。我们将 batch size 推到 TPUv3 pod 的内存上限，可以在 76 分钟内完成 BERT 训练 (表 1)。

表 1：我们使用 SQuAD-v1 的 F1 score 作为精度指标。F1 的基线成绩是由 BERT 的公共 github 提供的预训练模型 (BERT- large) 实现的 (截止到 2019 年 2 月 1 日)。我们在实验中使用 tpuv3。我们使用了与基线相同的设置：总 epochs 的前 9/10 使用序列长度128，最后 1/10 使用序列长度 512。所有的实验运行相同数量的 epochs。
4 M7 r7 T3 `: H5 d8 W3 L( m6 t9 y深度神经网络的训练是十分耗时的。目前，减少训练时间最有效的方法是使用多个芯片(如 CPU、GPU 和 TPU) 来并行化 SGD 变体的优化过程。由于前向传播和反向传播中不同层之间的数据依赖关系，使得跨层的并行化效率并不高。相反，研究人员在每次迭代中并行化小批量中的数据点。如果确定了训练的 epochs 的数量，那么线性地增大batch size 意味着会线性地减少迭代次数 (即更新权重的次数)。为了最小化训练时间，最大化 batch size 将是理想的。
8 Q! B: S( f! Q. k! g: j" e然而，大批量的训练是困难的。例如，使用大小为 512 的 batch size 训练在 ImageNet上训练 AlexNet，能实现 80% 以上的 top-5 测试精度。但将 batch size 扩大到 4096之后，直接训练可能只能获得 50% ~ 60% 的 top 5 精度。
" |2 t; v9 B9 x- C5 T6 P6 HKeskar 等人 (10) 认为在大批量训练中存在一个泛化差距 (generalization gap)。Hoffer等人 (6) 认为，训练时间越长，泛化差距越小。然而，训练时间过长意味着进行大批量训练就没有好处了。
6 C' B$ J# z1 F3 u& \' V因此，大批量训练的目标是在一定数量的 epochs 内达到可观的精度。通过设计一系列的学习率计划表，研究者已经可以将 ImageNet 训练的 batch size 扩大到 32K，并且精度损失较小。据我们所知， Ying et al. 实现了目前最快的 ImageNet 训练速度，并且达到了 76+% 的 top-1 精度。通过使用 LARS 优化器，将 batch size 扩展到 32K,，Ying等人使用 TPUv3 Pod，在 2.2 分钟内完成了 ResNet-50 的 ImageNet 训练。（最新，富士通研究院刷新了这一速度，将 ImageNet 训练时间降到 74.7 秒）
BERT 是目前最先进的深度学习语言模型。BERT 建立在用于语言理解的深度双向transformers 之上。对 BERT 来说，当将 batch size 扩大到非常大时 (例如超过8192)，以前的 large-batch 训练技术效果并不好。BERT 的预训练也需要很长时间才能完成 (使用 16 个 TPUv3 芯片大约需要 3 天)。
$ V- I! H* k8 |6 b9 p为了扩大 BERT 的 batch size，本文提出 LAMB 优化器。LAMB 支持自适应 element-wise updating 和精确的逐层修正 (layer-wise correction)。
7 g+ }  D7 O( l. p' P2 B7 ILAMB 是一个适用于小批量和大批量的通用优化器。用户只需要调整学习率，不需要调其他超参数。使用 LAMB，我们可以将 BERT 预训练的批大小扩大到 64K，而不会丢失准确性。
& F5 T! x" D5 L: ?) R. UBERT 预训练包括两个阶段：
+ Q; L/ ^* w; w7 k7 Y! A(1) 前 9/10 的 epochs 使用 128 的序列长度；
( O0 V) d* {& D5 X2 W0 k+ H8 {(2) 后 1/10 epochs 使用 512 的序列长度。
! l- \: y" m( ^% `% s3 k: t
, @) @+ T. r: F: f4 s: cbaseline 需要 100 万次迭代来完成 BERT 预训练，但我们只需要 8599 次迭代，这使我们能够将 BERT 训练时间从 3 天减少到 76 分钟。
! A4 C3 o8 Y$ r+ Z' k7 v0 ?* N( x我们将批大小推到了 TPU Pod 的硬件极限。批大小大于 32768(序列长度为 512) 的话将耗尽内存。批大小大于 65536(序列长度为 128) 则不会带来任何加速。我们的优化器可以将批大小扩大到 128k，甚至更大。由于硬件限制，序列长度为 512 的设置下，我们在批大小达到 32768 时停下，在序列长度为 128 的设置下，批大小达到 65536 时停止。
+ C. k5 G$ h  z( A$ B6 B0 m3 l本文中所有的 BERT 模型都指 BERT-Large 模型。为了进行公平的比较，本文所有的实验都运行相同数量的 epochs(即固定数量的浮点运算)。我们的结果如表 1 所示。
' w; ^2 z4 S( f; N  aLAMB优化器 LAMB的全称是Layer-wise Adaptive Moments optimizer for Batch training。
7 h; L4 c. g8 w9 ^BERT 训练的基线使用权重衰减的 Adam 作为优化器，这是 Adam 优化器的一个变体。另一个成功应用于大批量卷积神经网络训练的自适应优化器是 LARS。
这些优化器启发我们提出了新的优化器，用于大批量 BERT 训练。我们提出的 LAMB 优化器的概述如算法 1 所示。
7 r$ V* Z  P  \; {
实验和结果
常规训练
( N" L3 F  ~- k! tTPU 是浮点运算的强大计算硬件。我们在所有的实验中都使用了 TPUv3。TPUv3 Pod有 1024 个芯片，可以为混合精度计算提供超过 100 petaflops 的性能。我们的结果如表 1 所示。基线模型在预训练时使用 Wikipedia 和 BooksCorpus 数据集。
9 f$ X! [  Z" J我们使用了与原始 BERT 模型相同的数据集，即 Wikipedia 和 BooksCorpus，分别有2.5B 和 8 亿单词。原始 BERT 模型的作者首先以 128 的序列长度进行了 900k 次迭代训练，然后以 512 的序列长度进行了 100k 迭代训练。
( P( s# |5 X6 |, k; z8 i. P0 o16 个 TPUv3 芯片的总训练时间约为 3 天。我们使用 SQuAD-v1 的 F1 分数作为精度指标。F1 得分越高，准确度越高。斯坦福问答数据集 (SQuAD) 是一个阅读理解数据集，包含众包工作者从维基百科的文章中提出的问题，每一个问题的答案都是对应阅读文章的一段文字，或者该问题无法回答。我们从 BERT 的公开 GitHub 库上下载了预训练好的模型。
) K7 Q$ x; p" b" X8 t: k1 A+ ~使用作者提供的脚本，baseline 的 F1 得分为 90.395。在我们的代码中，我们使用了BERT 的作者提供的数据集和基线模型，只修改了优化器。通过使用 LAMB 优化器，我们能够在批大小为 32768 的 15625 次迭代中获得 91.460 的 F1 分数 (序列长度为 128的 14063 次迭代和序列长度为 512 的 1562 次迭代)。
我们把训练时间从 3 天减少到100 分钟左右。我们将批大小推到了 TPU Pod 的硬件极限。批大小大于 32768 时 (序列长度为 512) 将导致 TPU Pod 耗尽内存。
我们实现了 76.7% 的弱扩展效率 (49.1 倍的加速，64 倍的计算资源)。由于我们在 TPU Pod 上使用同步数据并行来进行分布式训练，因此在互连上传输梯度会带来通信开销。梯度的大小与训练后的模型相同。
7 V) M; d* X0 W6 G0 LMixed-Batch 训练
6 A* M8 G& ]; d7 z. W1 f0 K) S如前所述，BERT 预训练包括两个阶段：
( S, M: t$ b8 x6 C(1) 前 9/10 的 epoch 使用 128 的序列长度，
(2) 最后 1/10 的 epoch 使用 512 的序列长度。
对于第二阶段，由于内存限制，TPUv3 Pod 上的最大批大小为 32768，因此我们将第二阶段在批大小达到 32768 时停止。
对于第一阶段，由于内存限制，TPUv3 Pod 上的最大批大小是 131072。但是，当我们将批大小从 65536 增加到 131072 时，并没有看到加速，因此我们在第一阶段批大小达到 65536 时停止。
- J7 c' t$ [- z7 Z此前，Smith 等人也研究了混合批训练。但是，他们在训练中增大了批大小，而我们减小了批大小。
4 h: m: L$ n. o5 y2 F% r我们能够从头到尾充分利用硬件资源。Smith 等人的研究只在最后阶段充分利用了硬件资源。增加批大小可以 warm-up 和稳定优化过程，但是减小批大小会给优化过程带来混乱，导致训练不收敛。
& O: I4 @2 h- ?) `! e( [在实验中，我们发现了一种有助于稳定第二阶段优化的方法。由于我们切换到一个不同的优化问题，有必要重新 warm-up 优化过程。在第二阶段，我们没有降低学习率，而是将学习率从零开始增加 (re-warm-up)。
* Y# @6 W! [  L) Z通过这些改变，我们只需要 8599 次迭代，可以在 76 分钟左右完成 BERT 训练，实现了 101.8% 的弱缩放效率 (weak scaling efficiency)，提速 65.2 倍，利用了 64 倍的计算资源。
# r9 A" p- b; A4 @2 m* ~结论
! S8 c% K" f6 m8 [1 ?; xLarge batch 技术是加快神经网络深度训练的关键。在本文中，我们提出了支持adaptive element-wise updating 和 layer-wise correction 的 LAMB 优化器。LAMB是一个通用的优化器，适用于小批量和大批量。通过使用 LAMB，我们可以将 BERT 预训练的 batch size 扩展到 64K，而不会丢失准确性。我们将 BERT 的训练时间从 3 天减少到 76 分钟左右，并将批大小推到了 TPU Pod 的硬件极限。我们正在研究 LAMB优化器的理论分析。
【2019 新智元 AI 技术峰会精彩回顾】
2019 年 3 月 27 日，新智元再汇 AI 之力，在北京泰富酒店举办 AI 开年盛典 ——2019 新智元 AI 技术峰会。峰会以 “智能云・芯世界 “为主题，聚焦智能云和 AI 芯片的发展，重塑未来 AI 世界格局。
同时，新智元在峰会现场权威发布若干 AI 白皮书，聚焦产业链的创新活跃，评述 AI 独角兽影响力，助力中国在世界级的 AI 竞争中实现超越。
! V/ V- ]0 e1 Z0 H* R( H: m
* I- t/ P9 K( e2 t) v1 u. G9 x4 U+ w

来源：http://www.yidianzixun.com/article/0LepFuGW
1 h$ e6 v5 ^7 b+ [: d. V7 v免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！