A Closer Look at Codistillation for Distributed Training
Authors: Shagun Sodhani
Olivier Delalleau
Mido Assran
Koustuv Sinha
Nicolas Ballas
Michael Rabbat
Venue:  (2021)
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.

Under review as a conference paper at ICLR 2021
A CLOSER LOOK AT CODISTILLATION FOR DISTRIBUTED TRAINING
Anonymous authors Paper under double-blind review
ABSTRACT
Codistillation has been proposed as mechanism to share knowledge among concurrently trained models by encouraging them to represent the same func- tion through an auxiliary loss. This contrasts with the more commonly used fully-synchronous data-parallel stochastic gradient descent methods, where dif- ferent model replicas average their gradients (or parameters) every iteration and thus maintain identical parameters. We investigate codistillation in a dis- tributed 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 models trained with synchronous data-parallel methods, despite using a much weaker synchronization mechanism. These findings hold across range of batch sizes and learning schedules, well as different kinds of models and datasets. Obtaining this level of accuracy. however, requires properly accounting for the regularization effect of codistilla- tion, 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,
INTRODUCTION
Several recent improvements in the performance of machine learning models can be attributed to scaling the training of neural network models (He et al.] 20166) Goyal et al.) 2017 Vaswani al.] 2017 Devlin et al., 2018 Shoeybi et al.) 2019, Huang et al., 2019 Kaplan et al. 2020 Lepikhin et al. 2020 Brown et al. 2020). Most approaches to scaling up training leverage some form of data parallelism (using multiple workers to compute gradients on different training samples in parallel), and the most common approach to data-parallel training is synchronous first-order optimization.
In synchronous data-parallel training, several replicas of neural network model are created, each on a different worker. The workers process different mini-batches locally at each step using an opti- mizer such Stochastic Gradient Descent (SGD) Adam (Kingma & Ba 2015), and the replicas synchronize (i.c., average either gradicnts or parameters) at every step by communicating either with centralized parameter server (Li et al.) 2014) using a reduce (Goyal al. (017). More computing resources can be used in parallel by increasing the number workers, effectively in- creasing the batch size used compute stochastic gradient. Increasing the batch size reduces the gradient's variance and ideally makes it possible to increase the learning rate in proportion to the number of workers. By doing so, the number of steps required to reach given model quality is also reduced in proportion to the number of workers, and near-linear speedup is achieved (Goyal al.) 2017). However, it has been observed that the linear learning rate scaling strategy leads to performance degradation for very large batch sizes (Goyal et al. (017), and even with more prin- cipled learning rate scaling mechanisms, synchronous SGD with larger batches eventually yields diminishing returns (Johnson et al. 2020).
Synchronous data-parallel methods ensure that all models are precisely synchronized at every step during training. This incurs substantial communication overhead, which increases with the number of replicas, and can quickly become bottleneck limiting the processing units' utilization (e.g.. GPU or TPU), especially when devices communicate A number of approaches have been proposed to reduce communication overhead, including using mixed-presd mixediprecision (Jia al. 2018) other forms of compression (Alistarh al. 2017 Bernstein
commodity interconnects such as Ethernet.

---

Under review as a conference paper at ICLR 2021
et al.
2018).
reducing the frequency of synchronization not occur after every optimizer step (Stich)
2018
Yu et al.) 2019), using cossip-based methods for approximate distributed averaging (Lian
al.
2017
Assran et al. 2019 Wang et al. 2020), or using some combination thereof (Wang &
Joshi
2018:
Koloskova et al. 2020).
Codistillation is an elegant alternative approach distributed training with reduced communication overhead (Anil et a 2018). Rather than synchronizing models to have the same weights, codistilla- tion seeks share information by having the models represent the same function (i.e., input-output mapping). Codistillation accomplishes this by incorporating a distillation-like loss that penalizes the predictions made by one model on batch of training samples for deviating from the predictions made by other models on the same batch. In practice, a worker updating one model can compute the predictions made by another model by reading checkpoints of the other model and performing an additional forward pass. Previous work has demonstrated that codistillation is quite tolerant to asynchronous execution using stale checkpoints, e.g., using another model's 's checkpoint from up to 50 updates ago without observing significant drop in accuracy (Anil et al. 2018).
Anil et al. (2018) focuses on the very large batch setting. For example, when training a ResNet-50 on ImageNet, codistilling two models and using batch size 16k per model achieves substantially better performance than training with synchronous SGD and batch size 32k, although the final accuracy still significantly lower than that achieved by synchronous SGD with smaller batch size (e.g., 8k or smaller). This performance boost attributed to an ensembling-like effect introduced by the codistillation loss.
In this paper we study codistillation at moderate batch sizes, i.e., before the performance syn- chronous SGD begins degrade. We demonstrate that it is possible use codistillation in this regime without losing accuracy. For example, when training a ResNet-50 on ImageNet, we show that codistilling two models, each model using batch size of 256, achieves comparable perfor- mance to training a single model using synchronous SGD with batch size 512. Furthermore, this holds across a range of batch sizes.
Achieving this performance parity involves modifying the way that explicit regularization is used in conjunction with codistillation. This modification stems from new insights into codistillation. Specifically, we demonstrate that codistillation has regularizing effect. Thus, while increasing the batch size in synchronous SGD helps training by reducing the gradient variance, we conjecture that codistillation helps in complementary manner via this regularization. Because it has regular- izing effect, care needs to taken when using codistillation in conjunction with other forms of regularization, such as L2 regularization (weight decay), to avoid over-regularizing. We also evalu- ate the sensitivity of codistillation to different checkpoints and learning rate schedule.
like the frequency of reading new
Overall, our findings complement previous work on codistillation (Anil et al. 2018 Zhang et al. 2018). We summarize below main contributions:
1. To the best of knowledge, demonstrate for the first time that models trained with codistillation can perform as well as models trained with traditional parallel SGD methods even when trained with the same number of workers and same number of updates, despite using much weaker synchronization mechanism (Section 3). Previous work at the inter- section of codistillation and distributed training used extremely large batch sizes and more workers than the parallel SGD counterparts.
2. Complementing the existing work codistillation, show that codistillation acts regularizer (Section 4p. Our work demonstrates that explicitly accounting for its regular- ization effect is key ingredient using codistillation without losing accuracy (compared to parallel SGD methods).
CODISTILLATION: BACKGROUND AND RELATED WORK
Codistillation is proposed as a mechanism for sharing information between multiple models being trained concurrently (Anil et al. 2018 Zhang al. 2018). In typical multi-phas distillation, first, a teacher model is trained using standard supervised learning, and then a student model is trained to predict the outputs of the teacher model without any updating of the teacher. In contrast, when two
2

---

Under review as a conference paper at ICLR 2021
Algorithm 1: Codistillation
Input: Loss function Ly, 9) and codistillation loss function Díy. 3/) Input: Model architecture fa( and initial model parameters {0] i= -......." nu
Input: Number of iterations K, learning rates {17 LK fk=1 and penalty coefficients ]kk
for :,...:K, do
for i ,....,,.. do
2, y = get www.eattmiibbthhhhcc
o.+ = 0% Vo, [VV.Jo (x)) + a a n-I Ezzt D (o
end
end
or more models codistill, there only one phase, and addition to minimizing the usual supervised loss on the training data, additional loss term is used to share information between models by encouraging each model to make similar predictions to the other(s).
Codistillation, as described in Zhang al. (2018), is shown in Algorithm [ Here, 2 2 models are trained concurrently. The ith model is updated on line 4 by taking a gradient step to minimize the combination of standard supervised loss funetion L (e.g., cross-entropy MSE) and a distillation- like loss D which penalizes differences between the predictions made by model and those made by model j, averaged
all other models j A i. /Zhang al. (2018) and AAiil et al. (2018) both report using Kullback-Lieble (KL) divergence for 1) in their experiments, although Anil et al. (2018) mentions that other options are possible (c.g.. mean squared error between the logits of different models). Zhang et al. (2018) does not explicitly include a penalty parameter a", and instead (implicitly) takes a* for all kt
Anil et al. (2018) focuses on codistillation as mechanism for scaling distributed training, and demonstrates that using codistillation allows for better scaling with larger batch sizes than conven- tional synchronous parallel SGD. When different models reside on different GPUs, implementing line requires that the devices communicate, Anil al. (20T8) proposes to have workers exchange parameter checkpoints, and these checkpoints updated periodically rather than after every itera- tion. This leads to some delay, i.e., actually using 05 with k' < k in line 4 However, it is argued that such delay tolerable for codistillation because large changes in parameters 0 do not necessar- ily correspond large changes to the function fe. Experimentally. it is reported that exchanging checkpoints every 50 updates does not noticeably degrade performance in their setup. Anil et al. (2018) also reports using burn-in period since, "in the beginning of training, the distillation term in the loss is not very useful or may even be counterproductive" During the burn-in period of the first K < K steps, the models train independently; i.e., a² = 0 for k 1...., K. Below we find that codistillation can tolerate large delays (e.g., updating checkpoints after hundreds or thousands of updates) without any burn-in period and without any significant loss in accuracy.
One common application of distillation to avoid the computational overhead of deploying ensemble. Motivated by this, previous works (Anil et al. 2018 Zhang et al., 2018) explore the connection between codistillation and ensemble methods, showing that two models trained using codistillation achieve performance close to that of ensembling two independentlly-rainee -rrannne models. Anil et al. (2018) also conducts an experiment that compares codistilling two models with simply using label smoothing regularizer when training one model, and reports that codistillation performs significantly better.
train a single model mimic an ensemble's predictions
CODISTILLATION, SYNCHRONOUS SGD, AND REGULARIZATION
The aim of this section is to demonstrate that codistillation can be used as a mechanism for dis- tributed training without losing accuracy, even when used in conjunction with batch sizes that are not particularly large. To achieve this, we also show that regularization needs to be handled care- fully. Typical training regimes the common practice is to use L2 regularization (weight decay) when training convolutional net- works (Goyal et al. 2017) and to label smoothing when training neural machine translation
some form of regularization to avoid overfitting, For example,
3

---

Under review as a conference paper at ICLR 2021
2012
1000
Vallaaton Accurott
Tranking to2s
1000
1ree
el,Ndec 5020
codas (3250)
53
300
200
168
750
Number of updotes (x10*)
(x10%)
uodates (X20%)
(a) Training Loss
(b) Top-1 Validation Accuracy
(c) Change in parameter values from initialization
Figure Comparing all.reduee and codistillation for training ResNet50 ImageNet. We codistill two models using batch size 256 for each, and the model trained with all-reduce uses batch size 512. All experiments are run on 16 GPUs. We report the (a) training loss and (b) top-1 validation accuracy. We observe that compared to all reduce, the model trained with codistilla- tion underfits, obtaining higher training loss and lower top-1 accuracy. (c) Examining the difference in parameters from initialization over the course of training further suggests that codistillation has regularizing effect that impacts performance, since parameters remain closer to their initial values.
models (Ou et al. 2018). Directly incorporating codistillation top of existing training pipelines can lead to underfitting, We begin this section by illustrating this point.
3.1 CODISTILLATION "OUT THE BOX" CAN WEER-REGULARIZE
We train ResNet50 (He al. 2016) model perform image classification the ImageNet dataset (Russakovsky et al. 22155) To begin, we compare training a single synchronous SGD model using batch size 512 with codistilling two models using batch size 256 per model. In both cases we perform training using 16 GPUs and batch size 32 per GPU. Similar to Anil et al. (2018), we follow the training procedure proposed Goyal al. (017 for the learning rate warmup, learning rate schedule, L2 regularization, and other hyperparameters. Specifically, synchronous SGD training (all_reduce) is run for 90 epochs. The learning rate starts at 0.1, regardless of batch size, and warmed over the first five epochs to 0.1 x b/256, where b is the effective batch size across all workers (e.g., 512 when using synchronous SGD and 16 GPUs). The learning rate then decreased by factor of 0.1 the beginning of epochs 30, 60, and 80.
When training with codistillation, workers read the checkpoints of the other model once every 2500 updates. At this frequency of checkpoint reading, we did not observe any benefit to using burn-in period. We experimented with reading checkpoints more frequently and also did not observe any improvement in performance. When training with two models and batch size 256, use the same total number of GPUs (16) as an a all.redrddd all_reduue model with batch size 512, as well as the same number of gradient updates per GPU (with each update processing mini-batch of size 32). Hence, when training using codistillation we only train the two models for 45 epochs (so that the number of updates per codistilled model is the same the number of updates performed by all.reduce using the same total number of GPUs) and scale all milestones in the learning rate schedule by half, i.e., warming up over the first 2.5 epochs, and decreasing epochs 15, 30, and 40.
In Fig. ((a) and [(b) we observe that the model trained with codistillation achieves higher training loss and lower top-I accuracy (around 2% lower) compared to the model trained with synchronous SGD. The codistillation performance is calculated using one of the two models (model 1) being codistilled. This suggests that although the codistillation loss may provide useful training signal, without other modifications it is not sufficient to achieve the same accuracy synchronous SGD using the same number GPUs.
Since both training loss and validation accuracy are affected, we hypothesize that the problem stems from ovuu-rggulaiization. To further investigate this we plot the change in model weights, relative to their initialization, over the course of training in Fig. [(c] We observe that parameters do not move as far in parameter space when training with codistillation as they do when using synchronous SGD, providing further evidence that codistillation acts regularizer.
4

---

Under review as a conference paper at ICLR 2021
70-
Validetion Accuracy
Velidation
78.6
30-
200325
20-
redure (512)
reduce (512)
Eation (2425))
latton (2256)
10
5
50
100
150
200
50
100
150
200
Number of updates (x10*)
Number of updates (x10*)
(a) ResNet50
(b) ResNeXt101
Figure 2: Comparing al reduce and codistillation for ResNet50 and ResNeXt-101 models on the ImageNet dataset, with decreasing weight decay and a shifted learning rate decay schedule compared to Goyal et (2017). Both methods achieve similar values of top-1 validation accuracy.
3.2 ACCOUNTING FOR CODISTILLATION REGULARIZATION EFFECT BRIDGES THE GAP
Motivated by the observations above, we conduct an experiment using less explicit regularization when training with codistillation. Goyal al. (2017) recommends using constant L2 regularization of 10-4 throughout training. We propose to start with the same initial value (10- reduce it to 10 after the first learning rate decay, and reduce it further to 0 after the second learning rate decay. When training with codistillation, we observe that the model's training loss saturates slower and we shift the learning rate schedule by a few epochs to account for this (from 15, 30, 40 to 18, 38, 44).
In Fig. [(a) we compare the top-1 validation accuracy on ImageNet for the ResNer50 model trained with codistillation and al all_reduce, with the above-mentioned modifications to the codistilla- tion setup. Accounting for codistillation's regularization effeec helps narrow the performance gap, with all reduce now performing only marginally better (76.1% 1% for all reduce vs 75.9% for codistillation). We verify that with constant L.2 regularization of 10-4 throughout training, codis- tillation performs much worse than all_reduca implying that the improved performance is not only from the modified learning rate schedule; see Fig. 12 in the Appendix. To validate that the improvement is not specific to the ResNet50 architecture, we repeat the experiment using the larger ResNeXt101, ResNet152 (Fig. in Appendix) and ResNext152 (Fig. i0 Appendix) architectures. In Fig. 2(b) we observe that codistillation performs marginally better than all all.reduce, We be- lieve that the larger ResNeXt101 may benefit more from codistillation due having more capacity than the ResNet50. The training loss curves for both models are shown in Fig. g in the Appendix.
In synchronous SGD (all reduce), more workers can be added increase the effective batch size (summed across all workers). This reduces the gradient's variance and the model can be trained with a larger learning rate and fewer steps, while maintaining a similar level of accuracy (Goyal et al. (017). In Fig. 3add Table A.2 we demonstrate similar effect with codistillation. As one doubles the batch size per worker (and hence doubles the effective batch size), the learning rate can also be doubled and the model reaches similar performance in half the number of steps.
The experiments above use the step-wise learning rate schedule described in Goyal et al. (2017). We want to ascertain that our findings not dependent on this specific learning rate schedule. Hence we train the ResNet50 and ResNeXt101 models with the cosine learning rate schedule (He et al.) 2019). In Fig. E observe that the final validation performance for the two approaches is very close, confirming that codistillation works consistently across different learning rate schedules. The corresponding training loss plots shown in Fig. T in the Appendix.
So far, all our experiments have focused on the ImageNet dataset and the ResNet family of mod- els. Next, we evaluate the codistillation mechanism for neural machine translation (NMT). Specif- ically, we train the "big" transformer model (Vaswani et al.) 2017) (6 blocks in the encoder and decoder networks) on the WMT' 16 En-De translation dataset, following the setup described in Oul et al.) (2018). We note that for codistillation, reducing the explicit regularization is again important achieving performance comparable to a all.reduce (see details Appendix, Section A.
5

---

Under review as a conference paper at ICLR 2021
codattiacc (2x125)
codiatizzcoo (2x23))
codiailizccon (2x512)
Validation Accuracy
Training
codiit Fution 12128
codiit Fltion (2x25))
codst llaton (2512)
codie laaron 1270
codie Batton (722000)
100
760 300
400
100
200
303
450
Number of updates (x10*)
Number updates (x10P)
(a) Training Loss
(b) Validation Accuracy
Figure 3: Codistillation scales well across multiple values of batch size per worker. Each time we double the batch size per worker, we scale the learning rate schedule by factor of two and perform half the number of updates. We do not observe any significant degradation in (a) training loss (b) validation accuracy across wide range of batch sizes. The training loss and validation accuracy for different batch sizes are also provided in Table A.2
dalidtionn Accuracy
dation
60.0
77
30-
20-
10-
50
Number of updates (x10*)
Number of updates (x10")
(a) ResNet50
ResNeXt101
Figure 4: Comparing a all_reduce and codistillation for ResNet50 and ResNeXt101 models on the ImageNet dataset, using a cosine learning rate schedule (He et al. 2019). We observe that the final validation performance for the two approaches is very close, confirming that codistillation works consistently across different learning rate schedules.
This observation line with our previous ob- servations on convolutional models. In Fig. 5 we observe that the model trained with codis- tillation reaches similar validation loss as the model trained with a all. _rrduce. This con- firms that codistillation with adjusted regular- ization schedule also extends to NMT. The cor- responding training plot is shown in Fig. 14 in the Appendix.
al reduce codistillaion
Validation Loss
2125
There are several important aspects of the above results that are worth emphasizing. (i) Unlike prior work on codistillation, we train both se- tups with the same number of workers for the same number of steps. We believe this con- straint to be particularly important to better un- derstand the impact of codistillation on many practical distributed training workflows. (ii) Our objective is not to show that codistillation outperform reduce, and initial
2500 5000 7500 10000 12500 15000 17500 Number of updates
Figure 5: Comparing all.reduce and codis- tillation using "big" transformer model on 16 En-De dataset. Models trained using all_reduce and codistillation both reach sim- ilar NLL on the validation dataset.

---

Under review as a conference paper at ICLR 2021
recuse rodadtiation (22266) -
roduce (312)
1.0
rodelliation (2:550
Validation Accuracy
70-
65
Training
0.4
60-
0.2
Fraction of dataset used for training
Fraction of dataset used for training
(a) Final Validation Top-1 Accuracy
(b) Final Training Loss
Figure 6: Final validation top-1 accuracy and training loss of ResNet-50 model when trained using fraction of the full training data. We observe smaller fractions of training data are used (and model starts overfittin), codistillation setup increasingly improves over the a all Treduce setup in terms of validation accuracy.
expectation was actually that replacing all.re 1._reduce by codistillation would degrade the perfor- mance when using the same number of updates (as suggested in Fig. [ since codistillation is much weaker form of synchronization. With these experiments, we want to quantify this perfor- mance loss and understand how we can minimize it. Interestingly, we find that the performance gap can be narrowed down significantly, with codistillation even outperforming all.reduce in some settings (Fig. 4bb.. (iii) We believe that previous works did not investigate the regularization effect of codistillation because they focused on the ensembling effect obtained by increasing the number of workers. Being aware of the regularization effect and explicitly accounting for it is important when using codistillation in practice, as seen in Fig.
CODISTILLAAION HELPS REDUCE OVERFITTING
Section iillustrated that applying codistillation addition to other forms of regularization can result in over-regularizing, and progressively reducing the explicit regularization helps improve training and generalization. This section further explores the regularizing effect of codistillation by examin- ing settings where expect the model to overfit.
We simulate a scenario where the ResNet50 model is likely to overfit a controlled way by training the model using only fraction of the full ImageNet training set. When training on (1/k)th of the training set, we multiply the number of epochs by k so that the total number of model updates performed is the same when training with the full training set. The learning rate and weight decay schedules are also modified accordingly. As we train on less data, we expect the model to overfit, i.e. to obtain lower training loss and lower validation accuracy.
In Figure 6 we indeed observe that overfitting occurs when using less training data. However, the overfitting is less severe when training using codistillation, providing further support for the hypothesis that codistillation acts as regularizer.
An interesting side-effect of this observation that codistillation could be an interesting alternative to all reduce for training over-parameterized models. We further verify this hypothesis by train- ing a "big" Transformer model for small NMT dataset (IWSLT 2014 German-Fnglist translation dataset that contains 153K training sentences, 7K development sentences, and 7K test sentences). The model trained with codistillation achieves a validation NLL of 2.31 whereas the model trained with all alllree reduce reaches validation NLL of 2.37.
ROBUSTNESS TO UPDATE FREQUENCY AND LEARNING RATE
To further understand the codistillation setup, we conduct additional ablations examining the effect of some syper--arrmmetr choices.

---

Under review as a conference paper at ICLR 2021
P=01
70
70-
Velidetion Accurasy
Validetion Accuracy
63
90-
90-
30-
30-
50
100
150
200
50
100 150
200
Number of updates (x10*)
Number of updates (x10%)
(a) Vary Frequency Exchanging Check- points
Vary Learning Rate
Figure 7: Evaluating the robustness of codistillation setup by varying the frequency of exchanging checkpoints (left) and learning rate (right) when training ResNet50 ImageNet dataset. While there are some gains by using a higher frequency, the results do not degrade too much with lower frequency. The learning
be increased from 0.1 to 0.2 without hurting the performance.
In distributed training, model weights exchanged regularly between the training workers. In synchronous SGD, model parameters are communicated every gradient update, via a _reducee, which introduces the large communication overhead. In codistillation, the frequency of reading model checkpoints should be treated as a hnyer-paramette and should be tuned for the specific use cases. ZZhang et al., (2018) always uses the most recent version of the model for computing the codistillation loss, equivalent to reading new checkpoint at every iteration. Anil et al. (2018) reports exchanging checkpoints every updates.
In the ImageNet results presented in Sections 3 and exchange the model checkpoints once every 2500 updates. In Fig. [(a) we evaluate how the validation performance for the ResNet models varies with the frequency of exchanging the checkpoints). The curve corresponding to 1 is the baseline setup, with checkpoints exchanged every 2500 updates. Label 2 denotes the case where checkpoints are exchanged twice frequently (every 1250 updates), and other curves similarly show 0.5x and 4 x. The corresponding training loss curves are shown in Fig. 15(a) in the Appendix. We observe that the setup quite robust the frequency of exchange. Although there some small improvements by using a higher frequency, the results do not degrade too much with lower frequency. Similar trends hold for the Transformer model (for language translation task), as shown in Fig. 16(a) in the Appendix.
In Fig. 7(b) we vary the base learning rate when training the codistillation model on the ImageNet dataset, using the setup proposed in Goyal et al. (2017) and report the accuracy on the validation dataset. We observe that increasing the learning rate from 0.1 to 0.2 has almost no impact on the validation accuracy. Beyond that, higher learning rates hurt the validation accuracy more. The corresponding plot with the training loss shown Fig. [T(bb] in the Appendix.
CONCLUSION
In this work, demonstrate for the first time (to the best of our knowledge), that even for moderate batch sizes, models trained with codistillation can perform as well as models trained with traditional parallel SGD methods (while using comparable computational resources). Complementing existing works, we show that codistillation acts as regularizer, and that accounting for this property is essential when using codistillation. There are several exciting and potentially impactful directions for extending our understanding of codistillation. One important direction is to codistill more than two models. The loosely synchronized nature of codistillation potentially allows using different topologies for determining which models codistill together. Zhang et al. (2018) does consider one version of the setup, where all models codistill with each other, but they do so in a non-distributed setting. Another possible future work will be consider codistillation between models that different, e.g., in terms of architecture, capacity, or performance the standard knowledge distillation based approaches.
and evaluate if that improves
8

---

Under review as a conference paper at ICLR 2021
REFERENCES
Dan Alistarh, Demjan Grubic, Jerry Li, Ryota Tomioka, and Milan Vojnovic. Qsgd: Communication-effcientt sgd via gradient quantization and encoding. Advances in Neural Information Processing Systems, pp. 1709-1720 2017.
Rohan Anil, Gabriel Pereyra, Alexandre Passos, Robert Ormandi, George E. Dahl, and Geoffrey E. Hinton. Large scale distributed neural network training through online distillation. Interna- tional Conference on Learning Representations (ICLR), 2018.
Mahmoud Assran, Nicolas Loizou, Nicolas Ballas, and Mike Rabbat. Stochastic gradient push for distributed deep learning. volume 97 of Proceedings of Machine Learning Research, pp. 344-353, Long Beach, California, USA, 09-15 Jun 2019. PMLR.
Jeremy Bernstein, Yu-X Yu-Xiang Wang, Kamyar Avizzadenesheli, and Anima Anandkumar. signsgd: Compressed optimisation for non-convex problems. arXiv preprint arXiv: 04002, 2018.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 14165, 2020.
Jacob Devlin, Ming- Weig-Wei ee Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.0488 04800, 2018.
Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, An- drew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in hour. arXiv preprint arXiv: 1706.02677 2017.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, and Mu Li. Bag of tricks for image classification with convolutional neural networks. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 558-567, 2019.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, HyoukJoong Lee, Jiquan Ngiam, Quoc V Le, Yonghui Wu, al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Advances in neural information processing systems, pp- 103-112, 2019.
Xianyan Jia, Shutao Song, Wei He, Yangzihao Wang, Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang, Liwei Yu, Tiegang Chen, Guangxiao Hu, Shaohuai Shi, and Xi- aowen Chu. Highly scalable deep learning training system with mixed-precision: Training ima- genet in four minutes. arXiv preprint arXiv: 11255, 2018.
Tyler B. Johnson, Pulkit Agrawal, Haijie Gu, and Carlos Guestrin. AdaScale SGD: A user-friendly algorithm for distributed training. In International Conference on Machine Learning, 2020.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Anastasia Koloskova, Tao Lin, Sebastian U Stich, and Martin Jaggi. Decentralized deep learning with arbitrary communication compression. In International Conference on Learning Represen- tations, 2020.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu. Dehao Chen. Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.

---

Under review as a conference paper at ICLR 2021
Mu Li, David G. Andersen, Alexander Smola, and Kai Yu. Communication efficient distributed machine learning with the parameter server. In Advances in Neural Information Processing Sys- tems, 2014.
Xiangru Lian, Ce Zhang, Huan Zhang, Cho-Jui Hsieh, Wei Zhang, and Ji Liu. Can decentralized algorithms outperform centralized algorithms? a case study for decentralized parallel stochastic gradient descent. In Advances Neural Information Processing Systems, pp. 5330-5340, 2017.
Myle Ott, Sergey Edunov, David Grangier, and Michael Auli. Scaling neural machine translation. arXiv preprint arXiv: 1806.00188 2018.
Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng. David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of NAACL-HLT 2019: Demonstrations, 2019.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Viti,o, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang. Andrej Karpathy, Aditya Khosla, Michael Bernstein, al. Imagenet large scale visual recognition challenge. International Journal of Computer Vision, 115(3):211-252, 2015.
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LMM: Training multi-billion parameter language models using GPU model parallelism. arXiv preprint arXiv: 1909: 00055,, 2019.
Sebastian U Stich. Local SGD converges fast and communicates little.
arXiv preprint
vrXiv:1805.000 2018
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008 6008, 2017.
Jianyu Wang and Gauri Joshi. Cooperative SGD: A unified framework for the design and analysis of communication-efficenn SGD algorithms. arXiv preprint arXiv:1808. 070000 2018
Jianyu Wang, Vinayak Tantia, Nicolas Ballas, and Michael Rabbat. SlowMo: Improving fommunication-efffcient distributed SGD with slow momentum. In International Conference on Learning Representations, 2020.
Hao Yu, Sen Yang. and Shenghuo Zhu. Parallel restarted sgd with faster convergence and less
communication: Demystifying why model averaging works for deep learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pp. 5693-5700, 2019.
Ying Zhang, Tao Xiang, Timothy M. Hospedales, and Huchuan Lu. Deep mutual learning. In IEEE Conf. on Computer Vision and Pattern Recognition (CVPR), pp. 4320-4328, 2018.
10

---

Under review as a conference paper at ICLR 2021
A APPENDIX
A.1 IMPLEMENTATION DETAILS
The code for all experiments is implemented using PyTorch 1.5 (Paszke et al. 2017). All experi- ments are run for 3 seeds and plots report the average all seeds. The codistillation loss is the mean squared error between the logits of the two models (without centering them, as we found in preliminary experiments that centering yielded similar results).
The experiments involving the ImageNet dataset are based on the setup proposed in Goyal et al. 2017). We consider Two schedules control the value of the codistillation loss coefficient a the constant schedule, a* = for all k; and the step-wise schedule where we increase a every time we decrease the learning rate. We did not find any significant difference in the performance of the two cases and report the performance for the case where a increases by 0.5 every time we decrease the learning rate.
The experiments for translation tasks use the FairSeq library (Ou et al. 2019). These experiments are based on the setup proposed Ott al. (2018). We consider two schedules to control the value of a² the constant schedule and the exponential schedule where we increase a by a factor y = 1.1 every epoch. We find in Fig. 17 that the exponential schedule works better than the constant alpha schedule (2. 12 vs 1. validation perplexity), and this is the schedule we use in all other experiments on this task. When training the translation models with codistillation, we explicitly reduce the amount of regularization (to account for regularization added by codistillation) by removing the label smoothing loss.
A.2 ADDITIONAL RESULTS ON IMAGENET
all reduce (512) codistilation (2x256)
aeLuue (512)
codistilation (2x256)
5-
Training Loss
Training
1.2
0.6
1.0
200
225
200.255
50
100
150
200
50
100
150
200
Number of updates (x10²)
Number updates (x10²)
(a) ResNe150
(b) ResNeXt101
Figure 8: Comparing the training loss for the all.reduce and codistillation setups for ResNet50 and ResNeXt- 101 models (respectively) on the ImageNet dataset, with decreasing weight decay and a shifted learning rate decay scheduled compared to Goyal et al. (2017). While the redure setup reaches a lower training loss, the performance the validation dataset (in terms of top-1 validation accuracy) is very similar for the two setups (for both models) as shown in Fig. 2
In Fig. 8. we plot the training loss for the ResNet50 (left) and ResNeXt101 (right) models for the ImageNet dataset with some variations to the setup proposed Goyal et al. (011)) ). as Section 3.2 Specifically, Gooaa et al. (2017) recommends using constant 7.2 weight decay set to 10-4 through- out training. Keeping the initial value of this weight decay to 10 we reduce it to 10-5 after the first learning rate decay and further to 0 after the second learning rate decay. We also modify the learning rate schedule from Goyal et al. (211)). which based on how the training loss changes (and saturates) during training. Due the regularization effect of codistillation, the model's 's training loss saturates slower and we shift the schedule by a few epochs to account for this (from 15, 30, 40 to 18, 38, 44). We note that for both models, the a 11 reduce setup reaches
observe that
11

---

Under review as a conference paper at ICLR 2021
lower training loss. Despite this, the performance the validation dataset (in terms of top-1 validation accuracy) is very similar for the two setups (for both models) as shown in Fig.
all reduce (512)
5-
codistilation
70-
Validation Accuracy
Training Loss
60
0.75
175200225
175200725
codistillatinn 12 256)
50 100 150
200
50
100
150
200
Number updates (x10²)
Number updates (x10²)
(a) Training Loss
(b) Validation Accuracy
Figure 9: Comparing the training loss and validation accuracy (respectively) for the all.reduce and codistillation setups on the ImageNet dataset, with decreasing weight decay and shifted learn- ing rate decay scheduled compared to Goyal et al. (2017). While the all all-reduce setup reaches lower training loss, the performance on the validation dataset (in terms of top-1 validation accuracy) is very similar for the two setups.
6-
reduce (512)
codistilation (2 2561
5-
Validation Accuracy
Training Loss
60
50
80.0
77.5
175200
all reduce (512)
codistilaaion 256)
200
Number of updates (x10")
Number of updates (x10')
(a) Training Loss
(b) Validation Accuracy
Figure 10: Comparing the training loss and validation accuracy (respectively) for the alleallrddee reduce and codistillation setups on the ImageNet dataset, with decreasing weight decay and shifted learn- ing rate decay scheduled compared to Goyal al. (2017). While the a all.redee l_rrduuce setup reaches lower training loss, the performance on the validation dataset (in terms of top-1 validation accuracy) very similar for the two setups.
we show the training performance for ResNet50 and ResNeXt101 models when using the cosine learning rate schedule (He al. 2019). Again, while the a 1ll_reduee setup reaches lower training loss, the performance on the validation dataset (in terms of top-1 validation accuracy) very similar for the two setups (for both models) as shown in Fig.
Fig. II
In Fig. 12
we verify that the good performance observed in Fig. is not due only to the shifted learning rate schedule. We observe that, indeed, with this shifted schedule, if we keep the L2 weight decay to a constant value of 10-4 throughout training, then codistillation is significantly worse than both in terms of training loss and validation accuracy.
In synchronous SGD (all reduce), more workers can be added to increase the effective batch size (summed across all workers). The increased batch size reduces the gradient's variance and the
12

---

Under review as a conference paper at ICLR 2021
6-
reduce (512)
all reduce (512)
(2x256)
Training Loss
4-
1.0
2-
250
200
50
100
150
200
250
300
50
100
150
200
250
300
Number of updates (x10²)
Number of updates (x10²)
(a) ResNet50
(b) ResNeXt101
Figure 11: Comparing the training loss for the all reduce and codistillation setups for ResNet50 and ResNeXt-100 101 models (respectively) on the ImageNet dataset, using the cosine learning rate schedule proposed in He et al. (2019). While the all-reduce setup reaches lower training loss, the performance on the validation dataset (in terms top-1 validation accuracy) very similar for the two setups (for both models) shown in Fig.
reuuu (512)
codistilation (2x256)
Validation Accuracy
60
Training Loss
200
50
30
reduce
codistilation (2x256)
50
200
Number of updates (x10"))
Number of updates (x10")
(a) Training Loss
(b) Validation Accuracy
Figure 12: Comparing the training loss (left) and validation accuracy (right) of the all. reduuce and codistillation setups with the ResNet50 model on the ImageNet dataset. The training setup is the same as Fig. 2(a) with one difference the L2 weight decay is kept at a constant value 10-4 throughout training. These plots show that when we do not account for the regularization effect of codistillation, the codistillation model's performance lags behind a l_rueduce.
model be trained with a larger learning rate and fewer steps, while maintaining similar level of accuracy (Goyal et al. 2017). In Fig. and Table A.2 we demonstrate similar effect with codistillation, As one doubles the batch size per worker (and hence doubles the effective batch size), the learning rate can also be doubled and the model reaches similar performance in half the number of steps.
Batch Size
Training Loss
Validation Accuracy
2 x 128
1.12
75.61
2 x 256
1.08
75.92
2 x 512
1.12
75.51
1024
1.13
75.45
2 x 2048
1.14
75.26
13

---

200 Under review as a conference paper at ICLR 2021
0.5
8.05
Validetion Accuracy
Training Loxs
1.20
1.15
RES
32-
210 220
20
50
100
150
700
50
100
150
Number of updates (x10*)
Number of updates (x10*)
(a) Training Loss
(b) Validation Accuracy
Figure 13: Evaluating codistillation setup by varying the value of ag. Note that the value is not changed during training and hence is denoted by
A.3 EFFECT OF VARYING a
In Fig 13 we evaluate codistillation setup by varying the value of ag. Note that the value is changed during training and hence is denoted by
A.4 ADDITIONAL RESULTS WMT 16 EN-DE
Codist llation
Training
2500 5000
7500 10000 12500 15000 17500 Number of updates
Figure 14: Comparing a 1_r_rducce and codistillation setups using the "big" transformer model on the WMT 16 En-De dataset. The model trained using codistillation performs worse in terms of training performance, but generalizes well to the validation dataset, as seen in Fig.
In Fig. 14 we plot the training loss of the a all reduce and codistillation setups on the WMT" 16 En-De dataset. We observe that codistillation has higher training loss, but almost matches the validation loss of all_reduce as shown in Fig. 5
A.5 ADDITIONAL RESULTS CODISTILLATION ROBUSTNESS
To further understand the codistillation setup, we conduct additional ablations examining the effect of some hyper-parameter choices.
In Fig. 15(a) we compare the training performance of ResNet50 models trained on the ImageNet dataset, as we vary the frequency of exchanging the checkpoints. The curve corresponding to 1x the baseline setup, with checkpoints exchanged every 2500 updates. Label 2x denotes the case where checkpoints are exchanged twice as frequently (every 1250 updates), and other curves sim- ilarly show 0.5x and 4x. The training loss is robust changes the frequency, reaching very similar values across the whole range. The corresponding top-1 validation accuracy is shown in Fig. (a)
In Fig. [5(b) we vary the base learning rate when training the codistillation model on the ImageNet dataset, using the setup proposed in Goyal et al. (2017) and report the loss the training dataset.
14

---

Under review as a conference paper at ICLR 2021
0.5x
Training Loss
Training Loss
220
50
100
150
700
50
100
150
200
Number of updates (x10*)
Number of updates (x10P)
(a) Vary Frequency Exchanging Check- points
Vary Learning Rate
Figure 15: Evaluating the training loss robustness the codistillation setup by varying the frequency of exchanging checkpoints (left) and learning rate (right) when training ResNet50 model the ImageNet dataset. The training loss robust to changes to the frequency, reaching very similar values across the whole range. The learning rate be increased from 0.1 to 0.2 without hurting the training performance.
We observe that increasing the learning rate from 0.1 to 0.2 has almost impact the training loss. Beyond that, higher learning rates affect the training loss more. The corresponding plot with the top-I validation accuracy is shown in Fig. 7(b) and exhibits similar trend.
In Fig. [6 we vary the frequency of exchanging the model checkpoints for the Transformer model. We note that more frequent updates lead a higher validation loss initially, but training continues, this setup eventually outperforms (slightly) the one with less frequent updates. Overall, the model's performance does not decline drastically
when exchange checkpoints 2.5 times less often.
In Fig. 17 both the initial value a (codistillation loss coefficient) and (scaling factor for o) varied when training the Transformer model. Since we are using a geometric progression for increasing the value of a, changing 7 alone can increase a to very large values, thus making training unstable. Henee, for each y E [1, 1.1, 1.5) we vary the initial a to achieve optimal results. The case of y = 1 corresponds to keeping a constant. This case slightly under-perte mnder-performs as compared to the cases where we vary a.
24
Trequency frequency
frequency frequency
32
Validation Loos
10-
Training
2,122
200
2:22
2500 0
5000
7000 10300 12500 15000 17500 Number of updates
2300 10000
5000
7300
12500 15000 17500
Number of updates
Training Loss
Validation Loss
Figure 16: Evaluating the robustness of codistillation for the Transformer model by varying the frequency exchanging checkpoints. The performance does not degrade when exchanging check- points less frequently.
15

---

Under review as a conference paper at ICLR 2021
1.1
Validation Less
--
20-
Training
232
5000 7500
5000
7500
10000 15000
Number of updates
Number of updates
(a) Training Loss
(b) Validation Loss
Figure 17: Evaluating the robustness of codistillation for the Transformer model by varying the initial value of a and y. Given the nature geometric progression, we perform a sweep over a to achieve optimal performance for each y E {1, 1.1, The of 1 corresponds to keeping a constant. This case slightly under- ander-perforss as compared to the cases where we increase a over time.
16


Reviews:
[{"review_id": "lsfBSka53Gw", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Update: I have read the authors' response, and have decided to keep my score as-is\n\nClaims:\n\n1) Codistillation achieves accuracy close to synchronous gradient descent if the model's built-in regularizers (e.g. L2) are decreased to zero while training\n\n2) Codistillation therefore acts as a regularizer, and should be explicitly accounted for during training\n\n3) Codistillation works for moderate batch sizes, not just large-batch\n\nPros:\n\n-Interesting observations about codistillation's empirical behavior, which contributes to the overall literature on distributed training\n\n-Codistillation imposes much lower network costs since model parameters only have to be exchanged once every few thousand updates (rather than every mini-batch of updates)\n\nCons:\n\n-There are problems with the systems setting, which undermine the strength and validity of the claims:\n\n1) The setting is not fully clear: given that only n=2 codistillation replicas are used throughout the paper, how are the GPUs arranged? My guess is that there are 2 physical machines with 8 GPUs each, and within each machine, the GPUs exchange their updates synchronously. This information should be clearly presented.\n\n2) The overheads of synchronous distributed training increase with higher machine/replica count, so codistillation should be really be studied with many replicas, which is where it has the greatest potential to have an advantage. Note that, depending on the model, network and GPU hardware, synchronous distributed training on just 2 replicas/machines may not even incur any network overhead, which completely eliminates the overhead-reducing advantages of codistillation. Hence, the paper needs to study codistillation with at least n=4 to 16 replicas, verifying that the regularization effect of codistillation still holds with more replicas, and that the proposed means of controlling regularization effects (namely, reducing the built-in regularizers) is still effective. Thought experiment: what happens if the regularization effect of having 16 replicas is so strong (because there are 16 terms in the auxiliary loss) that reducing the built-in regularizers is no longer sufficient?\n\n3) The increased compute overhead of codistillation (due to auxiliary loss needing to evaluate every of the n model replicas every update) is not investigated, so it is unclear how it stacks up against the decrease in network costs. Runtime benchmarks versus all-reduce could have been provided and analyzed.\n\n-Experiments are performed on only two models, ResNet and Transformer. Comparable papers that focus on empirical evidence usually test their hypotheses on five or more models, in order to demonstrate that the technique is truly generic. Testing on non-residual CNN architectures (e.g. VGG) and RNN architectures would have made the results more convincing.\n\n-No systematic method is proposed/evaluated for reducing the model's built-in regularizers. Even if no theory is presented, I imagine the paper could have presented at least a rate function with controllable hyperparameter (similar to learning rate), and systematically conducted experiments for different values of the hyperparameter. Currently, each model's builtin regularizers are reduced in an ad hoc fashion, so the evidence is only of anecdotal quality.\n\nSuggestions for improvement:\n\n-In Section 4, the authors may want to clarify whether the model's built-in regularizers (L2 etc.) are or are not being reduced - it was not clearly stated, and I am guessing it is the latter case.\n\n-The paper would not be complete without a discussion of Federated Learning (McMahan et al., 2017), which provides another means of reducing the network overhead of synchronous distributed training, among other benefits. Even if experiments against Federated Learning are not conducted, it should be discussed in the related work.\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "VpVOBN-uOyA", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary of paper:\nThe paper studies the concept of codistillation in data parallel distributed training. In this setting, the standard minibatch SGD algorithm requires exchange of models in every update of every node. Recent work in distributed training has studied \"local SGD\", where models are exchanged at frequent (usually periodic) intervals after a bunch of local updates. This paper studies an alternative, called \"codistillation\". The idea is that at a given node, say node $i$, the local model updates are regularized by the most recent models at nodes $\\{j, j \\neq i\\}$ through an appropriately modified loss term. Specifically, the loss term bias the model at node $i$ towards having similar classification outcomes on the training data as the (most recent) local estimate of the model at nodes $\\{j, j \\neq i\\}.\n\nThe paper argues that this codistillation approach has  a regularizing effect, and if the other regularization approaches are combined with codistillation in a balanced manner so as to avoid over-regularization, then the performance can be similar to minibatch SGD, with much lower global communication.\n\nPros:\n--> The paper's results are certainly intriguing and likely to lead to further investigations of codistillation as an alternative or complementary approach to local SGD for reducing communications in distributed training. \n --> The results are promising and the experimental results seem quite comprehensive. \n\nCons:\n--> I am not sure completely understand/buy the reasoning provided in the paper. The regularization in imagenet - as I understand it - biases towards chosing models with lower weights (e.g., sparser models). The regularization effect of co-distillation is towards chosing models that are not too far away from the initial point - indeed 1(c) shows that the codistilled models are not too far away from the original models.  However, the paper (e.g, Sections 3.1, 3.2) seems to suggest these two forms of regularization are similar/equivalent, and one form can be increased and the other form can be lowered.  I find this explanation not very convincing. \n--> I think the regularization parameters (e.g., what exactly is the constant L2 regularization, what is the factor of 10^{-4} in Section 3.2) needs to be explained in more detail in the paper for completeness...at this point, the reading requires knowledge of the details of Goyal (2017) for the reader to know what exactly is implemented by the authors.\n--> The result that codistillation with local updates can lead to lower loss as compared to minibatch SGD is almost too surprising to be true. Common intuition suggests that the best achievable performance in terms of accuracy per update that any distributed approach can get is that of minibatch SGD, which is what one would do if all the data was centralized at one node. So I am suspcious of the improvements. I think there are two possible reasons, that the authors might clarify (or suggest other reasons)\na) The benchmark minibatch SGD model is not regularized properly, or something else is ineffectively done in this benchmark estimate (e.g., perhaps hyper parameters are not as well chosen).\nb) Perhaps I do not understand what exactly is done in the minibatch SGD setting, e.g., is the minibatch SGD done with one round of global communication in every model update, or is there some adjustment made to make the communication costs comparable?\n\nSummary:\nAlthough not all aspects of the paper are convincing,  the results of the paper are intriguing and it is likely to lead to further investigations along this topic.  I am on the borderline on this paper, perhaps leaning slightly more towards acceptance.\n \n\n====== Comments after author response ====\n\nThanks for the detailed response. Having read the update and the other reviews, I have lowered my score. \n I am still left unconvinced on (at least) two aspects.\n\n--> It is unclear to me as to how a method can improve upon minibatch-SGD can actually have better generalization. While nothing theoretically rules this out, it possible means that that either (A) the model architecture was sub-optimal and there is room for improvement, or (B) the optimization found a minimum that is not very good. In either case, there appears to be a different \"lesson\"  than the one described in the paper. This is certainly worth exploring in more detail.\n\n--> I agree with some of the other reviewers that the experimental results can be improved, specifically, the usage of just 2 replicas (now clarified in the paper) is limiting, and a more detailed analysis regarding how to tune the regularization parameters.\n\nI do think that the paper is on track towards an interesting discovery, but I would like to see a deeper/more detailed analysis to be convinced.  \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "1z_AssDWqsQ", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper aims to have a closer look at the role of codistillation for distributed training. Authors provided an answer with their empirical observations. That is, codistillation acts as a regularizer, since the distance between the learned model and the initialization is smaller than sync SGD without codistillation. Then, the authors claim that the codistillation may over-regularize and study how to modify the training configurations to avoid it. There are further discussions on the overfitting and robustness to hyper-parameters in sec 4 and sec 5. \n\n\nComparing to sync SGD, codistillation bring in auxillary loss enforcing all the local workers learn from each other in some degree. Thus, it is intuitively a regularization term. Deep understanding on how it takes effect will help us better trade off the bias and variance in distributed training. Unfortunately, this paper does not achieve this goal. \n\nMost of the claims in this paper are not surprising, with limited insights. 1) It is common sense that, regularization may “over-regularize” and can “reduce overfitting”. The empirical observations to demonstrate these are not necessary in sec 3.1 and sec 4. 2) To “bridge the (over-regularize) gap” (In sec 3.2), authors report the empirical behavior of distributed training with codistillation, with difference choice of learning rate, batch size, concise learning rate schedule, tasks (NMT), and update frequency (in sec 5). The proposed modification of training is minor for a specific setting, i.e., to decreasing learning rate from 10^{-4} to 10^{-5} on ImageNet for ResNet50. \n\nAuthors ignore the factors which are more direct to trade off bias and variance. 1) The coefficient of the auxiliary loss \\alpha can adjust the regularization. I have not seen the discussion on it. What is the value for it in the experiments? Did I miss anything? 2) For me, the number of local workers may alleviate or enlarge the regularization of codistillation. It seems that, most of the experiments are conducted with 2 models/workers. \n\nOverall, this paper is easy to follow, but the results (from the experiments) are non-surprising, with limited insights. The modification proposed is minor and did’t get significant improvement of performance.\n\n\nOverall, this paper proposed methods to make CNN to be symmetric as (NS euqation) dynamics. The technical contribution need be highlighted.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "1xZNxZ0qetv", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work analyzes the effect of co-distillation for distributed training under moderate batch sizes. Using distillation-like techniques to improve synchronous SGD training is an interesting direction. And the paper carefully analyzed this setting while using the same amount of compute, which is not done by prior work to my knowledge. In addition, the writing is good and easy to follow.\n\nHowever, I think the work can be improved in the following aspects:\n\nW1: The setup of the work is not very realistic. \n\nThe authors proposed to analyze distributed training (using two machines) with very small batch sizes (e.g., 256). These two constraints don't appear to be commonly used for ML tasks where distributed training help. E.g., for CV, Exploring the Limits of Weakly Supervised Pretraining (Dhruv et al., 2018) was already using batch sizes of 8,064 back in 2018; modern NLP tasks use significantly larger batch sizes (up to 4M in GPipe, Huang et al., 2018). Even if hardware is a limiting factor, showing how results scale with more than 2 model copies would help a lot in demonstrating the potential value of this work in distributed settings.\n\nW2: The contribution of this work seems to be incremental.\n\nConsidering distillation / co-distillation's regularization effect is not exactly a novel direction. This line of thoughts originated at least from Distilling the Knowledge in a Neural Network (Hilton et al., 2014) and was discussed in other papers like Collaborative Learning for Deep Neural Networks (Song et al., 2018).\n\nW3: Experiments could be done on top of stronger baselines. \n\nThe authors proposed that co-distillation benefits models with higher capacity more (e.g., 3.2: ResNeXt10(1) may benefit more from codistillation due to having more capacity than ResNet50), but there are no CV experiments demonstrating that this scales with more parameters done on more complex architectures (e.g., SENet, AmoebaNet, etc).\n\nW4: In the NLP/NMT case, there are no experiments showing that this benefits self-attention architectures besides a remark in Section 4. For small datasets like IWSLT'14 DE->EN, it's known that using fewer parameters  (e.g., 256d) vs the default 1024 in transformer-big produces much better results. See e.g., Tied Transformers: Neural Machine Translation with Shared Encoder and Decoder (AAAI'19). Comparing the regularization effect with 6 layers / 1024 dimension as a baseline setup feels like cherry picking.  For NMT it's also much more common to report BLEU vs perplexity.\n\nW5: if the paper's goal is to demonstrate co-distillation as a viable distributed training alternative, it would be also valuable to compare it with other related training methods like gossip SGD etc.\n\n\n====== Edit after author response ====\n\nI have read authors' responses and have decided to keep my score as is.\n\nThe additional experiments haven't addressed problem formulation issues (W1, W2, W5) if the paper is positioned as more of a theoretical work; if the paper is positioned as a more of an experimental work, the baselines used (W3, W4) need to be improved with proper hparams settings.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a novel and interesting approach called co-distillation for distributed training. The main idea is to add a regularizer in order to encourage local models to be consistent with the global objective. Although the idea is a promising alternative to local-update SGD, the approach is mostly empirical. The claim that co-distillation helps reduce overfitting could be better justified by theoretical analysis, in addition to the experimental results. I hope that the reviewers' constructive comments will help improve the paper for future re-submission.
Decision: Reject