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 applications such as real-time adaptability to device drift and environmental variation, user customization, and federated learning across devices. In this work, we address four key challenges for training on edge devices with non-volatile memory: low weight update density, weight quantization, low auxiliary memory, and online learning. We present a low-rank training scheme that addresses these four 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 updates. 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 (Konecný 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 10 6 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. We demonstrated the potential for SKS to solve the major challenges facing online training on NVM-based edge devices: low write density and low auxiliary memory. SKS is a computationallyefficient, memory-light algorithm capable of decoupling batch size from auxiliary memory, allowing larger effective batch sizes, and consequently lower write densities. Additionally, we noted that SKS may allow for training under severe weight quantization constraints as rudimentary gradient accumulations are handled by the L, R matrices, which can have high bitwidths (as opposed to SGD, which may squash small gradients to 0). We found expressions for when SKS might have better convergence properties. Across a variety of online adaptation problems and a large-scale transfer learning demonstration, SKS was shown to match or exceed the performance of SGD while using a small fraction of the number of updates. Finally, we suspect that these techniques could be applied to a broader range of problems. Auxiliary memory minimization may be analogous to communication minimization in training strategies such as federated learning, where gradient compression is important.
We use Kronecker sum approximations for low-rank training to address challenges in training neural networks on edge devices that utilize emerging memory technologies.
{"src": "scitldr-AIC_train", "paper_id": "SkeXL0NKwH"}