Large scale distributed neural network training through online distillation

Rohan Anil; Gabriel Pereyra; Alexandre Passos; Robert Ormandi; George E. Dahl; Geoffrey E. Hinton

Source-linked AI summary

Large scale distributed neural network training through online distillation

Rohan Anil, Gabriel Pereyra, Alexandre Passos, Robert Ormandi, George E. Dahl, Geoffrey E. Hinton

arXiv:1804.03235v2Apr 9, 2018cs.LGcs.AIstat.ML

TL;DR

The paper addresses the difficulty of using ensembles and distillation at industrial scale because of test-time cost, training complexity, and distributed-SGD scalability limits. It introduces codistillation, which trains model copies simultaneously while exchanging predictive information, including from stale checkpoints. Across large-scale experiments, codistillation is reported to improve accuracy, accelerate training beyond SGD’s scaling limit, and reduce prediction churn, with open questions about model topologies and teacher accuracy.

Problem
Ensembling and distillation can improve model quality, but their test-time cost, pipeline complexity, and limited distributed-SGD scalability hinder practical large-scale use.
Method
Codistillation simultaneously trains same-architecture models on the same dataset, adding a loss that makes each model match the others’ predictions during training.
Results
Codistillation improves accuracy and speeds training beyond distributed SGD’s useful parallelism limit, while retaining ensemble-like reproducibility benefits.
Takeaways & Limitations
Codistillation offers a simpler, communication-efficient way to use additional machines for large-scale training and reduce prediction churn.
Takeaways & Limitations
The paper mainly studies pairs of codistilling models and leaves alternative communication topologies and teacher-prediction accuracy limits unexplored.

Abstract

from arXiv · show

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

1 INTRODUCTION

Large-scale distributed SGD encounters infrastructure and optimization limits that make adding workers progressively less useful. The paper introduces codistillation as a simpler way to use additional parallelism while improving accuracy, training speed, and prediction reproducibility.

Motivation: Distributed synchronous and asynchronous SGD suffer diminishing training-time improvements as more machines are added.Communication overhead, latency tails, batch-size limits, and gradient interference constrain scalability.
Motivation: Ensembling can use extra machines to improve accuracy and prediction stability, but increases test-time cost.This can conflict with latency or other operational cost constraints.
Contribution: Codistillation trains multiple model copies in parallel by matching each model’s predictions to the average predictions of the others.It is presented as a simpler online alternative to ensemble distillation.
Contribution: Compared with distributed SGD, codistillation improves accuracy and speeds training by using computational resources beyond SGD’s useful parallelism limit.The reported benefits include ensemble-like distillation without increasing training time.
Contribution: Codistillation also preserves ensemble-like reproducibility benefits by reducing prediction churn across retrained models.The paper identifies this as useful when testing and launching model versions non-disruptively.
Contribution: The paper’s primary contribution is experimental validation of codistillation at scale, alongside practical design and implementation recommendations.The introduction distinguishes this focus from related simultaneous-distillation algorithms.

2 CODISTILLATION

Codistillation applies distillation during simultaneous training of multiple same-architecture models, allowing them to exchange predictive information with infrequent communication. Its design tolerates stale teacher predictions and can be combined with distributed SGD for communication-efficient scaling.

Concept: Distillation transfers ensemble-like model-quality benefits by training a student to match a teacher’s predictions.The teacher is traditionally trained before the student, but the paper considers simultaneous alternatives.
Codistillation: Codistillation uses the same architecture and dataset for all models while applying the distillation loss before convergence.Its defining characteristic is simultaneous training of each model and its teacher.
Codistillation: The distillation term encourages agreement between model predictions and can use squared-logit error, KL divergence, or another agreement measure.The experiments use cross-entropy error with the teacher predictive distribution as soft targets.
Staleness: Stale predictions can remain useful because weight changes affect only a subset of predictions near convergence, while weights and gradients are not directly comparable across networks.The paper reports little adverse effect from predictions tens of thousands of asynchronous updates or 800k examples old.
Staleness: The tolerance of stale predictions enables a communication strategy that is less intensive than synchronous or asynchronous SGD.Lower-precision stored parameters may also keep codistillation’s additional computational cost close to independent training.
Distributed implementation: Each worker can train on a local data subset, periodically exchange parameter checkpoints, and use the freshest available copies for codistillation.Independent worker groups can exchange checkpoints while sharing gradient information internally.
Distributed implementation: Codistillation requires workers to read other models’ checkpoints only rarely, unlike distributed SGD’s per-iteration model-sized communication.Model parameters can serve as a compressed representation of data-relevant information.

3 EXPERIMENTS AND RESULTS

The experiments evaluate codistillation’s scalability on neural language modeling and validate selected large-scale findings on ImageNet. Smaller experiments separately examine prediction churn and other algorithmic properties, while scalability tests deliberately reach distributed SGD’s limits.

Experimental design: Neural language modeling provides a representative, expensive large-scale test bed because web text offers vast quantities of training data.The paper also uses ImageNet to test whether results depend on peculiarities of language modeling.
Experimental design: Scalability experiments use datasets large enough that expressive models cannot feasibly reach convergence with existing SGD parallelization strategies.This design makes potential scaling improvements more visible.
Experimental design: Smaller experiments study prediction churn and other codistillation properties, while scalability experiments must reach distributed SGD’s limits.The paper distinguishes these goals because they require different experimental scales.

3.1 DATA SETS AND MODELS

The experiments use Common Crawl language modeling and ImageNet classification, with implementation choices tailored to large-scale distributed training. Common Crawl provides a massive filtered text corpus, while ImageNet follows a synchronous SGD baseline setup.

Common Crawl: Common Crawl preprocessing retained roughly 915 million English documents totaling 20TB of text.Documents were selected for long paragraphs to preserve somewhat longer-range dependencies.
Common Crawl: Common Crawl experiments did not use the entire filtered dataset because training large language models on all of it remained infeasible.
Common Crawl: The Common Crawl language model uses the ADAM optimizer in all experiments.
ImageNet: ImageNet experiments closely follow Goyal et al.’s setup, using fully synchronous SGD and a batch size of 16384 baseline configuration achieving 75% accuracy.The cited setup uses the same learning-rate scaling and schedule as the reference configuration.
Criteo: The Criteo benchmark contains roughly 43 million examples with 13 integer and 26 categorical features for binary click-through-rate prediction.The model is a ReLU feedforward network with hidden layers of 2560, 1024, and 256 units, trained with Adagrad at learning rate 0.001.

3.2 REACHING THE LIMITS OF DISTRIBUTED SGD FOR TRAINING RNNS ON COMMON CRAWL

On the Common Crawl language-modeling task, synchronous SGD gained from more workers only up to 128 GPUs, establishing the baseline limit for further scalability comparisons.

Asynchronous SGD: Asynchronous SGD with 32 and 128 workers was difficult to stabilize, with large-worker RNN training tending to diverge.The experiments therefore focused on synchronous SGD to reduce dependence on infrastructure and implementation specifics.
Synchronous SGD scaling: Increasing synchronous SGD workers reduced steps to best validation error only through 128 workers, with no further improvement at 256 workers.At 256 workers, step time degraded substantially, making the additional workers counterproductive in practice.
Synchronous SGD scaling: The practical scalability limit arose primarily from diminishing returns caused by increasing effective batch size, rather than step-time degradation alone.
Baseline: Synchronous SGD with 128 workers was the strongest tested baseline for training time and final accuracy.Subsequent experiments compare against 128-worker synchronous SGD and use synchronous SGD within codistillation.

3.3 CODISTILLATION WITH SYNCHRONOUS SGD

Two-way codistillation extends synchronous training beyond the 128-GPU SGD limit by exchanging model knowledge, while improving learning curves on Common Crawl and ImageNet.

Common Crawl: Two-way codistillation uses two groups of 128 GPUs running synchronous SGD and periodically exchanging checkpoints.This configuration targets the extra GPUs that did not improve basic 128-GPU synchronous SGD training time.
Common Crawl: Codistillation is compared with synchronous SGD, label smoothing baselines, and a two-network ensemble.The label-smoothing comparisons test whether codistillation’s gains can be explained by smoothing output distributions.
Common Crawl: 2X fewer steps are needed for codistillation to reach the baseline’s best validation error on Common Crawl.Codistillation also reaches a lower final error, and its implementation adds no step time for the tested model because computation overlaps.
ImageNet: 5250 steps reach 75% accuracy with ImageNet codistillation, compared with 7250 steps for the baseline.At 7250 steps, codistillation reaches 75.6% accuracy.

3.4 PREDICTION STALENESS SENSITIVITY

Codistillation remains useful with delayed checkpoint exchange and reaches comparable validation error faster than multi-phase distillation. Its practical scope is bounded by incomplete coverage of very large datasets.

Prediction Staleness Sensitivity: 50-step checkpoint exchanges remain feasible, while intervals beyond 50 steps slightly degrade the learning curve.The experiments tested exchange delays of 50, 100, and 250 steps on Common Crawl.
Codistillation vs Multi-Phase Distillation Variants: 10K steps were sufficient for two-way codistillation to reach roughly the validation error achieved by a 27K-step multi-phase procedure.The multi-phase procedure used 18K ensemble-training steps followed by 9K distillation steps.
Codistillation vs Multi-Phase Distillation Variants: Codistillation reduces training-pipeline complexity by using the same architecture and dataset for all models, avoiding teacher selection and cross-phase dependencies.This symmetry also simplifies tuning, rollback, and reproducibility relative to traditional multi-phase distillation.
Codistillation vs Multi-Phase Distillation Variants: 2.3% of Common Crawl was visited even with codistillation and 256 GPUs, leaving the experiments far from full dataset coverage.The passage frames larger models and more data as a possible route to better results.

3.5 REDUCING PREDICTION CHURN WITH CODISTILLATION

Neural networks with identical architecture and data can produce different predictions despite similar aggregate performance. Codistillation is evaluated as a way to reduce this prediction churn while retaining single-model serving cost.

Reducing Prediction Churn with Codistillation: Similar validation and test performance does not prevent independently trained neural networks from making substantially different predictions and mistakes.Prediction differences can arise from initialization, data order, infrastructure, or small model changes.
Reducing Prediction Churn with Codistillation: Model averaging tends to make predictions more consistent across retrains and minor modifications to base models.This motivates testing whether codistillation provides a similar reproducibility benefit.
Reducing Prediction Churn with Codistillation: 35%: codistillation reduced prediction churn, achieving results similar to ensembling without increasing serving costs.The evaluation measured mean absolute prediction differences between retrains on the Criteo dataset.

4 DISCUSSION AND FUTURE WORK

The discussion presents codistillation as a flexible training-time use of distillation and identifies unresolved design questions. These include broader communication topologies, teacher accuracy, and extracting information beyond predictions.

Discussion and Future Work: Codistillation may support more than pairwise model arrangements, but fully connected and ring topologies remain unexplored.The authors note that fully connected graphs might make models too similar too quickly.
Discussion and Future Work: The required accuracy of teacher predictions was not studied, and aggressive teacher-model quantization is proposed as a possible cost reduction.The proposed direction is to make codistillation nearly as cheap as normal training for very large models.
Discussion and Future Work: The experiments repeat each case five times and report the mean half the range.This describes the reported experimental aggregation procedure.
Discussion and Future Work: Codistillation can sometimes let poorly performing models learn faster than independently trained models, but the ideal teacher properties remain unknown.The authors identify characterizing suitable teacher models as future work.
Discussion and Future Work: The method currently extracts predictions from checkpoints, while richer checkpoint information could potentially provide stronger regularization.The passage also suggests possible applications to federated learning.