测量神经网络数据并行训练的极限

lixia01

在过去十年中，神经网络在各种各样的预测任务中都取得了最先进的成果，包括图像分类、机器翻译和语音识别。这些成功至少在一定程度上得益于硬件和软件的改进，这些改进显著加速了神经网络的训练。更快的训练直接导致了模型质量的显著提高，既可以处理更多的训练数据，也允许研究人员更快地尝试新的想法和配置。今天，像Cloud TPU Pods这样的硬件发展正在迅速增加可用于神经网络训练的计算量，这增加了利用额外计算使神经网络训练更快并促进模型质量更大改进的可能性。但我们究竟应该如何利用这种前所未有的计算量，我们是否应该总是期望更多的计算来促进更快的训练？
利用大规模计算能力的最常见方法是将计算分配到不同的处理器之间并同时执行这些计算。在训练神经网络时，实现此目标的主要方法是模型并行（涉及将神经网络分布在不同的处理器上）和数据并行（涉及将训练示例分布在不同的处理器上并并行计算对神经网络的更新）。虽然模型并行可以训练大于单个处理器可以支持的神经网络，但它通常需要根据可用的硬件定制模型架构。相比之下，数据并行与模型无关，适用于任何神经网络架构——它是并行化神经网络训练最简单、使用最广泛的技术。对于最常见的神经网络训练算法（同步随机梯度下降及其变体），数据并行的规模对应于批处理大小，即用于计算对神经网络的每次更新的训练示例数量。但是这种并行化的极限是什么，什么时候我们应该期望看到大幅加速？在“测量数据并行在神经网络训练中的影响
”中“，我们通过使用三种不同的优化算法（“优化器”）在七个不同数据集上的六种不同类型的神经网络上运行实验，研究了批次大小和训练时间之间的关系。总的来说，我们在约 450 个工作负载上训练了超过 100K 个独立模型，并观察到在我们测试的所有工作负载中，批次大小和训练时间之间存在看似普遍的关系。我们还研究了这种关系如何随数据集、神经网络架构和优化器而变化，并发现工作负载之间存在极大的差异。此外，我们很高兴与研究界分享我们的原始数据以供进一步分析。这些数据包括超过 7100 万个模型评估，构成了我们训练的所有 100K+ 个独立模型的训练曲线，可用于重现我们论文中的所有 24 个图。
批次大小和训练时间之间的普遍关系在理想的数据并行系统中，处理器之间同步的时间可以忽略不计，训练时间可以通过训练步骤
的数量（对神经网络参数的更新）来衡量。在此假设下，我们观察到批量大小和训练时间之间关系的三种不同的扩展机制：一种是“完美扩展”机制，其中批量大小加倍会使达到目标样本外误差所需的训练步数减半；其次是“收益递减”机制；最后是“最大数据并行”机制，其中进一步增加批量大小不会减少训练时间，即使假设硬件理想化。
虽然批次大小和训练时间之间的基本关系似乎是普遍的，但我们发现不同扩展方案之间的转换点在不同的神经网络架构和数据集中存在巨大差异。这意味着，虽然简单的数据并行性可以在当今硬件的极限（例如Cloud TPU Pods）下为某些工作负载提供大幅加速，甚至超越这一极限，但某些工作负载需要超越简单的数据并行性才能从当今存在的最大规模硬件中受益，更不用说尚未构建的硬件了。例如，在上图中，CIFAR-10上的ResNet-8无法从大于 1,024 的批次大小中受益，而ImageNet上的ResNet-50 则可以继续从将批次大小增加到至少 65,536 中受益。优化工作负载 如果可以预测哪些工作负载从数据并行训练中受益最多，那么就可以定制他们的工作负载以最大限度地利用可用的硬件。然而，我们的结果表明，这通常并不简单，因为最大可用批大小至少在某种程度上取决于工作负载的各个方面：神经网络架构、数据集和优化器。例如，即使在使用相同优化器对相同数据集进行训练时，某些神经网络架构也可以从比其他架构更大的批大小中受益。虽然这种影响有时取决于网络的宽度和深度，但不同类型的网络之间不一致，有些网络甚至没有明显的“宽度”和“深度”概念。虽然我们发现某些数据集可以从比其他数据集更大的批大小中受益，但这些差异并不总是由数据集的大小来解释——有时较小的数据集比较大的数据集从较大的批大小中受益更多。
也许我们最有希望的发现是，即使对优化算法进行微小的更改（例如允许随机梯度下降中的动量），也可以显著改善训练在增加批量大小时扩展的效果。这增加了设计新优化器或测试我们未考虑的优化器的扩展属性的可能性，以找到可以最大程度利用增加的数据并行性的优化器。
未来工作
通过增加批量大小来利用额外的数据并行性是一种在一系列工作负载中产生有价值的加速的简单方法，但是，对于我们尝试的所有工作负载，在最先进硬件的限制范围内，好处会减少。但是，我们的结果表明，一些优化算法可能能够在许多模型和数据集中一致地扩展完美的扩展机制。未来的工作可以对我们尝试的几个密切相关的优化器以外的其他优化器进行相同的测量，看看是否有任何现有的优化器在许多问题上扩展了完美的扩展。
致谢
本研究的作者包括 Chris Shallue、Jaehoon Lee、Joe Antognini、Jascha Sohl-Dickstein、Roy Frostig 和 George Dahl（Chris 和 Jaehoon 贡献相同）。许多研究人员已经在此领域开展了工作，我们在此基础上进行了研究，因此请参阅我们的论文以了解相关工作的完整讨论。