[论文解读] Large scale distributed neural network training through online distillation

Rohan Anil, Gabriel Pereyra+4 more|arXiv (Cornell University)|Apr 9, 2018

Machine Learning and Data Classification被引用 206

一句话总结

Codistillation 通过使用陈旧预测的蒸馏项，在并行训练的多个模型副本时实现更快的训练，超越 SGD 限制并在不增加额外测试时成本的情况下提升可重复性。它在 Common Crawl language modeling、ImageNet 和 Criteo 上得到验证。

ABSTRACT

Techniques such as ensembling and distillation promise model quality improvements when paired with almost any base model. However, due to increased test-time cost (for ensembles) and increased complexity of the training pipeline (for distillation), these techniques are challenging to use in industrial settings. In this paper we explore a variant of distillation which is relatively straightforward to use as it does not require a complicated multi-stage setup or many new hyperparameters. Our first claim is that online distillation enables us to use extra parallelism to fit very large datasets about twice as fast. Crucially, we can still speed up training even after we have already reached the point at which additional parallelism provides no benefit for synchronous or asynchronous stochastic gradient descent. Two neural networks trained on disjoint subsets of the data can share knowledge by encouraging each model to agree with the predictions the other model would have made. These predictions can come from a stale version of the other model so they can be safely computed using weights that only rarely get transmitted. Our second claim is that online distillation is a cost-effective way to make the exact predictions of a model dramatically more reproducible. We support our claims using experiments on the Criteo Display Ad Challenge dataset, ImageNet, and the largest to-date dataset used for neural language modeling, containing $6\ imes 10^{11}$ tokens and based on the Common Crawl repository of web data.

研究动机与目标

• 为超出分布式 SGD 的实际极限的大型神经网络提供可扩展的训练动机。
• 引入 codistillation，一种在线蒸馏变体，使用对多个模型的同步训练。
• 证明 codistillation 通过利用额外的并行性在不增加测试时成本的情况下加速训练。
• 展示相比集成或离线蒸馏，改进的可重复性和降低的预测波动。
• 提供关于 codistillation 的设计选择和实现考虑的实际指南。

提出的方法

• 在本地分区数据上并行训练 n 份模型副本，且不进行集中化梯度共享。
• 在每个模型的目标函数中加入蒸馏损失项，促使其与其他模型的平均预测保持一致。
• 在初始烧入期之后启用蒸馏项，以保留模型多样性。
• 可选地将 codistillation 与标准分布式 SGD 结合，在组之间交换检查点。
• 讨论替代方案，例如使用预测服务器来交换预测而非权重检查点。
• 展示 codistillation 对使用陈旧预测具有鲁棒性，且需要的额外通信很少。

实验结果

研究问题

• RQ1 在线 codistillation 是否能够在仅分布式 SGD 的能力之上加速训练？
• RQ2 相较于 SGD、标签平滑或集成等基线，codistillation 是否能保持或提高最终模型准确性？
• RQ3 使用陈旧预测如何影响训练稳定性和最终性能？
• RQ4 codistillation 能否在再训练和版本之间降低预测波动？
• RQ5 哪些实际设计选择（数据分区、检查点交换频率）能够最大化 codistillation 的收益？

主要发现

• 使用 128 GPU 的双向 codistillation 将训练步骤减少约一半，以达到基线相同的验证误差，并且最终误差可能更低。
• 在 Common Crawl language modeling 上，双向 codistillation 靠近双向集成的训练曲线，在大约一半步骤内实现类似或更好的准确度。
• 在 ImageNet 上，双向 codistillation 在 5250 步达到 75% 的准确率，相较于基线的 7250 步，确认了减少的训练步数。
• Codistillation 对陈旧预测具有良好的容忍度；将检查点重载间隔提高到一定程度只带来轻微的性能下降。
• 使用 codistillation 时预测波动减少 35%，在不增加服务成本的情况下达到与集成相似的可重复性。
• 在 codistilling 模型之间使用不同数据子集相较于使用相同数据，带来更大收益，表明对不同数据部分的信息共享更有效。