Auto Byte

专注未来出行及智能汽车科技

微信扫一扫获取更多资讯

Science AI

关注人工智能与其他前沿技术、基础学科的交叉研究与融合发展

微信扫一扫获取更多资讯

作者:Yang You、Jing Li等

76分钟训练BERT!谷歌大脑新型优化器LAMB加速大批量训练

去年,谷歌发布了基于双向 Transformer 的大规模预训练语言模型 BERT 并开源。该模型参数量非常大——3 亿,训练时间也很长。近日,来自谷歌大脑的研究者提出了一种新型优化器 LAMB,有效降低 BERT 预训练时间,76 分钟即可完成 BERT 预训练!

尽管 BERT效果惊人,但它所需的计算量非常大,原作者在论文中也表示每次只能预测 15% 的词,因此模型收敛得非常慢。如果我们想保留这种 Mask 机制,那么就需要寻找另一种加速方法了。

当时,BERT 的作者在 Reddit 上也表示预训练的计算量非常大,Jacob 说:「OpenAI 的 Transformer 有 12 层、768 个隐藏单元,他们使用 8 块 P100 在 8 亿词量的数据集上训练 40 个 Epoch 需要一个月,而 BERT-Large 模型有 24 层、2014 个隐藏单元,它们在有 33 亿词量的数据集上需要训练 40 个 Epoch,因此在 8 块 P100 上可能需要 1 年?16 Cloud TPU 已经是非常大的计算力了。」

而在谷歌大脑的这篇新论文中,研究者提出新型优化器 LAMB,通过使用 65536/32768 的批量大小,他们只需要 8599 次迭代、76 分钟就能完成 BERT 预训练。总体而言,相比基线 BERT-Large 用 16 块 TPU 芯片,LAMB 训练 BERT-Large 用了一个 TPU v3 Pod(1024 块 TPU 芯片),因此时间也由 3 天降低为 76 分钟。

不过在 BERT 原论文中,训练 BERT-Large 使用了 64 块 TPU 芯片(16 Cloud TPU),它需要花四天时间完成训练。

论文:Reducing BERT Pre-Training Time from 3 Days to 76 Minutes

论文地址:https://arxiv.org/abs/1904.00962

摘要:大批量训练是加速大型分布式系统中深度神经网络训练的关键。但是,大批量训练难度很大,因为它会产生泛化差距(generalization gap),直接优化通常会导致测试集准确率受损。BERT [4] 是当前最优的深度学习模型,它基于用于语言理解的深度双向 transformer 而构建。当我们扩展批量大小时(比如批量大小超过 8192),之前的大批量训练技术在 BERT 上的执行性能并不好。BERT 预训练需要大量时间(使用 16 个 TPUv3 训练 3 天)。

为了解决这个问题,来自谷歌大脑的研究者提出了一种新型优化器 LAMB,可在不损害准确率的情况下将批量大小扩展至 65536。LAMB 是一款通用优化器,它适用于小批量和大批量,且除了学习率以外其他超参数均无需调整。基线 BERT-Large 模型的预训练需要 100 万次迭代,而 LAMB 使用 65536/32768 的批量大小,仅需 8599 次迭代。研究者将批量大小扩展到 TPUv3 pod 的内存极限,在 76 分钟内完成了 BERT 的训练。

具体来讲,LAMB 优化器支持自适应元素级更新(adaptive element-wise updating)和准确的逐层修正(layer-wise correction)。LAMB 可将 BERT 预训练的批量大小扩展到 64K,且不会造成准确率损失。BERT 预训练包括两个阶段:1)前 9/10 的训练 epoch 使用 128 的序列长度,2)最后 1/10 的训练 epoch 使用 512 的序列长度。

基线 BERT-Large 模型的预训练需要 100 万次迭代,研究者仅使用 8599 次迭代就完成了预训练,将训练时间从 3 天缩短到 76 分钟。该研究使用的训练批量大小接近 TPUv3 pod 的内存极限。LAMB 优化器可以将批量大小扩展到 128k 甚至更大,由于硬件限制,研究者在序列长度为 512 时使用的批量大小为 32768,在序列长度为 128 时使用的批量大小为 65536。该论文中的 BERT 模型均指 BERT-Large。为公平起见,研究中所有实验均运行同样数量的 epoch(即固定数量的浮点运算)。实验结果见下表。

表 1:该研究使用 SQuAD-v1 的 F1 分数作为准确率度量。

在上表中,基线 F1 分数来自于 BERT 开源 GitHub 中 BERT-Large 预训练模型的对应分数。实验所用硬件为 TPUv3,实验设置与基线 BERT-Large 相同:前 9/10 的训练 epoch 使用 128 的序列长度,最后 1/10 的训练 epoch 使用 512 的序列长度。所有实验均运行同样数量的 epoch。

LAMB (Layer-wise Adaptive Moments optimizer for Batch training)

BERT 基线模型的训练使用 Adam with weight decay(Adam 优化器的变体)作为优化器 [15]。另一个成功用于大批量卷积神经网络训练的自适应优化器是 LARS [21]。这些优化器启发了研究者提出用于大批量 BERT 训练的新型优化器 LAMB。LAMB 优化器详见算法 1。

实验

常规训练

TPU 是强大的浮点运算计算硬件。研究者在所有实验中均使用 TPUv3。每个 TPUv3 pod 拥有 1024 个芯片,可提供超过 100 petaflops 的混合精度计算。实验结果见表 1。基线模型预训练过程中使用了 Wikipedia 和 BooksCorpus 数据集。研究者使用与开源 BERT 模型相同的数据集进行预训练,即包含 2.5B 单词的 Wikipedia 数据集和包含 800M 单词的 BooksCorpus 数据集。

BERT 作者首先以 128 的序列长度进行 900k 次迭代,然后以 512 的序列长度进行 100k 次迭代。在 16 块 TPUv3 上训练的总时间大约为 3 天。该研究使用 SQuAD-v1 的 F1 分数作为准确率度量。较高的 F1 分数意味着较好的准确率。研究者下载了 BERT 开源项目中提供的预训练模型。使用 BERT 作者提供的脚本,基线模型得到了 90.395 的 F1 分数。

该研究中,研究者使用 BERT 作者提供的数据集和基线模型,仅改变了优化器。使用新型优化器 LAMB 后,研究者以 32768 的批量大小进行了 15625 次迭代,得到了 91.460 的 F1 分数(用于序列长度 128 的迭代有 14063 次,用于序列长度 512 的迭代有 1562 次)。研究者将 BERT 训练时间从 3 天缩短到大约 100 分钟。

该研究取得了 76.7% 的弱可扩展性效率(weak scaling efficiency)。研究者在 TPU Pod 上使用了分布式训练的同步数据并行化,因此梯度迁移会产生通信开销。这些梯度与训练后的模型大小一样。在 ImageNet 数据集上训练 ResNet-50 时的弱可扩展性效率可达到 90+%,因为 ResNet-50 比 BERT参数少得多(25 million vs 300 million)。LAMB 优化器无需用户调整超参数,用户只需输入学习率即可。

混合批次训练(Mixed-Batch Training)

如前所述,BERT 预训练主要分为两部分,1)前面 9/10 的 Epoch 使用 128 的序列长度;2)最后 1/10 的 Epoch 使用 512 的序列长度进行训练。对于第二阶段而言,因为内存限制,TPUv3 Pod 上最大的批量大小为 32768,因此第二阶段使用的批大小为 32768。对于第一阶段,受限于内存,TPUv3 Pod 上最大的批量大小为 131072。然而,研究者将批大小从 65536 增加到 131072 时,并没有发现明显的加速,因此研究者将阶段 1 的批大小定为 65536。

此前,也有其它研究者讨论了混合批次训练,然而他们是在训练中增大批量大小;而本研究是降低批大小,这样他们能从开始到结束充分利用硬件资源。增大批大小能起到 warm-up 初始训练和稳定最优化过程的效果,但是降低批大小可能会带来最优化过程的紊乱,并导致训练的不收敛

在研究者的实验中,他们发现有一些技术能稳定第二阶段的最优化过程。因为这两阶段会切换到不同的最优化问题,因此有必要重新 warm-up 最优化过程。研究者没有在第二阶段进行学习率衰减,而是将学习率从零开始增加(re-warm-up)。和第一阶段一样,研究者在 re-warm-up 之后执行学习率衰减。因此,研究者能够以 8599 次迭代完成 BERT 训练,且仅用时 76 分钟,达到了 101.8% 的弱可扩展性效率。

理论BERT语言模型分布式计算
6
相关数据
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

基于Transformer 的双向编码器表征技术

BERT是谷歌发布的基于双向 Transformer的大规模预训练语言模型,该预训练模型能高效抽取文本信息并应用于各种NLP任务,并刷新了 11 项 NLP 任务的当前最优性能记录。BERT的全称是基于Transformer的双向编码器表征,其中“双向”表示模型在处理某一个词时,它能同时利用前面的词和后面的词两部分信息。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

收敛技术

在数学,计算机科学和逻辑学中,收敛指的是不同的变换序列在有限的时间内达到一个结论(变换终止),并且得出的结论是独立于达到它的路径(他们是融合的)。 通俗来说,收敛通常是指在训练期间达到的一种状态,即经过一定次数的迭代之后,训练损失和验证损失在每次迭代中的变化都非常小或根本没有变化。也就是说,如果采用当前数据进行额外的训练将无法改进模型,模型即达到收敛状态。在深度学习中,损失值有时会在最终下降之前的多次迭代中保持不变或几乎保持不变,暂时形成收敛的假象。

学习率技术

在使用不同优化器(例如随机梯度下降,Adam)神经网络相关训练中,学习速率作为一个超参数控制了权重更新的幅度,以及训练的速度和精度。学习速率太大容易导致目标(代价)函数波动较大从而难以找到最优,而弱学习速率设置太小,则会导致收敛过慢耗时太长

超参数技术

在机器学习中,超参数是在学习过程开始之前设置其值的参数。 相反,其他参数的值是通过训练得出的。 不同的模型训练算法需要不同的超参数,一些简单的算法(如普通最小二乘回归)不需要。 给定这些超参数,训练算法从数据中学习参数。相同种类的机器学习模型可能需要不同的超参数来适应不同的数据模式,并且必须对其进行调整以便模型能够最优地解决机器学习问题。 在实际应用中一般需要对超参数进行优化,以找到一个超参数元组(tuple),由这些超参数元组形成一个最优化模型,该模型可以将在给定的独立数据上预定义的损失函数最小化。

运筹优化技术

最优化问题(英语:Optimization problem)在数学与计算机科学领域中,是从所有可行解中寻找最优良的解的问题。根据变数是连续的或离散的,最佳化问题可分为两类:连续最佳化问题与组合优化。

卷积神经网络技术

卷积神经网路(Convolutional Neural Network, CNN)是一种前馈神经网络,它的人工神经元可以响应一部分覆盖范围内的周围单元,对于大型图像处理有出色表现。卷积神经网路由一个或多个卷积层和顶端的全连通层(对应经典的神经网路)组成,同时也包括关联权重和池化层(pooling layer)。这一结构使得卷积神经网路能够利用输入数据的二维结构。与其他深度学习结构相比,卷积神经网路在图像和语音识别方面能够给出更好的结果。这一模型也可以使用反向传播算法进行训练。相比较其他深度、前馈神经网路,卷积神经网路需要考量的参数更少,使之成为一种颇具吸引力的深度学习结构。 卷积网络是一种专门用于处理具有已知的、网格状拓扑的数据的神经网络。例如时间序列数据,它可以被认为是以一定时间间隔采样的一维网格,又如图像数据,其可以被认为是二维像素网格。

准确率技术

分类模型的正确预测所占的比例。在多类别分类中,准确率的定义为:正确的预测数/样本总数。 在二元分类中,准确率的定义为:(真正例数+真负例数)/样本总数

批次技术

模型训练的一次迭代(即一次梯度更新)中使用的样本集。

深度神经网络技术

深度神经网络(DNN)是深度学习的一种框架,它是一种具备至少一个隐层的神经网络。与浅层神经网络类似,深度神经网络也能够为复杂非线性系统提供建模,但多出的层次为模型提供了更高的抽象层次,因而提高了模型的能力。

优化器技术

优化器基类提供了计算梯度loss的方法,并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法,如梯度下降和Adagrad。 优化器是提供了一个可以使用各种优化算法的接口,可以让用户直接调用一些经典的优化算法,如梯度下降法等等。优化器(optimizers)类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类,但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer(tensorflow下的优化器包)等等这些算法。

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用,如语音识别,机器翻译,词性标注,句法分析和资讯检索。由于字词与句子都是任意组合的长度,因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题),也使得在语料库中估算字串的机率变得很困难,这也是要使用近似的平滑n元语法(N-gram)模型之原因。

推荐文章
暂无评论
暂无评论~