Large scale distributed neural network training through online distillation (1804.03235v2)

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 LLMing, containing $6\times 10^{11}$ tokens and based on the Common Crawl repository of web data.

• The paper presents codistillation, an online distillation method that accelerates distributed training by nearly two times relative to standard SGD.
• The approach adjusts each model's loss function to incorporate averaged predictions, reducing training steps by up to 27.6% on datasets like ImageNet.
• Empirical tests on large datasets demonstrate that codistillation enhances both training efficiency and model stability in real-world applications.

An Evaluation of Large-Scale Distributed Neural Network Training via Online Distillation

This paper presents an investigation into the use of a variant of knowledge distillation, termed codistillation, to facilitate large-scale distributed neural network training. Its primary focus is the discovery that online distillation enables efficient model training on substantial datasets, in contrast to the limitations encountered with distributed stochastic gradient descent (SGD).

The authors posit that when parallelism in existing distributed SGD implementations is pushed to its limits, online distillation can significantly speed up training by around two times. This remark is supported through experiments on extensive datasets such as the Criteo Display Ad Challenge dataset, ImageNet, and a Common Crawl-based neural LLMing dataset, demonstrating notable improvements in both training time and model stability.

The Codistillation Methodology

Codistillation differs from traditional distillation as it doesn’t necessitate a multi-phase setup. The concurrent training of $n$ model replicas includes a modification to each replica's loss function to match average predictions from other models. This system provides an avenue to leverage additional computational resources effectively, even when conventional parallelism with SGD yields no extra benefits.

Critical to this approach is codistillation's efficiency: it harnesses delays better than comparable algorithms, allowing for more tolerant use of stale teacher predictions without degrading result quality. This adaptive characteristic is accomplished by allowing distributed workers to periodically share computational checkpoints and use them to refine training through the codistillation protocol.

Empirical Validation and Comparative Results

The paper includes empirical tests which illustrate that codistillation's use in distributed training significantly reduces the necessary training steps relative to traditional SGD methods. For instance, in a Common Crawl LLMing exercise, two-way codistillation cut required training steps in half when compared to the optimal throughput achieved by 128-worker synchronous SGD, maintaining competitive accuracy levels to that of a two-way ensemble.

On ImageNet, the employment of codistillation enabled achieving 75% top-1 accuracy in significantly fewer steps — reducing steps by 27.6% in comparison to a non-codistilled model. Thus, codistillation emerges as a scalable alternative due to its efficiency in communication overhead and computation use.

Implications and Future Research Directions

The insights drawn from this research delineate practical ramifications for deploying distributed neural networks in industrial applications. Codistillation successfully circumvents some scalability barriers presented by traditional asynchronous and synchronous distributed SGD procedures, especially for models bound by infrastructure constraints and needing efficient communication strategies for model updates.

Theoretical avenues deserve exploration as the "teaching" model's role in influencing outcomes becomes more salient. Understanding how suboptimal models codistill to improve collective network accuracy should raise pertinent questions about model architecture and training method adaptations. Additionally, exploring alternative network topologies and even the impact of quantizing predictions shared across learners in the codistillation framework could hold promise for enhancing process efficiency and predicting training trajectory outcomes.

In synthesis, the proposal and analysis of codistillation as a distributed training tool reflect a progressive step towards addressing prevalent challenges in scaling neural network training, offering an adaptable and computationally effective strategy within expansive data environments.