A Closer Look at Codistillation for Distributed Training

1 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. , 2016 ; Goyal et al. , 2017 ; Vaswani et 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 a neural network model are created , each on a different worker . The workers process different mini-batches locally at each step using an optimizer such as Stochastic Gradient Descent ( SGD ) or Adam ( Kingma & Ba , 2015 ) , and the replicas synchronize ( i.e. , average either gradients or parameters ) at every step by communicating either with a centralized parameter server ( Li et al. , 2014 ) or using all reduce ( Goyal et al. , 2017 ) . More computing resources can be used in parallel by increasing the number of workers , effectively increasing the batch size used to compute a 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 a given model quality is also reduced in proportion to the number of workers , and a near-linear speedup is achieved ( Goyal et 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. , 2017 ) , and even with more principled 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 a bottleneck limiting the processing units ’ utilization ( e.g. , GPU or TPU ) , especially when devices communicate over commodity interconnects such as Ethernet . A number of approaches have been proposed to reduce communication overhead , including using mixed-precision ( Jia et al. , 2018 ) or other forms of compression ( Alistarh et al. , 2017 ; Bernstein et al. , 2018 ) , reducing the frequency of synchronization to not occur after every optimizer step ( Stich , 2018 ; Yu et al. , 2019 ) , using gossip-based methods for approximate distributed averaging ( Lian et 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 to distributed training with reduced communication overhead ( Anil et al. , 2018 ) . Rather than synchronizing models to have the same weights , codistillation seeks to 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 a 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 checkpoint from up to 50 updates ago without observing a 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 is still significantly lower than that achieved by synchronous SGD with a smaller batch size ( e.g. , 8k or smaller ) . This performance boost is 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 of synchronous SGD begins to degrade . We demonstrate that it is possible to 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 a batch size of 256 , achieves comparable performance 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 a regularizing effect . Thus , while increasing the batch size in synchronous SGD helps training by reducing the gradient variance , we conjecture that codistillation helps in a complementary manner via this regularization . Because it has a regularizing effect , care needs to be taken when using codistillation in conjunction with other forms of regularization , such as L2 regularization ( weight decay ) , to avoid over-regularizing . We also evaluate the sensitivity of codistillation to different hyper-parameters like the frequency of reading new checkpoints and learning rate schedule . Overall , our findings complement previous work on codistillation ( Anil et al. , 2018 ; Zhang et al. , 2018 ) . We summarize below our main contributions : 1 . To the best of our knowledge , we 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 a much weaker synchronization mechanism ( Section 3 ) . Previous work at the intersection of codistillation and distributed training used extremely large batch sizes and more workers than the parallel SGD counterparts . 2 . Complementing the existing work on codistillation , we show that codistillation acts as a regularizer ( Section 4 ) . Our work demonstrates that explicitly accounting for its regularization effect is a key ingredient to using codistillation without losing accuracy ( compared to parallel SGD methods ) . 2 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 et al. , 2018 ) . In typical multi-phase 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 Algorithm 1 : Codistillation Input : Loss function L ( y , ŷ ) and codistillation loss function D ( y , y′ ) Input : Model architecture fθ ( x ) and initial model parameters { θ1i : i = 1 , . . . , n } Input : Number of iterations K , learning rates { ηk } Kk=1 , and penalty coefficients { αk } Kk=1 1 for k = 1 , . . . , K do 2 for i = 1 , . . . , n do 3 x , y = get next minibatch ( ) 4 θk+1i = θ k i − ηk∇θi ( L ( y , fθki ( x ) ) + α k 1 n−1 ∑ j 6=iD ( fθki ( x ) , fθkj ( x ) ) ) 5 end 6 end or more models codistill , there is only one phase , and in addition to minimizing the usual supervised loss on the training data , an 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 et al . ( 2018 ) , is shown in Algorithm 1 . Here , n ≥ 2 models are trained concurrently . The ith model is updated on line 4 by taking a gradient step to minimize the combination of a standard supervised loss function L ( e.g. , cross-entropy or MSE ) and a distillationlike loss D which penalizes differences between the predictions made by model i and those made by model j , averaged over all other models j 6= i. Zhang et al . ( 2018 ) and Anil et al . ( 2018 ) both report using Kullback-Liebler ( KL ) divergence for D in their experiments , although Anil et al . ( 2018 ) mentions that other options are possible ( e.g. , mean squared error between the logits of different models ) . Zhang et al . ( 2018 ) does not explicitly include a penalty parameter αk , and instead ( implicitly ) takes αk = 1 for all k. Anil et al . ( 2018 ) focuses on codistillation as a mechanism for scaling distributed training , and demonstrates that using codistillation allows for better scaling with larger batch sizes than conventional synchronous parallel SGD . When different models reside on different GPUs , implementing line 4 requires that the devices communicate . Anil et al . ( 2018 ) proposes to have workers exchange parameter checkpoints , and these checkpoints are updated periodically rather than after every iteration . This leads to some delay , i.e. , actually using θk ′ j with k ′ < k in line 4 . However , it is argued that such delay is tolerable for codistillation because large changes in parameters θ do not necessarily correspond to large changes to the function fθ . 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 a 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. , αk = 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 is to train a single model to mimic an ensemble ’ s predictions to avoid the computational overhead of deploying an 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 independently-trained models . Anil et al . ( 2018 ) also conducts an experiment that compares codistilling two models with simply using label smoothing as a regularizer when training one model , and reports that codistillation performs significantly better . 3 CODISTILLATION , SYNCHRONOUS SGD , AND REGULARIZATION . The aim of this section is to demonstrate that codistillation can be used as a mechanism for distributed 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 carefully . Typical training regimes use some form of regularization to avoid overfitting . For example , the common practice is to use L2 regularization ( weight decay ) when training convolutional networks ( Goyal et al. , 2017 ) and to use label smoothing when training neural machine translation models ( Ott et al. , 2018 ) . Directly incorporating codistillation on top of existing training pipelines can lead to underfitting . We begin this section by illustrating this point .

Summary: 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.