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

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

Advanced Neural Network Applications参考文献 21被引用 152

一句话总结

本文提出了共蒸馏（codistillation）—一种在线蒸馏变体，在并行训练多个模型并鼓励它们的预测达成一致的基础上，实现超越标准分布式 SGD 的更快扩展，并在不增加测试时成本的情况下提高可重复性。

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 的可扩展训练的动机。
• 提出一种简单的在线蒸馏方法（共蒸馏），允许多个模型并行训练并共享知识。
• 证明共蒸馏在不增加测试时成本的情况下能加快训练并提高准确性。
• 探讨将在现实世界的大规模数据集上应用共蒸馏的实际注意事项和设计选择。

提出的方法

• 在相互独立的数据子集上训练 n 份模型，并加入蒸馏项以鼓励与其他模型的平均预测达成一致。
• 使用其他模型的陈旧预测来计算蒸馏损失，从而实现低通信需求。
• 在初始预热阶段后启用蒸馏，保留基础损失并添加蒸馏项，形成一个组合目标。
• 展示共蒸馏可以通过对工作节点分组并跨组交换检查点来与分布式 SGD 结合。
• 提供实现方面的考虑因素，如通过检查点或预测服务器进行通信，以及对陈旧预测的容忍度。
• 将共蒸馏与集成和离线蒸馏进行比较，聚焦于训练时效率和可重复性。

实验结果

研究问题

• RQ1 在线共蒸馏是否能够通过利用额外并行性，在不需要多阶段蒸馏管线的情况下实现比分布式 SGD 更快的训练？
• RQ2 与传统蒸馏和集成相比，共蒸馏是否能保持或提高最终模型的准确性和可重复性？
• RQ3 陈旧预测和通信策略如何影响共蒸馏在大规模设置中的有效性和实用性？
• RQ4 哪些实际设计选择可以在最大化共蒸馏收益的同时最小化流水线复杂度？

主要发现

• 共蒸馏通过在传统 SGD 出现收益递减时仍能有效利用更多计算资源，实现更快的训练。
• 双向共蒸馏将训练时间大约降低约 2 倍，以达到与基线相同的验证误差，并且可以接近集成的性能。
• 共蒸馏提供与集成类似的可重复性优势，减少预测波动而不增加服务成本。
• 让共蒸馏模型使用不同的数据子集带来比使用相同数据更强的收益，表明跨数据分区的知识转移取得成功。
• 共蒸馏在陈旧预测下仍然有效，并且可以与同步或异步 SGD 集成，通信成本可控。