[Paper Review] A Closer Look at Codistillation for Distributed Training

Shagun Sodhani, Olivier Delalleau+4 more|arXiv (Cornell University)|Oct 6, 2020

Topic Modeling20 references4 citations

TL;DR

This paper proposes codistillation as a communication-efficient alternative to standard data-parallel training, where multiple models are trained in parallel and synchronized via knowledge distillation loss rather than parameter averaging. It demonstrates that 2-way codistillation achieves comparable accuracy to standard data parallelism using 2X devices, while reducing communication by up to 1000×, though scaling to n>2 models faces challenges tied to model and dataset characteristics.

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.

Motivation & Objective

• To investigate whether codistillation can serve as a scalable, communication-efficient alternative to standard data-parallel training.
• To evaluate the performance and communication efficiency of 2-way codistillation compared to traditional all-reduce based data parallelism.
• To understand the limitations and challenges of scaling codistillation beyond 2 models (n>2) in distributed training.
• To analyze how codistillation affects training dynamics, including its regularization-like effects.
• To explore the conditions under which n-way codistillation can outperform 2-way codistillation or standard training.

Proposed method

• Codistillation trains multiple models in parallel, each on a subset of data, with a distillation loss that penalizes prediction disagreements between models.
• Instead of synchronizing model parameters via all-reduce, models are synchronized through functional agreement via knowledge distillation.
• The training objective combines standard cross-entropy loss with a symmetric distillation loss that encourages model outputs to match.
• Communication overhead is reduced by transmitting only distilled predictions or soft labels, not full gradients or parameters.
• The method is evaluated on ResNet50 (ImageNet) and a transformer-based machine translation model, with ablation studies on n-way codistillation.
• Hyperparameter α controls the distillation loss weight and is tuned for optimal performance.

Experimental results

Research questions

• RQ1 Can 2-way codistillation achieve comparable accuracy to standard data-parallel training while reducing communication by orders of magnitude?
• RQ2 How does codistillation scale with increasing numbers of devices per model, and does it maintain linear speedup?
• RQ3 Why does scaling codistillation to more than two models (n>2) fail to consistently improve performance across different workloads?
• RQ4 What role does the model architecture and dataset characteristics play in the effectiveness of n-way codistillation?
• RQ5 How does codistillation influence training dynamics, and what regularization effects does it introduce?

Key findings

• 2-way codistillation achieves comparable accuracy to standard data-parallel training when using 2X devices, while reducing communication by up to 1000× per iteration.
• The method enables linear scaling with respect to device count per model, halving training updates needed for convergence when doubling devices.
• Scaling codistillation to n>2 models leads to mixed results: performance improvements are observed in some cases but not consistently, and often no better than 2-way codistillation.
• The challenges in scaling beyond 2 models are strongly tied to model architecture and dataset-specific characteristics, such as generalization behavior and loss landscape.
• Codistillation introduces a regularization effect that influences training dynamics and must be accounted for to achieve optimal performance.
• Despite reduced communication, wall-clock speedup is not guaranteed and depends on model size and communication bandwidth, with larger models benefiting more.