Accelerating DNN Training through Selective Localized Learning 

Training Deep Neural Networks ( DNNs ) places immense compute requirements on the underlying hardware platforms , expending large amounts of time and energy . We propose LoCal+SGD , a new algorithmic approach to accelerate DNN training by selectively combining localized or Hebbian learning within a Stochastic Gradient Descent ( SGD ) based training framework . Back-propagation is a computationally expensive process that requires 2 Generalized Matrix Multiply ( GEMM ) operations to compute the error and weight gradients for each layer . We alleviate this by selectively updating some layers ’ weights using localized learning rules that require only 1 GEMM operation per layer . Further , since the weight update is performed during the forward pass itself , the layer activations for the mini-batch do not need to be stored until the backward pass , resulting in a reduced memory footprint . Localized updates can substantially boost training speed , but need to be used selectively and judiciously in order to preserve accuracy and convergence . We address this challenge through the design of a Learning Mode Selection Algorithm , where all layers start with SGD , and as epochs progress , layers gradually transition to localized learning . Specifically , for each epoch , the algorithm identifies a Localized→SGD transition layer , which delineates the network into two regions . Layers before the transition layer use localized updates , while the transition layer and later layers use gradient-based updates . The trend in the weight updates made to the transition layer across epochs is used to determine how the boundary between SGD and localized updates is shifted in future epochs . We also propose a low-cost weak supervision mechanism by controlling the learning rate of localized updates based on the overall training loss . We applied LoCal+SGD to 8 image recognition CNNs ( including ResNet50 and MobileNetV2 ) across 3 datasets ( Cifar10 , Cifar100 and ImageNet ) . Our measurements on a Nvidia GTX 1080Ti GPU demonstrate upto 1.5× improvement in end-to-end training time with ∼0.5 % loss in Top-1 classification accuracy . 1 INTRODUCTION . Deep Neural Networks ( DNNs ) have achieved continued success in many application domains involving images ( Krizhevsky et al. , 2017 ) , videos ( Ng et al. , 2015 ) , text ( Zhou et al. , 2015 ) and natural language ( Goldberg & Hirst , 2017 ) . However training state-of-the-art DNN models is computationally quite challenging , often requiring exa-FLOPs of compute as the models are quite complex and need to be trained using large datasets . Despite rapid improvements in the capabilities of GPUs and the advent of specialized accelerators , training large models using current platforms is still quite expensive and often takes days to even weeks . In this work , we aim to reduce the computational complexity of DNN training through a new algorithmic approach called LoCal+SGD1 , which alleviates the key performance bottlenecks in Stochastic Gradient Descent ( SGD ) through selective use of localized or Hebbian learning . Computational Bottlenecks in DNN Training . DNNs are trained in a supervised manner using gradient-descent based cost minimization techniques such as SGD ( Bottou , 2010 ) or Adam ( Kingma & Ba , 2015 ) . The training inputs ( typically grouped into minibatches ) are iteratively forward propagated ( FP ) and back propagated ( BP ) through the DNN layers to compute weight updates that push the network parameters in the direction that decreases the overall classification loss . 1In addition to combining localized and SGD based learning , LoCal+SGD is Low-Calorie SGD or SGD with reduced computational requirements Back-propagation is computationally expensive , accounting for 65-75 % of the total training time on GPUs . This is attributed to two key factors : ( i ) BP involves 2 Generalized Matrix Multiply ( GEMM ) operations , one to propagate the error across layers and the other to compute the weight gradients , and ( ii ) when training on distributed systems using data/model parallelism ( Dean et al. , 2012b ; Krizhevsky et al. , 2012 ) , aggregation of weight gradients/errors across devices incurs significant communication overhead . Further , BP through auxiliary ops such as batch normalization are also more expensive than FP . Prior Efforts on Efficient DNN Training . Prior research efforts to improve DNN training time can be grouped into a few directions . One group of efforts enable larger scales of parallelism in DNN training through learning rate tuning ( You et al. , 2017a ; Goyal et al. , 2017 ; You et al. , 2017b ) and asynchronous weight updates ( Dean et al. , 2012a ) . Another class of efforts employ importancebased sample selection during training , wherein ‘ easier ’ training samples are selectively discarded to improve runtime ( Jiang et al. , 2019 ; Zhang et al. , 2019 ) . Finally , model quantization ( Sun et al. , 2019 ) and pruning ( Lym et al. , 2019 ) can lead to significant runtime benefits during training by enabling the use of reduced-bitwidth processing elements . LoCal+SGD : Combining SGD with Localized Learning . Complementary to the aforementioned efforts , we propose a new approach , LoCal+SGD , to alleviate the performance bottlenecks in DNN training , while preserving model accuracy . Our hybrid approach combines Hebbian or localized learning ( Hebb ) with SGD by selectively applying it in specific layers and epochs . Localized learning rules ( Hebb ; Oja , 1982 ; Zhong , 2005 ) utilize a single feed-forward weight update to learn the feature representations , eschewing BP . Careful formulation of the localized learning rule can result in ∼2× computation savings compared to SGD and also significantly reduces memory footprint as activations from FP needed not be retained until BP . The reduction in memory footprint can in turn allow increasing the batch size during training , which leads to further runtime savings due to better compute utilization and reduced communication costs . It is worth noting that localized learning has been actively explored in the context of unsupervised learning ( Chen et al. , 2020 ; van den Oord et al. , 2018 ; Hénaff et al. , 2019 ) . Further , there has been active research efforts on neuro-scientific learning rules ( Lee et al. , 2015 ; Nøkland , 2016 ) . Our work is orthogonal to such efforts and represents a new application of localized learning in a fully supervised context , wherein we selectively combine it within an SGD framework to achieve computational savings . Preserving model accuracy and convergence with LoCal+SGD requires localized updates to be applied judiciously i.e. , only to selected layers in certain epochs . We address this challenge through the design of a learning mode selection algorithm . At the start training , the selection algorithm initializes the learning mode of all layers to SGD , and as training progresses determines the layers that transition to localized learning . Specifically , for each epoch , the algorithm identifies a Localized→SGD transition layer , which delineates the network into two regions . Layers before the transition layer use localized updates , while subsequent layers use gradient-based updates . This allows BP to stop at the transition layer , as layers before it have no use for the back-propagated errors . The algorithm takes advantage of the magnitude of the weight updates of the Localized→SGD transition layer in deciding the new position of the boundary every epoch . Further , we provide weak supervision by tweaking the learning rate of locally updated layers based on overall training loss . Contributions : To the best of our knowledge , LoCal+SGD is the first effort that combines localized learning ( an unsupervised learning technique ) within a supervised SGD context to reduced computational costs while maintaining classification accuracy . This favorable tradeoff is achieved by LoCal+SGD through a Learning Mode Selection Algorithm that applies localized learning to selected layers and epochs . Further improvement is achieved through the use of weak supervision by modulating the learning rate of locally updated layers based on the overall training loss . Across 8 image recognition CNNs ( including ResNet50 and MobileNet ) and 3 datasets ( Cifar10 , Cifar100 and ImageNet ) , we demonstrate that LoCal+SGD achieves up to 1.5× improvement in training time with ∼0.5 % Top-1 accuracy loss on a Nvidia GTX 1080Ti GPU . 2 LoCal+SGD : COMBINING SGD WITH SELECTIVE LOCALIZED LEARNING The key idea in LoCal+SGD is to apply localized learning to selected layers and epochs during DNN training to improve the overall execution time , without incurring loss in accuracy . The following components are critical to the effectiveness of LoCal+SGD : • Localized Learning Rule Formulation . We formulate a computationally efficient localized learning rule and highlight the clear runtime benefits when compared to SGD . • Learning Mode Selection Algorithm . We propose a learning mode selection algorithm that chooses between localized learning and SGD-based learning for each layer in every epoch , based on the potential impact on accuracy and computational benefits . • Weak Supervision . We propose a weak supervision technique , which comprises of a low- cost supervision signal communicated to the localized learning layers in each epoch . The signal modulates the learning rates of these layers based on the rate of change of the overall classification loss . In the following sub-sections , we describe the salient aspects of these components in greater detail . 2.1 EFFICIENT LOCALIZED LEARNING . Localized learning has been extensively explored in the context of unsupervised learning , demonstrating success on small ( < = 3 layer ) networks using relatively simpler datasets ( e.g . MNIST , Cifar-10 ) ( LeCun & Cortes , 2010 ; Krizhevsky et al. , a ) ) with an accuracy gap that is yet to be bridged on larger datasets ( e.g . ResNet50 or MobileNetV2 on ImageNet ( Deng et al. , 2009 ) ) . First proposed in ( Hebb ) , the key intuition behind localized learning rules is to encourage correlations between neurons that have similar activation patterns . Equation 1 depicts the Hebbian weight update proposed in ( Hebb ) , for a synapse with weight W , connecting a pair of input and output neurons whose activation values are represented by x and y respectively , with η as the learning rate . 4W = η · x · y ( 1 ) Considerable research has gone into evolving this equation over the years to improve the performance of localized learning ( Oja , 1982 ; Zhong , 2005 ) . However , many of the proposed rules are computationally complex , or are difficult to parallelize on modern hardware platforms such as GPUs and TPUs . Since our primarily goal is improving DNN training time , we adopt the computationally simple localized learning rule presented in Equation 1 . The learning rule in Equation 1 assumes a distinct synapse between each input and output neuron pair . While its application to fully-connected ( fc ) layers is straightforward , we need to consider the sharing of weights between neuron pairs in convolutional ( conv ) layers . For updating a shared weight of a conv layer , we calculate the individual updates due to each pair of pre- and post-synaptic neurons sharing the weight and sum all such updates . This essentially reduces to a convolution operation between the input and output activations of the layer and can be expressed by Equation 3 in Figure 1 . For further computational efficiency improvement , unlike Equation 1 , we consider the pre-activation-function values of the outputs i.e. , zl instead of their post activation value al . Further , we normalize the localized update values as shown in Equation 4 of Figure 1 , as it was observed to achieve better convergence in practice . Overall , we utilize Equations 3 and 4 from Figure 1 to perform the weight updates in all layers that are earlier than the Localized→SGD transition layer during a certain epoch . All other layers continue to be updated using SGD-based BP , expressed by Equations 5-7 in Figure 1 . SGD updates are applied to batch-normalization layers present after the Localized→SGD transition layer , and are otherwise skipped . Clearly , Equation 3 has the same computational complexity as Equation 6 of SGD-based BP for conv and fc layers . Thus , from Figure 1 , we can directly infer that our localized learning rule will be considerable faster than SGD-based BP . In practice , we measured this improvement to be more than 2× on a NVIDIA GTX 1080Ti GPU for the ImageNet-ResNet50 benchmark , across all conv and fc layers . In addition to the computational complexity , the memory footprint of SGD-based BP is also higher . This is because DNN software frameworks commonly store all activation values computed during FP to avoid recomputing al−1 , the input activations to the layers , used in Equation 6 of SGD-based BP . In contrast , the localized update for a layer can be performed as soon as the FP through the layer is complete . The activation tensor al of layer L can be discarded or over-written as soon as FP proceeds to the next layer in the network , thereby freeing up a significant portion of on-device memory during training . In turn , this can allow larger minibatch sizes to be accommodated on a given hardware platform , when the localized updates are applied on a sufficient number of layers .

Summary: This paper try to leverage the benefit of Hebb learning to reduce CNN training time cost. In order to achieve this,  a learning mode selection algorithm is proposed to progressively increase  number of layers using Hebb learning. The writing of this paper is good and the idea is also interesting, however, the experimental part should be improved: