[Paper Review] 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 Classification206 citations

TL;DR

Codistillation trains multiple model copies in parallel by a distillation term using stale predictions, enabling faster training beyond SGD limits and improving reproducibility without extra test-time costs. It is validated on Common Crawl language modeling, ImageNet, and 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.

Motivation & Objective

• Motivate scalable training for large neural networks beyond the practical limits of distributed SGD.
• Introduce codistillation, an online distillation variant that uses simultaneous training of multiple models.
• Show that codistillation speeds up training by leveraging extra parallelism without adding test-time costs.
• Demonstrate improved reproducibility and reduced prediction churn compared to ensembles or offline distillation.
• Provide practical guidance on design choices and implementation considerations for codistillation.

Proposed method

• Train n copies of a model in parallel on locally partitioned data with no centralized gradient sharing.
• Add a distillation loss term to each model's objective that encourages agreement with the average predictions of the other models.
• Enable the distillation term after an initial burn-in period to preserve model diversity.
• Optionally combine codistillation with standard distributed SGD, exchanging checkpoints between groups.
• Discuss alternatives such as a prediction server to exchange predictions instead of weight checkpoints.
• Show that codistillation is robust to using stale predictions and requires minimal additional communication.

Experimental results

Research questions

• RQ1 Can online codistillation accelerate training beyond what is achievable with distributed SGD alone?
• RQ2 Does codistillation preserve or improve final model accuracy compared to baselines like SGD, label smoothing, or ensembles?
• RQ3 How does the use of stale predictions affect training stability and final performance?
• RQ4 Can codistillation reduce prediction churn across retrains and versions of a model?
• RQ5 What practical design choices (data partitioning, checkpoint exchange frequency) maximize benefits of codistillation?

Key findings

• Two-way codistillation with 128 GPUs reduces training steps by about half to reach the same validation error as the baseline, and can reach lower final error.
• On Common Crawl language modeling, two-way codistillation approaches the training curve of a two-way ensemble, achieving similar or better accuracy in roughly half the steps.
• On ImageNet, two-way codistillation reaches 75% accuracy in 5250 steps versus 7250 steps for the baseline, confirming reduced training steps.
• Codistillation tolerates stale predictions well; increasing checkpoint reload intervals up to a point yields only mild degradation.
• Prediction churn is reduced by 35% with codistillation, achieving reproducibility similar to ensembling without increasing serving costs.
• Using different data subsets between codistilling models yields greater gains than using the same data, indicating effective information sharing about diverse data portions.