减少模型半数内存用量:百度&英伟达提出混合精度训练法

机器之心编译

选自：arXiv

参与：李泽南、蒋思源

深度学习的模型正在变得越来越复杂，所需要的计算资源也在变得越来越多，在开发更加强大的硬件的同时，很多人也在致力于改进算法。最近，来自百度和英伟达共同提出了一种混合精度训练的新方法，研究者称，通过使用半精度浮点数部分代替单精度，其多种深度学习模型在训练时的内存占用量减少了接近一半，同时训练速度也有所提升。

深度学习已经在很多领域里展现了自己的实力，从图像识别、语言建模、机器翻译到语音识别。目前，技术的发展主要遵从两大方向：使用更大的数据集进行训练或增加模型的复杂度。例如，Hannun 等人在 2014 年提出的语音识别模型经过了 5000 小时的训练，而更最近的声学模型则经过了接近 12,000 小时的训练（Amodei 等人，2016），Soltau 等人则使用了更大的数据集，训练了 125,000 小时。与此同时，Hannun 等人在 2014 年使用了 1100 万个参数，而在 2016 年 Amodei 等人的研究中，双向 RNN 模型增长到了 6700 万个参数，最近的门控循环单元（GRU）模型则更是增加到了 1.16 亿个参数。

越大的模型通常需要在训练时消耗更多的计算资源和内存。这些需求可以通过减少精度表示和计算量来缩减。任何程序的表现（速度），包括神经网络训练和推断都受到三种条件的限制：运算带宽、内存带宽和延迟。降低精度可以解决其中的两个问题。通过使用较少位的数值表示，在处理同样数据时我们需要读/写的内容就更短，内存带宽的压力就变得更小了。与此同时，计算时间也会因为数据的简化而减少。在最近的研究中，半精度（half-precision）方法可以让 GPU 效能提升 2-8 倍（相对于单精度）。在提升速度以外，低精度的格式也减少了训练时的内存使用量。

现代深度学习训练系统通常使用单精度（FP32）格式。在百度和英伟达最近发表的论文中，研究人员试图在降低精度的同时保证模型的准确性。具体来说，研究人员使用 IEEE 半精度格式（FP16）训练了多种神经网络。相对于 FP32 来说，FP16 的动态范围（dynamic range）更窄。为了防止准确度降低，研究人员引入了两种新技术：以 FP32 格式保留原版权重；loss-scaling，将梯度最小化逼近零。使用这些技术，该研究证明了新的方法可以在多种网络架构和应用中训练出准确度和 FP32 格式相同的模型。

深度学习模型的混合精度训练

实验结果包含使用卷积和循环神经网络架构，并训练分类、回归和生成的任务。应用包含图像分类、图像生成、物体检测、语言模型、机器翻译和语音识别。此外，使用新的技术无需改变模型和训练超参数。

图 1. 混合精度（Mixed precision）训练在单层中的迭代。

研究人员在多种深度学习模型上运行了新方法，并与常用方法进行了对比：

• 基线（FP32）：单精度（Single-precision）存储用于激活、权重和梯度。计算也使用单精度单元。

• 混合精度（Mixed Precision/MP）：FP16 用于存储和计算。权重、激活和梯度使用 FP16，FP32 原版用于权重更新。5 个 Loss-scaling 用于一些应用。使用 FP16 算法的实验应用了英伟达新一代芯片中的 Tensor Core，将 FP16 运算累加到 FP32 上，并用于卷积、全连接层和循环层中的矩阵乘法。

基线实验在英伟达 Maxwell 或 Pascal 架构的 GPU 上运行。混合精度实验则使用了最新的 Volta V100 来将 FP16 结果累加至 FP32。其中，混合精度语音识别实验在 FP16 精度上只使用了 Maxwell GPU 以模拟非 Volta 架构上的 Tensor Core 操作。研究人员在这种模式下训练了多种神经网络，以确保模型准确度与 Volta V100 GPU 训练混合精度的结果相同。

图 2. 左图显示了三个实验的结果：基线（FP32）、带 FP32 原版权重的拟 FP16、无 FP32 原版权重的拟 FP16。右图显示了 FP32 汉语语音识别训练权重的直方图。在模型所有层的训练期间，每 4000 次迭代采样一次梯度。

表 1. ILSVRC12 分类任务上，基线与混合精度方法的准确率对比。

研究人员表示，在未来，他们计划将这一方向扩展至生成模型如文本-语音系统和深度强化学习应用。此外，自动 loss-scaling 系数选择也可以进一步简化混合精度训练方法的复杂度。Loss-scaling 系数可由权重梯度的溢出、跳过更新来进行增减。

论文：Mixed Precision Training

论文链接：https://arxiv.org/abs/1710.03740

摘要：深度神经网络已经在各类应用场景中取得了成功。通常，越大的神经网络可以获得越准确的结果。但随着模型尺寸的增长，用于训练的内存和计算资源的需求也随之增加。在本论文中，我们介绍了一种使用半精度浮点数（half precision floating point numbers）训练深度神经网络的新技术。在我们的技术中，权重、激活值和梯度都被以 IEEE 半精度格式存储。与单精度数字相比，半精度浮点数具有较小的数值范围。

我们提出了两种新技术来解决信息丢失的问题。首先，我们提出在每个优化器步之后维护累加梯度权重的单精度原版（FP32）。这种单精度原版在训练中可以转变为半精度格式。其次，我们提出适当减小损失以处理半精度梯度信息丢失的方法。我们证明了这种方法适用于多种模型，包括卷积神经网络、循环神经网络和生成对抗网络。这种技术在超过 1 亿参数，被大数据集训练的大规模模型上非常有效。使用这种方式，我们可以在深度学习模型上减少近乎一半的内存消耗。在未来，我们可以期待半精度硬件单元（half-precision hardware units）带来更多的计算加速效果。