Low Rank Training of Deep Neural Networks for Emerging Memory Technology

The recent success of neural networks for solving difficult decision tasks has incentivized incorporating smart decision making “ at the edge. ” However , this work has traditionally focused on neural network inference , rather than training , due to memory and compute limitations , especially in emerging non-volatile memory systems , where writes are energetically costly and reduce lifespan . Yet , the ability to train at the edge is becoming increasingly important as it enables real-time adaptability to device drift and environmental variation , user customization , and federated learning across devices . In this work , we address two key challenges for training on edge devices with non-volatile memory : low write density and low auxiliary memory . We present a low-rank training scheme that addresses these challenges while maintaining computational efficiency . We then demonstrate the technique on a representative convolutional neural network across several adaptation problems , where it out-performs standard SGD both in accuracy and in number of weight writes . 1 INTRODUCTION . Deep neural networks have shown remarkable performance on a variety of challenging inference tasks . As the energy efficiency of deep-learning inference accelerators improves , some models are now being deployed directly to edge devices to take advantage of increased privacy , reduced network bandwidth , and lower inference latency . Despite edge deployment , training happens predominately in the cloud . This limits the privacy advantages of running models on-device and results in static models that do not adapt to evolving data distributions in the field . Efforts aimed at on-device training address some of these challenges . Federated learning aims to keep data on-device by training models in a distributed fashion ( Konecný et al. , 2016 ) . On-device model customization has been achieved by techniques such as weight-imprinting ( Qi et al. , 2018 ) , or by retraining limited sets of layers . On-chip training has also been demonstrated for handling hardware imperfections ( Zhang et al. , 2017 ; Gonugondla et al. , 2018 ) . Despite this progress with small models , on-chip training of larger models is bottlenecked by the limited memory size and compute horsepower of edge processors . Emerging non-volatile ( NVM ) memories such as resistive random access memory ( RRAM ) have shown great promise for energy and area-efficient inference ( Yu , 2018 ) . However , on-chip training requires a large number of writes to the memory , and RRAM writes cost significantly more energy than reads ( e.g. , 10.9 pJ/bit versus 1.76 pJ/bit ( Wu et al. , 2019 ) ) . Additionally , RRAM endurance is on the order of 106 writes ( Grossi et al. , 2019 ) , shortening the lifetime of a device due to memory writes for on-chip training . In this paper , we present an online training scheme amenable to NVM memories to enable next generation edge devices . Our contributions are ( 1 ) an algorithm called Streaming Kronecker Sum Approximation ( SKS ) , and its analysis , which addresses the two key challenges of low write density and low auxiliary memory ; ( 2 ) two techniques “ gradient max-norm ” and “ streaming batch norm ” to help training specifically in the online setting ; ( 3 ) a suite of adaptation experiments to demonstrate the advantages of our approach . 2 RELATED WORK . Efficient training for resistive arrays . Several works have aimed at improving the efficiency of training algorithms on resistive arrays . Of the three weight-computations required in training ( forward , backprop , and weight update ) , weight updates are the hardest to parallelize using the array structure . Stochastic weight updates ( Gokmen & Vlasov , 2016 ) allow programming of all cells in a crossbar at once , as opposed to row/column-wise updating . Online Manhattan rule updating ( Zamanidoost et al. , 2015 ) can also be used to update all the weights at once . Several works have proposed new memory structures to improve the efficiency of training ( Soudry et al. , 2015 ; Ambrogio et al. , 2018 ) . The number of writes has also been quantified in the context of chip-in-the-loop training ( Yu et al. , 2016 ) . Distributed gradient descent . Distributed training in the data center is another problem that suffers from expensive weight updates . Here , the model is replicated onto many compute nodes and in each training iteration , the mini-batch is split across the nodes to compute gradients . The distributed gradients are then accumulated on a central node that computes the updated weights and broadcasts them . These systems can be limited by communication bandwidth , and compressed gradient techniques ( Aji & Heafield , 2017 ) have therefore been developed . In Lin et al . ( 2017 ) , the gradients are accumulated over multiple training iterations on each compute node and only gradients that exceed a threshold are communicated back to the central node . In the context of on-chip training with NVM , this method helps reduce the number of weight updates . However , the gradient accumulator requires as much memory as the weights themselves , which negates the density benefits of NVM . Low-Rank Training . Our work draws heavily from previous low-rank training schemes that have largely been developed for use in recurrent neural networks to uncouple the training memory requirements from the number of time steps inherent to the standard truncated backpropagation through time ( TBPTT ) training algorithm . Algorithms developed since then to address the memory problem include Real-Time Recurrent Learning ( RTRL ) ( Williams & Zipser , 1989 ) , Unbiased Online Recurrent Optimization ( UORO ) ( Tallec & Ollivier , 2017 ) , Kronecker Factored RTRL ( KF-RTRL ) ( Mujika et al. , 2018 ) , and Optimal Kronecker Sums ( OK ) ( Benzing et al. , 2019 ) . These latter few techniques rely on the weight gradients in a weight-vector product looking like a sum of outer products ( i.e. , Kronecker sums ) of input vectors with backpropagated errors . Instead of storing a growing number of these sums , they can be approximated with a low-rank representation involving fewer sums . 3 TRAINING NON-VOLATILE MEMORY . The meat of most deep learning systems are many weight matrix - activation vector productsW · a. Fully-connected ( dense ) layers use them explicitly : a [ ` ] = σ ( W [ ` ] a [ ` −1 ] + b [ ` ] ) for layer ` , where σ is a non-linear activation function ( more details are discussed in detail in Appendix C.1 ) . Recurrent neural networks use one or many matrix-vector products per recurrent cell . Convolutional layers can also be interpreted in terms of matrix-vector products by unrolling the input feature map into strided convolution-kernel-size slices . Then , each matrix-vector product takes one such input slice and maps it to all channels of the corresponding output pixel ( more details are discussed in Appendix C.2 ) . The ubiquity of matrix-vector products allows us to adapt the techniques discussed in “ Low-Rank Training ” of Section 2 to other network architectures . Instead of reducing the memory across time steps , we can reduce the memory across training samples in the case of a traditional feedforward neural network . However , in traditional training ( e.g. , on a GPU ) , this technique does not confer advantages . Traditional training platforms often have ample memory to store a batch of activations and backpropagated gradients , and the weight updates ∆W can be applied directly to the weights W once they are computed , allowing temporary activation memory to be deleted . The benefits of low-rank training only become apparent when looking at the challenges of proposed NVM devices : Low write density ( LWD ) . In NVM , writing to weights at every sample is costly in energy , time , and endurance . These concerns are exacerbated in multilevel cells , which require several steps of an iterative write-verify cycle to program the desired level . We therefore want to minimize the number of writes to NVM . Low auxiliary memory ( LAM ) . NVM is the densest form of memory . In 40nm technology , RRAM 1T-1R bitcells @ 0.085 um2 ( Chou et al. , 2018 ) are 2.8x smaller than 6T SRAM cells @ 0.242 um2 ( TSMC , 2019 ) . Therefore , NVM should be used to store the memory-intensive weights . By the same token , no other on-chip memory should come close to the size of the on-chip NVM . In particular , if our b−bit NVM stores a weight matrix of size no × ni , we should use at most r ( ni + no ) b auxiliary non-NVM memory , where r is a small constant . Despite these space limitations , the reason we might opt to use auxiliary ( large , high endurance , low energy ) memory is because there are places where writes are frequent , violating LWD if we were to use NVM . In the traditional minibatch SGD setting with batch size B , an upper limit on the write density per cell per sample is easily seen : 1/B . However , to store such a batch of updates without intermediate writes to NVM would require auxiliary memory proportional to B . Therefore , a trade-off becomes apparent . If B is reduced , LAM is satisfied at the cost of LWD . If B is raised , LWD is satisfied at the cost of LAM . Using low-rank training techniques , the auxiliary memory requirements are decoupled from the batch size , allowing us to increase B while satisfying both LWD and LAM1 . Additionally , because the low-rank representation uses so little memory , a larger bitwidth can be used , potentially allowing for gradient accumulation in a way that is not possible with low bitwidth NVM weights . In the next section , we elaborate on the low-rank training method . 4 LOW-RANK TRAINING METHOD . Let z ( i ) = Wa ( i ) + b be the standard affine transformation building block of some larger network , e.g. , y ( i ) p = fpost ( z ( i ) ) and a ( i ) = fpre ( x ( i ) ) with prediction loss L ( y ( i ) p , y ( i ) t ) , where ( x ( i ) , y ( i ) t ) is the ith training sample pair . Then weight gradient∇WL ( i ) = dz ( i ) ( a ( i ) ) > = dz ( i ) ⊗ a ( i ) where dz ( i ) = ∇z ( i ) L ( i ) . A minibatch SGD weight update accumulates this gradient over B samples : ∆W = −η ∑B i=1 dz ( i ) ⊗ a ( i ) for learning rate η . For a rank-r training scheme , approximate the sum ∑B i=1 dz ( i ) ⊗a ( i ) by iteratively updating two rank- r matrices L̃ ∈ Rno×r , R̃ ∈ Rni×r with each new outer product : L̃R̃ > ← rankReduce ( L̃R̃ > + dz ( i ) ⊗a ( i ) ) . Therefore , at each sample , we convert the rank-q = r+ 1 system L̃R̃ > +dz ( i ) ⊗a ( i ) into the rank-r L̃R̃ > . In the next sections , we discuss how to compute rankReduce . 4.1 OPTIMAL KRONECKER SUM APPROXIMATION ( OK ) . One option for rankReduce ( X ) to convert from rank q = r + 1 X to rank r is a minimum error estimator , which is implemented by selecting the top r components of a singular value decomposition ( SVD ) of X . However , a naı̈ve implementation is computationally infeasible and biased : E [ rankReduce ( X ) ] 6= X . Benzing et al . ( 2019 ) solves these problems by proposing a minimum variance unbiased estimator for rankReduce , which they call the OK algorithm2 . The OK algorithm can be understood in two key steps : first , an efficient method of computing the SVD of a Kronecker sum ; second , a method of splitting the singular value matrix Σ into two rank-r matrices whose outer product is a minimum-variance , unbiased estimate of Σ . Details can be found in their paper , however we include a high-level explanation in Sections 4.1.1 and 4.1.2 to aid our discussions . Note that our variable notation differs from Benzing et al . ( 2019 ) .

Summary: While inference on edge devices is a popular and well-studied problem in recent days, training on these devices comes with many challenges. This paper proposes a low-rank training schema that helps mitigate some of the critical challenges that occur during training models on NVM memory-based edge devices. Additionally, two techniques, namely streaming batch norm and gradient max norm, are proposed to help training in an online setting. The proposed method is mainly based on approximating the Kronecker sum and is largely inspired by (Benzing et al, ICML 2019, Optimal Kronecker-sum approximation of real time recurrent learning). The proposed approach provides a few optimizations that improves this performance further, and outperforms SGD in terms of accuracy and the number of weights updates in a limited experimental setting.