Scaling Neural Machine Translation

深度学习模型训练往往需要几天或是更多的时间才能在大型基准数据集上获得最优性能，本文提出了一种半精度及大批量训练方法，在单机8卡上实验训练速度提升了将近5倍，在wmt14英德实验中只用了5个小时即复现了transformer实验结果。!

论文地址:
https://arxiv.org/pdf/1806.00187.pdf
代码地址:
https://github.com/pytorch/fairseq/tree/master/examples/scaling_nmt

引言

近年来，随着很多高效结构的引入，神经机器翻译(nmt)取得了令人瞩目的进展。计算密集型模型采用单机训练往往需要数天才能达到理想性能，即使扩展到多机训练，也面临着多个挑战：如何保证分布式训练的效率，更大批量的训练是否会影响到优化及泛化性能。本文将针对上述挑战提出针对性的改进方法。

模型

半精度训练

深度学习模型常常包含数百万甚至是上亿个参数，一次训练时间耗时很长，这使得模型优化及参数调优或是推理都受到了限制。张量核(tensor core)的引入，使得半精度计算(fp16 computation)相对于单精度计算(fp32 computation)速度快了几倍。但是，半精度缩小了浮点数表示范围，可能会导致参数数值上溢(overflows)或是下溢(underflows)，比较好的办法是将参数数值调整缩放到fp16表示范围之内。为了保证模型精度和训练速度，即采用fp16执行前向及反向传播操作，但是权重更新和损失采用fp32。具体实现如下：

其中对loss做scale操作的原因时因为在梯度计算过程中，梯度取值范围在fp16表示范围之外，如图1所示，因此需要对梯度值做平移，使得大部分梯度值能够处于fp16表示范围之内。一个简单的方式即是反向传播之前调整loss，按照链式法则，反向传播确保所有梯度值按照相同的比例缩放，这样就避免了部分梯度值下溢或是上溢的情况。在权重更新之前，必须对权重梯度做unscale，以保持在fp32中训练的权重大小。在训练初始阶段，loss比较大，需要对减小loss保证梯度值不上溢；在训练结尾阶段，loss比较小，需要扩大loss，保证梯度值不下溢。

表1中给出了wmt14英德实验结果，采用半精度训练方式，在其他参数保持一致的情况下，训练速度提升了2.9倍，bleu值提升了0.3。

大批量训练

论文中采用累计梯度更新的方法实现了大批量训练，具体来说就是在权重更新之前迭代多次，平均多次迭代的loss，然后进行权重更新。这样做的好处之一削弱了gpu之间的通信导致的效率低下问题，同时，能够保证各worker之间处理的数据尽可能趋于同分布。

如表1所示，基于上文中半精度训练方式，采用累计梯度更新的方法，每16个batch梯度更新一次，速度相对于reference提高了3.2倍；进一步的，当学习率调整为原来2倍时，在尽量不损失精度的情况下，速度相对于reference提升了4.6倍。

并行训练

当使用多节点训练时，由于gpu之间的通信，训练时间会显著增加。反向传播的序列特性可以使得我们能够进一步提高多节点训练性能。反向传播顺序是从网络顶部向下传播至输入，因此，论文中提出在部分网络层梯度计算完成之后，将梯度添加至缓存区。一旦缓存区梯度大小达到预定义的阈值(论文中设置为150M)，后台线程即可同步缓存梯度。这种方法实现了gpu通信与反向传播的并行运行。表一实验中设置了在128卡，2倍学习率情况下，速度相对于reference提升了38.6倍，当采用并行训练方式时，速度提升为原来的44.7倍。

实验结果

论文设置了英德和英法两组实验，采用wmt14 newstest为验证数据，实验结果如下：

论文中通过优化后的训练方式，英德实验仅耗时85分钟，bleu值为29.3(baseline : 29.2)；英法实验仅耗时512分钟，bleu值为43.2(baseline: 41.9)。在加入4亿爬虫数据后，英德实验耗时539分钟，bleu值提升了0.5, 英法实验耗时794分钟，但是bleu值降低了1.1。

结论

本文探索了如何训练最先进的nmt模型，提出了三种加速训练的训练方式，半精度训练，大批量训练(累计梯度更新)，并行训练。分别了设置了英德英法两组实验，在不损失模型精度或是显著提升模型性能的情况下，大大提高了训练速度。