A Closer Look at Codistillation for Distributed Training

Shagun Sodhani, Olivier Delalleau, Mido Assran, Koustuv Sinha, Nicolas Ballas, Michael Rabbat

Keywords: Distributed Training, Distillation, Neural Networks, Deep Learning, Large-scale Learning

Abstract: Codistillation has been proposed as a mechanism to share knowledge among concurrently trained models by encouraging them to represent the same function through an auxiliary loss. This contrasts with the more commonly used fully-synchronous data-parallel stochastic gradient descent methods, where different model replicas average their gradients (or parameters) at every iteration and thus maintain identical parameters. We investigate codistillation in a distributed training setup, complementing previous work which focused on extremely large batch sizes. Surprisingly, we find that even at moderate batch sizes, models trained with codistillation can perform as well as models trained with synchronous data-parallel methods, despite using a much weaker synchronization mechanism. These findings hold across a range of batch sizes and learning rate schedules, as well as different kinds of models and datasets. Obtaining this level of accuracy, however, requires properly accounting for the regularization effect of codistillation, which we highlight through several empirical observations. Overall, this work contributes to a better understanding of codistillation and how to best take advantage of it in a distributed computing environment.

One-sentence Summary: We develop new insights about codistillation (a.k.a., "online distillation") and leverage them to show that codistillation can perform as well as models trained with synchronous data-parallel methods, even at moderate batch sizes.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics

Community Implementations: [ 4 code implementations](https://www.catalyzex.com/paper/arxiv:2010.02838/code)

Reviewed Version (pdf): https://openreview.net/references/pdf?id=Z9qhtuBph2