On-Device Training Under 256KB Memory
By Ji Lin, Ligeng Zhu, Wei-Ming Chen, Wei-Chen Wang, Chuang Gan, Song Han. Neural Information Processing Systems, 2022
317 citations recorded by the source corpus.
On-device training enables the model to adapt to new data collected from the sensors by fine-tuning a pre-trained model. Users can benefit from customized AI models without having to transfer the data to the cloud, protecting the privacy. However, the training memory consumption is prohibitive for IoT devices that have tiny memory resources. We propose an algorithm-system co-design framework to make on-device training possible with only 256KB of memory. On-device training faces two unique challenges: (1) the quantized graphs of neural networks are hard to optimize due to low bit-precision and the lack of normalization; (2) the limited hardware resource does not allow full back-propagation. To cope with the optimization difficulty, we propose Quantization-Aware Scaling to calibrate the gradient scales and stabilize 8-bit quantized training. To reduce the memory footprint, we propose Sparse Update to skip the gradient computation of less important layers and sub-tensors. The algorithm innovation is implemented by a lightweight training system, Tiny Training Engine, which prunes the backward computation graph to support sparse updates and offload the runtime auto-differentiation to compile time. Our framework is the first solution to enable tiny on-device training of convolutional neural networks under 256KB SRAM and 1MB Flash without auxiliary memory, using less than 1/1000 of the memory of PyTorch and TensorFlow while matching the accuracy on tinyML application VWW. Our study enables IoT devices not only to perform inference but also to continuously adapt to new data for on-device lifelong learning. A video demo can be found here: https://youtu.be/0pUFZYdoMY8.
Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition
View Paper: https://www.semanticscholar.org/paper/c2c6bb5c70488897afb0a96dfabe39a607605210
How researchers see this paper
The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.
In Artificial Intelligence
In Computer Vision and Pattern Recognition
Citations in the corpus
- OLLA: Optimizing the Lifetime and Location of Arrays to Reduce the Memory Usage of Neural Networks (2022, background): “[52] also mentioned reordering computations as a way to enable operator fusion and reduce the peak memory footprint while training.”