让神经网络训练速度加快4倍！谷歌大脑团队提出“数据回送”算法

在摩尔定律的暮色中，GPU 和其他硬件加速器极大地加速了神经网络的训练。但是，训练过程的前期阶段（如磁盘读写和数据预处理）并不在加速器上运行。随着加速器的不断改进，这些前期阶段所花费的时间将逐渐成为训练速度的瓶颈。谷歌大脑团队提出了“数据回送（data echoing）”算法，它减少了训练前期阶段的总计算量，并在加速器上游的计算占用训练时间时加快训练速度。“数据回送”复用训练前期阶段的中间输出，以利用闲置的计算空间。作者研究了在不同任务、不同数量的回送以及不同批尺寸时数据回送算法的表现。在所有的情况下，至少有一种数据回送算法可以用更少的上游计算达到模型的基线性能。在部分情况下，数据回送甚至可以将训练速度提升 4 倍。

在过去的十年中，神经网络的训练速度得到了极大的提升。研究人员能够使用更大的数据集训练更大的模型，并更快地探索新的想法，显著提高了模型的表现性能。随着摩尔定律的失效，通用处理器的速度已经不能满足要求，但特定的加速器通过优化特定操作得到显著的加速。例如，GPU 和 TPU 优化了高度并行的矩阵运算，而矩阵运算正是神经网络训练算法的核心组成部分。

然而，训练一个神经网络需要的不仅仅是在加速器上运行的操作。训练程序需要读取和解压缩训练数据，对其进行打乱（shuffle）、批处理，甚至转换或增强操作。这些步骤需要用到多个系统组件，包括 CPU、磁盘、网络带宽和内存带宽。这些通用操作涉及的组件太多，为它们设计专用的加速硬件是不切实际的。同时，加速器的改进远远超过了通用计算的改进，并且在加速器上运行的代码只占用整个训练时间的一小部分。因此，如果想使神经网络训练地更快，有以下两个方法：（1）使非加速器工作地更快，或（2）减少非加速器所需的工作量。选项（1）可能需要大量的工程工作或技术探索，给系统增加太多的复杂性。因此，作者将重点放在选项（2）上，探索减少训练前期阶段的工作量的方法。

神经网络的训练过程可以看作是一个数据流程，需要对计算进行缓冲和重叠。例如，图 1 显示了小批次随机梯度下降（SGD）及其变体的典型训练流程，这是训练神经网络的标准算法。训练程序首先读取和解码输入数据，然后对数据进行 shuffle，应用一组转换操作来增加数据，并将数据分成不同批次。最后，通过迭代更新网络参数降低损失函数值；此阶段称为“SGD 更新”。由于任何流程阶段的输出都可以缓冲，因此不同阶段的计算互相重叠，最慢的阶段将占用主要的训练时间。

图 1 经典神经网络训练流程图

图 2 数据回送插入点的上游和下游的覆盖计算时间

在这篇论文中，作者研究如何通过减少训练前期部分的时间来加速神经网络训练（图 2a）。作者将训练过程中第一部分的输出重复用于多个 SGD 更新，以利用空闲计算能力。这类算法称为 “数据回送”（data echoing），每个中间输出被使用的次数称为 回送因子（echoing factor）。

数据回送算法在训练流程中的某个位置（SGD 更新之前）插入重复阶段。如果上游任务（重复阶段之前）花费的时间超过下游任务（重复阶段之后）花费的时间，算法将回收下游闲置的计算能力，并提高模型的 SGD 更新率（图 2b）。通过改变插入点、回送因子和 shuffle 程度，可以得到不同的数据回送算法。

要点

1. 在不同的数据集和模型结构上，数据回送在取得竞争性表现的同时减少了上游计算量；
2. 支持大范围的回送因子；
3. 有效性取决于训练流程中的插入点；
4. 数据回送可以从回送后的 shuffle 中获益；
5. 数据回送与调优的基线模型取得相同的错误率。

其中 tupstream 是回送上游所有阶段所用的时间，tdownstream 是回送下游所有阶段所用的时间，e 是回送因子。假设 tupstream≥tdownstream，这是使用数据回送的主要动机。如果用 R=tupstream/tdownstream 表示上下游处理时间之比，那么完成一个上游步骤和 e 个下游步骤的时间对于从 e 到 R 的所有回送因子都是相同的。换句话说，额外的下游步骤是“无消耗的”，因为它们使用的是空闲的下游计算能力。

使用数据回送时，减少的训练时间取决于上游步骤和下游步骤之间的比重。一方面，由于重复数据的价值可能低于新鲜的数据，数据回送可能需要更多的下游 SGD 更新才能达到预期的表现。另一方面，每个下游步骤只需要 1/e（而不是 1）个上游步骤的时间。如果下游步骤增加的数量小于 e，上游步骤的总数（以及总训练时间）将减少。R 代表数据回送的最大可能加速，如果 e=R，并且重复数据和新数据价值相同，则可以实现最大加速。

考虑到训练过程中的每个操作都需要一些时间来执行，如果在 SGD 更新之前应用数据回送，R 能取最大值，但这将导致每个 epoch 多次使用同一批数据。然而，如果想在上游步骤和下游步骤之间找到更有利的权衡，更好的方式是较早的插入回送数据。

不同插入点的数据回送：

批处理前或批处理后的回送： 批处理前的回送是指数据在样本级别而不是批级别进行重复。这样增加了相邻批次不同的可能性，但代价是不能复制一批内的样本。批处理前回送的算法称为 样本回送（example echoing），批处理后回送的算法称为 批回送（batch echoing）。

数据增强前或数据增强后的回送： 数据增强前的回送允许重复数据进行不同的转换，能使重复数据更接近新数据。

数据回送的表现也受回送阶段后 shuffle 程度的影响。在适用的情况下，作者将 shuffle 作为缓冲区。缓冲区越大，shuffle 程度越高，训练算法能够近似于将整个训练集加载到内存中。

（编辑：ASP站长网）