Parallel Training of Deep Networks with Local Updates

1 INTRODUCTION . Backpropagation ( Rumelhart et al. , 1985 ) is by far the most common method used to train neural networks . Alternatives to backpropagation are typically used only when backpropagation is impractical due to a non-differentiable loss ( Schulman et al. , 2015 ) , non-smooth loss landscape ( Metz et al. , 2019 ) , or due to memory and/or compute requirements ( Ororbia et al. , 2020 ) . However , progress in deep learning is producing ever larger models in terms of parameter count and depth , in vision ( Hénaff et al. , 2019 ; Chen et al. , 2020 ) , language ( Radford et al. , 2019 ; Brown et al. , 2020 ) , and many other domains ( Silver et al. , 2017 ; Vinyals et al. , 2019 ; Berner et al. , 2019 ) . As model size increases , backpropagation incurs growing computational , memory , and synchronization overhead ( Ben-Nun & Hoefler , 2018 ) . This raises the question of whether there are more efficient training strategies , even for models and losses that are considered well matched to training by backpropagation . Much of the work on training large scale models focuses on designing compute infrastructure which makes backpropagation more efficient , despite growing model size ( Dean et al. , 2012b ; Chen et al. , 2015 ; Sergeev & Balso , 2018 ) . One of the most common ways to achieve efficient training of deep neural networks with backpropagation is to scale utilizing data parallelism ( Zhang et al. , 1989 ; Chen et al. , 2016 ) , training on bigger batch sizes spread across multiple devices . However , diminishing returns have been reported with this method for larger batch sizes , effectively wasting compute ( Goyal et al. , 2017 ; Masters & Luschi , 2018 ; Shallue et al. , 2018 ; McCandlish et al. , 2018 ) . Training based on pipeline parallelism has also been introduced , but still requires large batches for efficient training ( Petrowski et al. , 1993 ; Ben-Nun & Hoefler , 2018 ; Huang et al. , 2019 ) . Moreover , in addition to the limitation that in the forward pass each layer can only process the input data in sequence ( forward locking ) , the use of backpropagation implies that the network parameters of each layer can only be updated in turn after completing the full forward pass ( backward locking ) . This backward locking results in increased memory overhead , and precludes efficient parallel processing across layers ( Jaderberg et al. , 2017 ) . The challenges of scaling compute infrastructure to support deep networks trained with backpropagation motivate the need for alternative approaches to training deep neural networks . In this work , we explore how layer-wise local updates ( Belilovsky et al. , 2019a ; Löwe et al. , 2019 ; Xiong et al. , 2020 ) can help overcome these challenges and scale more efficiently with compute than backpropagation . With local updates , each layer is updated before even completing a full forward pass through the network . This remedies the forward and backward locking problems which harm memory efficiency and update latency in standard backprop . Layer-wise local updates are not proportional to gradients of the original loss , and are not even guaranteed to descend a loss function . Nevertheless , in practice they are effective at training neural networks . We refer to this approach of parallelizing compute , which is alternative and complementary to data and model parallelism , as local parallelism . Our investigation focuses on the trade-offs of using local update methods as opposed to global backpropagation . To summarize our contributions : ( i ) We provide the first large scale investigation into local update methods in both vision and language domains . We find training speedups ( as measured by the reduction in required sequential compute steps ) of up to 10× on simple MLPs , and 2× on Transformer architectures . These training speedups are the result of local training methods being able to leverage more parallel compute than backprop . ( ii ) We provide insight into how local parallelism methods work , and experimentally compare the similarity of their gradient and features to those from backprop . ( iii ) We demonstrate a prototype implementation of local parallelism for ResNets , and show up to a 40 % increase in sample throughput ( number of training points per second ) relative to backprop , due to higher hardware utilization . We believe that local parallelism will provide benefits whenever there are diminishing returns from data parallelism , and avoid stale weights from pipelined model parallelism . Additionally , we have released code showing an example of local parallelism , available at hiddenurl . 2 RELATED WORK . 2.1 PARALLELIZATION IN DEEP LEARNING . Scaling large models has led to the development of a number of techniques to train deep models in a parallel fashion ( Ben-Nun & Hoefler , 2018 ) , summarized in Figure 1 . Data Parallelism : Data Parallelism ( Zhang et al. , 1989 ) is an attempt to speed up training of a model by splitting the data among multiple identical models and training each model on a shard of the data independently . Data parallelism is effectively training with larger minibatches ( Kaplan et al. , 2020 ) . This creates issues around the consistency of a model which then needs to be synchronized ( Deng et al. , 2012 ; Dean et al. , 2012a ) . There are two main ways to synchronize weights across model copies : ( i ) Synchronous optimization , where data parallel training synchronizes at the end of every minibatch ( Das et al. , 2016 ; Chen et al. , 2016 ) , with a communication overhead that increases with the number of devices ; ( ii ) Asynchronous optimization that implements data parallel training with independent updates of local model parameters without global synchronization ( Niu et al. , 2011 ; Dean et al. , 2012a ) – this increases device utilization , but empirically gradients are computed on stale weights , which results in a poor sample efficiency and thus slower overall training time compared to synchronous optimization . Model Parallelism : Model Parallelism is used when a model is too large to fit in the memory of a single device and is instead spread over multiple processors ( Krizhevsky et al. , 2012 ; Shazeer et al. , 2018 ; Harlap et al. , 2018 ; Lepikhin et al. , 2020 ) . This is increasingly common as state of the art performance continues to improve with increasing model size ( Brown et al. , 2020 ) . Model parallelism unfortunately has a few downsides : ( i ) High communication costs – the total training time for larger networks can become dominated by communication costs ( Simonyan & Zisserman , 2015 ) , which in the worst case can grow quadratically with the number of devices , and can reach up to 85 % of the total training time of a large model such as VGG-16 ( Harlap et al. , 2018 ; Simonyan & Zisserman , 2015 ) ; ( ii ) Device under-utilization – forward propagation and backward propagation are both synchronous operations , which can result in processor under-utilization in model-parallel systems . This problem becomes worse as we increase the number of layers ( Ben-Nun & Hoefler , 2018 ; Jia et al. , 2014 ; Collobert et al. , 2011 ; Abadi et al. , 2016 ; Huang et al. , 2018 ) . Pipeline Parallelism : Due to the forward and backward locking , using multiple devices to process consecutive blocks of the deep model would make an inefficient use of the hardware resources . Pipelining ( Harlap et al. , 2018 ) concurrently passes multiple mini-batches to multiple layers on multiple devices . This increases device utilization but can introduce staleness and consistency issues which lead to unstable training . Harlap et al . ( 2018 ) alleviates the consistency issue by storing past versions of each layer . Huang et al . ( 2019 ) addresses the staleness issue by pipelining microbatches and synchronously updating at the end of each minibatch . Guan et al . ( 2019 ) builds on this work by introducing a weight prediction strategy and Yang et al . ( 2020 ) investigates to what extent the tradeoff between staleness/consistency and device utilization is necessary . Local updates on the other hand can keep device utilization high with both small and large batches and avoid the weight staleness problem . Local Learning Rules : Local learning describes a family of methods that perform parameter updates based only on local information , where locality is defined as dependence of neighboring neurons , layers , or groups of layers . The earliest local method we are aware of is Hebbian Learning ( Hebb , 1949 ) which has further been explored in BCM theory ( Izhikevich & Desai , 2003 ; Coesmans et al. , 2004 ) , Oja ’ s rule ( Oja , 1982 ) , Generalized Hebbian Learning ( Sanger , 1989 ) , and meta-learned local learning rules ( Bengio et al. , 1990 ; 1992 ; Metz et al. , 2018 ; Gu et al. , 2019 ) . Architectures like Hopfield Networks ( Hopfield , 1982 ) and Boltzmann Machines ( Ackley et al. , 1985 ) also employ a local update , and predate backprogation in deep learning . Modern variants of local training methods have attempted to bridge the performance gap with backpropagation . These include projection methods such as Hebbian learning rules for deep networks ( Krotov & Hopfield , 2019 ; Grinberg et al. , 2019 ; Ryali et al. , 2020 ) , and local layer-wise learning with auxiliary losses ( Belilovsky et al. , 2019a ; b ) . Most similar to our work is decoupled greedy layer-wise learning ( Belilovsky et al. , 2019b ; Löwe et al. , 2019 ) , which trained auxiliary image classifiers greedily , and local contrastive learning ( Xiong et al. , 2020 ) . These methods mainly focus on matching the performance of backpropagation with respect to training epochs , whereas our work focuses on tradeoffs . Finally , while not local in the sense that parallelized layers still optimize for the global objective , Huo et al . ( 2018b ) parallelize layers by caching gradients and using delayed gradient signals to overcome the backward locking problem and update decoupled layers in parallel . 3 LOCAL PARALLELISM . Given a deep neural network , we divide the layers into a sequence of J blocks , which may contain one or more layers . Each block is trained independently with an auxiliary objective , and receives the activations output by the previous block as input or , in the case of the first block , the data from the sampled minibatch . We consider five variants to train this sequence of J blocks : backpropagation , greedy local parallelism , overlapping local parallelism , and chunked local parallelism , as shown in Figure 2 . We also include a baseline method of just training the last , or last two , layers . In all of the local methods , training occurs by attaching objective functions to the end of each block and back propagating the signal locally into the corresponding block or blocks . In this work the auxiliary objective functions that we use take the same form as the global objective . For example , to train a classifier on CIFAR-10 , we attach auxiliary linear classifiers to each local block . See Belilovsky et al . ( 2019b ) for further discussion on the form of this objective . Backpropagation : In our notation , backpropagation groups all layers into one block and thus J = 1 . The parameters are updated with one instance of global error correction . While backpropagation ensures that all weights are updated according to the final output loss , it also suffers from forward and backward locking ( Jaderberg et al. , 2017 ) , an issue that local parallelized methods aim to resolve . Greedy local parallelism : A straightforward approach to enable local training is to attach an auxiliary network to each local layer , which generates predictions from the activations of hidden layers . After generating predictions , each local gradient is backpropagated to its respective local block , shown in Figure 2 ( b ) . The activations are then passed as input to the next layer . We refer to this approach , introduced in ( Belilovsky et al. , 2019b ) , as greedy . Greedy local parallelism is the most parallelizable of all the schemes we consider . However , a potential downside is that fully greedy updates force the layers to learn features that are only relevant to their local objective and preclude inter-layer communication , which may result in lower evaluation performance for the global objective , or worse generalization . Overlapping local parallelism : One issue with the purely greedy approach is that features learned for any individual block may not be useful for subsequent blocks , since there is no inter-block propagation of gradient . For this reason , we consider overlapping local architectures where the first layer of each block is also the last layer of the previous block , as shown in Figure 2 ( c ) , though overlapping of more layers is also possible . This redundancy enables inter-block propagation of gradient that is still local , since only neighboring blocks overlap . However , this comes at the cost of running additional backward passes . The overlapping architecture has appeared before in Xiong et al . ( 2020 ) , but was used only for contrastive losses . Ours is the first work to investigate overlapping local architectures for standard prediction objectives in computer vision and language . Overlapping updates are parallelizable , but come with the additional complexity of keeping duplicates of the overlapping components and averaging updates for these layers . Chunked local parallelism : The greedy architecture is maximally parallel in the sense that it distributes one layer per block . However , it is also possible to have fewer parallel blocks by combining multiple layers into one . We refer to this architecture , shown in Figure 2 ( d ) , as chunked local parallelism . This method trades off parallelizability and therefore throughput for an error signal that propagates through more consecutive layers . It differs from overlapping local parallelism by not needing to duplicate any layer . While previous work has investigated the asymptotic performance of chunked parallelism ( Belilovsky et al. , 2019b ) , ours is the first to consider the compute efficiency and parallelizability of local parallelism . By stacking multiple layers per each parallelized block , chunked parallelism sits between fully parallelized methods , such as greedy and overlapping updates , and fully sequential methods like backpropagation .

Summary: It is a very poorly written paper. Basic idea of finding a way to not have to wait for full forward pass is not new. Multiple research papers have been published from the extreme of using stale weight to some form of sub-network backdrop as a proxy for the full network. This paper proposed no new idea for local update. Prior work have all suffered with one or both of these two limitations: a) poor experimental framework, or b) not being able to meet the accuracy bar set by backprop. This work suffers from both.  Very poorly described experimental basis - and failing to come even close to the backprop accuracy target with any decent speedup claim. Former is my biggest concern. Section 6 starts with 'Here we show that performance gains of local parallelism can be realized on real hardware' - with near-zero description of any 'real' hardware, except a footnote on '1000 IPUs on a chip'.